Skip to content

Commit

Permalink
More on uniqueness context checking
Browse files Browse the repository at this point in the history
  • Loading branch information
Eric-Song-Nop committed Jul 30, 2024
1 parent 6d0e46e commit 7191560
Show file tree
Hide file tree
Showing 4 changed files with 129 additions and 51 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSimpleFunctionChecker
import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph

class UniqueDeclarationChecker(private val session: FirSession, private val config: PluginConfiguration) :
Expand All @@ -21,7 +22,8 @@ class UniqueDeclarationChecker(private val session: FirSession, private val conf
if (!config.checkUniqueness) return
val errorCollector = ErrorCollector()
try {
val uniqueCheckerContext = UniqueCheckerData(session, config, errorCollector)
val assignArgumentPos = declaration.controlFlowGraphReference?.controlFlowGraph?.usedForAssignOrCall() ?: emptyMap()
val uniqueCheckerContext = UniqueCheckerData(session, config, errorCollector, assignArgumentPos = assignArgumentPos)

val cfaChecker = UniqueCFA(uniqueCheckerContext)
declaration.controlFlowGraphReference?.controlFlowGraph?.let { cfaChecker.analyze(it, reporter, context) }
Expand All @@ -30,4 +32,46 @@ class UniqueDeclarationChecker(private val session: FirSession, private val conf
reporter.reportOn(declaration.source, PluginErrors.UNIQUENESS_VIOLATION, error, context)
}
}

// This function resolve nodes use as arguments or rhs for assignments
// When an node is used in this case, we need to track alias or uniqueness level information in this.data
private fun ControlFlowGraph.usedForAssignOrCall(): Map<CFGNode<*>, CFGNode<*>> {
val result: MutableMap<CFGNode<*>, CFGNode<*>> = mutableMapOf()

fun resolveArgument(node: CFGNode<*>): CFGNode<*> {
var curNode: CFGNode<*> = node
var exitNode = node
while (exitNode !is FunctionCallArgumentsExitNode) {
if (exitNode.previousNodes.isEmpty()) {
TODO("Throw correctly")
}
exitNode = exitNode.previousNodes[0]
if (exitNode.owner != node.owner) return exitNode
}
while (true) {
if (curNode.previousNodes.isEmpty()) {
TODO("Throw correctly")
}
curNode = curNode.previousNodes[0]
if (curNode is ExitNodeMarker && curNode is FunctionCallArgumentsExitNode)
curNode = resolveArgument(curNode)
else if (curNode is EnterNodeMarker) break
else result[curNode] = exitNode
}
return curNode
}

for (node in nodes) {
if (node is FunctionCallArgumentsExitNode) {
resolveArgument(node)
}
if (node is VariableDeclarationNode || node is VariableAssignmentNode) {
if (node.previousNodes.isEmpty()) {
TODO("Throw correctly")
}
result[node.previousNodes[0]] = node
}
}
return result
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,47 +14,43 @@ import org.jetbrains.kotlin.fir.analysis.checkers.cfa.FirControlFlowChecker
import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
import org.jetbrains.kotlin.fir.declarations.FirVariable
import org.jetbrains.kotlin.fir.declarations.hasAnnotation
import org.jetbrains.kotlin.fir.expressions.toResolvedCallableSymbol
import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
import org.jetbrains.kotlin.fir.references.symbol
import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol

sealed class LastOrPath {
data object LastNode : LastOrPath()
data class Path(val path: FirVariableSymbol<FirVariable>) : LastOrPath()
}

class UniqueCFA(private val data: UniqueCheckerData) : FirControlFlowChecker(MppCheckerKind.Common) {
class UniqueCFA(private val data: UniqueCheckerContext) : FirControlFlowChecker(MppCheckerKind.Common) {
override fun analyze(graph: ControlFlowGraph, reporter: DiagnosticReporter, context: CheckerContext) {
graph.traverseToFixedPoint(UniqueInfoCollector(data))
}

private class UniqueInfoCollector(private val data: UniqueCheckerData) :
PathAwareControlFlowGraphVisitor<LastOrPath, Set<UniqueLevel>>() {
private class UniqueInfoCollector(private val data: UniqueCheckerContext) :
PathAwareControlFlowGraphVisitor<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>() {

override fun mergeInfo(
a: ControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
b: ControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
a: ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
b: ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
node: CFGNode<*>,
): ControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
): ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
// TODO: No implemented yet, for now just return from one branch
return a
}

override fun visitNode(
node: CFGNode<*>,
data: PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
// Check if node is a FirVariable<Variable>

// TODO: I should remove something here
val dataForNode = data.transformValues { it.remove(LastOrPath.LastNode) }
return super.visitNode(node, dataForNode)
override fun visitEdge(
from: CFGNode<*>,
to: CFGNode<*>,
metadata: Edge,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val dataForEdge = super.visitEdge(from, to, metadata, data)
return dataForEdge
}

override fun visitFunctionEnterNode(
node: FunctionEnterNode,
data: PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val declaration = node.fir

// construct context and start control flow analysis
Expand All @@ -64,41 +60,84 @@ class UniqueCFA(private val data: UniqueCheckerData) : FirControlFlowChecker(Mpp
}
val valueParameters =
declaration.valueParameters.associate {
Pair<LastOrPath, Set<UniqueLevel>>(
LastOrPath.Path(it.symbol),
Pair<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>(
it.symbol,
if (it.hasAnnotation(this.data.uniqueId, this.data.session)) {
setOf(UniqueLevel.Unique)
} else {
setOf(UniqueLevel.Shared)
}
)
}.toPersistentMap()
return super.visitFunctionEnterNode(node, persistentMapOf(Pair(NormalPath, valueParameters)))
return persistentMapOf(Pair(NormalPath, valueParameters))
}

override fun visitVariableDeclarationNode(
node: VariableDeclarationNode,
data: PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
val dataForNode = super.visitVariableDeclarationNode(node, data)
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val dataForNode = data
val lSymbol = node.fir.symbol
val rhsUniqueLevel = data[NormalPath]?.get(LastOrPath.LastNode) ?: setOf(UniqueLevel.Shared)
val temp = dataForNode.transformValues { it.put(LastOrPath.Path(lSymbol), rhsUniqueLevel) }
.transformValues { it.remove(LastOrPath.LastNode) }
return temp
if (this.data.nodeUniqueLevel.contains(node)) {
val nodeUniqueLevel = this.data.nodeUniqueLevel[node] ?: emptyMap()
if (nodeUniqueLevel.isEmpty() || nodeUniqueLevel.all { (_, k) -> k.isEmpty() }) {
return dataForNode.transformValues { it.put(lSymbol, setOf(UniqueLevel.Shared)) }
}
return dataForNode.transformValues { it.put(lSymbol, nodeUniqueLevel.firstNotNullOf { (_, v) -> v }) }
} else {
return dataForNode.transformValues { it.put(lSymbol, setOf(UniqueLevel.Shared)) }
}
}

override fun visitVariableAssignmentNode(
node: VariableAssignmentNode,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
// FIXME: this function might be totally wrong, variable assignment might be totally different from variable declaration
val dataForNode = data
val lSymbol = node.fir.lValue.toResolvedCallableSymbol(this.data.session) as FirVariableSymbol
if (this.data.nodeUniqueLevel.contains(node)) {
val nodeUniqueLevel = this.data.nodeUniqueLevel[node] ?: emptyMap()
if (nodeUniqueLevel.isEmpty() || nodeUniqueLevel.all { (_, k) -> k.isEmpty() }) {
return dataForNode.transformValues { it.put(lSymbol, setOf(UniqueLevel.Shared)) }
}
return dataForNode.transformValues { it.put(lSymbol, nodeUniqueLevel.firstNotNullOf { (_, v) -> v }) }
} else {
return dataForNode.transformValues { it.put(lSymbol, setOf(UniqueLevel.Shared)) }
}
}

override fun visitFunctionCallArgumentsExitNode(
node: FunctionCallArgumentsExitNode,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
// TODO: this is where function call parameter checking should happen
if (this.data.nodeUniqueLevel.contains(node)) {
println("$node: ${this.data.nodeUniqueLevel[node]}")
}
return data
}

override fun visitQualifiedAccessNode(
node: QualifiedAccessNode,
data: PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>>
): PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
val dataForNode = visitNode(node, data)
if (node.fir.calleeReference is FirResolvedNamedReference) {
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val dataForNode = data
if (node.fir.calleeReference is FirResolvedNamedReference && this.data.assignArgumentPos.contains(node)) {
// When accessed node is used for assignment and is a path
// When accessed node is a unique variable need to change its uniqueness level to Top
// In all cases store the unique level
val usePos: CFGNode<*> = this.data.assignArgumentPos[node]!!
val callee = node.fir.calleeReference.symbol as FirVariableSymbol
val rhsUniqueLevel = data[NormalPath]?.get(LastOrPath.Path(callee)) ?: setOf(UniqueLevel.Shared)
return data.transformValues { it.put(LastOrPath.LastNode, rhsUniqueLevel) }
val rhsUniqueLevel = data[NormalPath]?.get(callee) ?: setOf(UniqueLevel.Shared)
if (this.data.nodeUniqueLevel[usePos].isNullOrEmpty()) {
this.data.nodeUniqueLevel[usePos] = mapOf(node to rhsUniqueLevel).toMutableMap()
} else {
this.data.nodeUniqueLevel[usePos]?.set(node, rhsUniqueLevel)
}
return dataForNode
}
return super.visitQualifiedAccessNode(node, data)
return dataForNode
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ package org.jetbrains.kotlin.formver

import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.declarations.FirDeclaration
import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
import org.jetbrains.kotlin.name.ClassId

Expand All @@ -15,8 +16,8 @@ interface UniqueCheckerContext {
val errorCollector: ErrorCollector
val session: FirSession
val uniqueId: ClassId
val nodeUniqueLevel: MutableMap<CFGNode<*>, MutableMap<CFGNode<*>, Set<UniqueLevel>>>
val assignArgumentPos: Map<CFGNode<*>, CFGNode<*>>

fun resolveUniqueAnnotation(declaration: FirDeclaration): UniqueLevel
fun getPathUniqueLevel(path: FirPropertySymbol): Set<UniqueLevel>
fun setPathUniqueLevel(path: FirPropertySymbol, level: UniqueLevel)
}
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ package org.jetbrains.kotlin.formver
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.declarations.FirDeclaration
import org.jetbrains.kotlin.fir.declarations.hasAnnotation
import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
Expand All @@ -17,7 +18,8 @@ class UniqueCheckerData(
override val session: FirSession,
override val config: PluginConfiguration,
override val errorCollector: ErrorCollector,
public val pathUniqueTable: MutableMap<FirPropertySymbol, Set<UniqueLevel>> = mutableMapOf(),
override val nodeUniqueLevel: MutableMap<CFGNode<*>, MutableMap<CFGNode<*>, Set<UniqueLevel>>> = mutableMapOf(),
override val assignArgumentPos: Map<CFGNode<*>, CFGNode<*>>
) : UniqueCheckerContext {

private fun getAnnotationId(name: String): ClassId =
Expand All @@ -32,12 +34,4 @@ class UniqueCheckerData(
}
return UniqueLevel.Shared
}

override fun getPathUniqueLevel(path: FirPropertySymbol): Set<UniqueLevel> {
return this.pathUniqueTable[path] ?: setOf(UniqueLevel.Shared)
}

override fun setPathUniqueLevel(path: FirPropertySymbol, level: UniqueLevel) {
this.pathUniqueTable[path] = setOf(level)
}
}

0 comments on commit 7191560

Please sign in to comment.