Skip to content

Commit

Permalink
Store uniquelevel of expression with stack of stack
Browse files Browse the repository at this point in the history
  • Loading branch information
Eric-Song-Nop committed Jul 30, 2024
1 parent 7191560 commit 358a1d5
Show file tree
Hide file tree
Showing 4 changed files with 80 additions and 100 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ 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 @@ -22,56 +21,11 @@ class UniqueDeclarationChecker(private val session: FirSession, private val conf
if (!config.checkUniqueness) return
val errorCollector = ErrorCollector()
try {
val assignArgumentPos = declaration.controlFlowGraphReference?.controlFlowGraph?.usedForAssignOrCall() ?: emptyMap()
val uniqueCheckerContext = UniqueCheckerData(session, config, errorCollector, assignArgumentPos = assignArgumentPos)

val cfaChecker = UniqueCFA(uniqueCheckerContext)
val cfaChecker = UniqueCFA(UniqueCheckerData(session, config, errorCollector))
declaration.controlFlowGraphReference?.controlFlowGraph?.let { cfaChecker.analyze(it, reporter, context) }
} catch (e: Exception) {
val error = errorCollector.formatErrorWithInfos(e.message ?: "No message provided")
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 @@ -5,27 +5,26 @@

package org.jetbrains.kotlin.formver

import kotlinx.collections.immutable.persistentMapOf
import kotlinx.collections.immutable.toPersistentMap
import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
import org.jetbrains.kotlin.fir.analysis.cfa.util.*
import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
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.arguments
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
import org.jetbrains.kotlin.utils.addToStdlib.popLast

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

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

override fun mergeInfo(
Expand All @@ -47,73 +46,94 @@ class UniqueCFA(private val data: UniqueCheckerContext) : FirControlFlowChecker(
return dataForEdge
}

override fun visitNode(
node: CFGNode<*>,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
if (node is EnterNodeMarker) {
context.uniqueStack.add(ArrayDeque())
} else if (node is ExitNodeMarker) {
context.uniqueStack.popLast().lastOrNull()?.let { context.uniqueStack.last().add(it) }
}
return data
}

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

// construct context and start control flow analysis
if (declaration.receiverParameter != null) {
TODO("Ignore receiver for now, not sure how to convert it into property symbol yet")
val receiverParameter = declaration.receiverParameter?.annotations ?: emptyList()
}
val valueParameters =
declaration.valueParameters.associate {
Pair<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>(
it.symbol,
if (it.hasAnnotation(this.data.uniqueId, this.data.session)) {
if (it.hasAnnotation(this.context.uniqueId, this.context.session)) {
setOf(UniqueLevel.Unique)
} else {
setOf(UniqueLevel.Shared)
}
)
}.toPersistentMap()
return persistentMapOf(Pair(NormalPath, valueParameters))
}.toMap()
return dataForNode.transformValues { it.putAll(valueParameters) }
}

override fun visitVariableDeclarationNode(
node: VariableDeclarationNode,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val dataForNode = data
val dataForNode = visitNode(node, data)
val lSymbol = node.fir.symbol
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 {
if (node.fir.initializer == null) {
return dataForNode.transformValues { it.put(lSymbol, setOf(UniqueLevel.Shared)) }
}
val rhsUnique = when (val last = context.uniqueStack.last().last()) {
is Level -> last.level
is Path -> dataForNode[NormalPath]?.get(last.symbol) ?: setOf(UniqueLevel.Shared)
}
return dataForNode.transformValues { it.put(lSymbol, rhsUnique) }
}

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)) }
val dataForNode = visitNode(node, data)
val lSymbol = node.fir.lValue.toResolvedCallableSymbol(context.session)
val rhsUnique = when (val last = context.uniqueStack.last().last()) {
is Level -> last.level
is Path -> dataForNode[NormalPath]?.get(last.symbol) ?: setOf(UniqueLevel.Shared)
}
return dataForNode.transformValues { it.put(lSymbol as FirVariableSymbol, rhsUnique) }
}

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]}")
// Here do not call visitNode
val arguments = node.fir.arguments
val argumentAliasOrUnique: MutableList<PathUnique> = mutableListOf()
for (i in arguments.indices) {
argumentAliasOrUnique.add(context.uniqueStack.last().popLast())
}
// FIXME: might not get annotation in this way
arguments.zip(argumentAliasOrUnique).forEach { (arg, varUnique) ->
val argShouldUnique = arg.hasAnnotation(context.uniqueId, context.session)
val passedUnique: Set<UniqueLevel> = when (varUnique) {
is Path -> data[NormalPath]?.get(varUnique.symbol) ?: setOf(UniqueLevel.Shared)
is Level -> varUnique.level
}
// TODO: more general comparison
if (argShouldUnique && passedUnique.any { it != UniqueLevel.Unique }) {
throw Exception("Functino Call Uniqueness Level not Match")
}
}
return data
}
Expand All @@ -122,22 +142,27 @@ class UniqueCFA(private val data: UniqueCheckerContext) : FirControlFlowChecker(
node: QualifiedAccessNode,
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(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
}
val dataForNode = visitNode(node, data)
val callee = node.fir.calleeReference.symbol as FirVariableSymbol
context.uniqueStack.last().add(Path(callee))
return dataForNode
}

override fun visitFunctionCallNode(
node: FunctionCallNode,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val uniqueLevel = mutableSetOf<UniqueLevel>()
// FIXME: might not be able to get the annotation for function declaration in this way
if (node.fir.hasAnnotation(this.context.uniqueId, this.context.session)) {
uniqueLevel.add(UniqueLevel.Unique)
} else {
uniqueLevel.add(UniqueLevel.Shared)
}
context.uniqueStack.last().add(Level(uniqueLevel))

// Function call does not change context
return data
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,21 @@ 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.fir.declarations.FirVariable
import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
import org.jetbrains.kotlin.name.ClassId

// TODO: find a better name for it
sealed class PathUnique
public data class Path(val symbol: FirVariableSymbol<FirVariable>) : PathUnique()
public data class Level(val level: Set<UniqueLevel>) : PathUnique()

interface UniqueCheckerContext {
val config: PluginConfiguration
val errorCollector: ErrorCollector
val session: FirSession
val uniqueId: ClassId
val nodeUniqueLevel: MutableMap<CFGNode<*>, MutableMap<CFGNode<*>, Set<UniqueLevel>>>
val assignArgumentPos: Map<CFGNode<*>, CFGNode<*>>
val uniqueStack:ArrayDeque<ArrayDeque<PathUnique>>

fun resolveUniqueAnnotation(declaration: FirDeclaration): UniqueLevel
}
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,6 @@ 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
import org.jetbrains.kotlin.name.Name
Expand All @@ -18,8 +16,7 @@ class UniqueCheckerData(
override val session: FirSession,
override val config: PluginConfiguration,
override val errorCollector: ErrorCollector,
override val nodeUniqueLevel: MutableMap<CFGNode<*>, MutableMap<CFGNode<*>, Set<UniqueLevel>>> = mutableMapOf(),
override val assignArgumentPos: Map<CFGNode<*>, CFGNode<*>>
override val uniqueStack: ArrayDeque<ArrayDeque<PathUnique>> = ArrayDeque(),
) : UniqueCheckerContext {

private fun getAnnotationId(name: String): ClassId =
Expand Down

0 comments on commit 358a1d5

Please sign in to comment.