create cfa checker for unique

plugins/formal-verification/formver.plugin/src/org/jetbrains/kotlin/formver/UniqueDeclarationChecker.kt

@@ -12,6 +12,8 @@ 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.controlFlowGraph
+import org.jetbrains.kotlin.text
 
 class UniqueDeclarationChecker(private val session: FirSession, private val config: PluginConfiguration) :
     FirSimpleFunctionChecker(MppCheckerKind.Common) {
@@ -20,9 +22,10 @@ class UniqueDeclarationChecker(private val session: FirSession, private val conf
         if (!config.checkUniqueness) return
         val errorCollector = ErrorCollector()
         try {
-            val uniqueCheckerContext = UniqueChecker(session, config, errorCollector)
-            declaration.accept(UniquenessCheckExceptionWrapper, uniqueCheckerContext)
+            val cfaChecker = UniqueCFA(UniqueCheckerData(session, config, errorCollector))
+            declaration.controlFlowGraphReference?.controlFlowGraph?.let { cfaChecker.analyze(it, reporter, context) }
         } catch (e: Exception) {
+            errorCollector.addErrorInfo("... while checking uniqueness level for ${declaration.source.text}")
             val error = errorCollector.formatErrorWithInfos(e.message ?: "No message provided")
             reporter.reportOn(declaration.source, PluginErrors.UNIQUENESS_VIOLATION, error, context)
         }

plugins/formal-verification/formver.uniqueness/build.gradle.kts

@@ -5,6 +5,7 @@ plugins {
 
 dependencies {
     compileOnly(project(":kotlin-formver-compiler-plugin.common"))
+    compileOnly(project(":compiler:fir:checkers"))
     compileOnly(project(":compiler:fir:cones"))
     compileOnly(project(":compiler:fir:tree"))
     compileOnly(project(":compiler:fir:plugin-utils"))

plugins/formal-verification/formver.uniqueness/src/org/jetbrains/kotlin/formver/UniqueCFA.kt

@@ -0,0 +1,161 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.formver
+
+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.expressions.toResolvedCallableSymbol
+import org.jetbrains.kotlin.fir.references.symbol
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+import org.jetbrains.kotlin.text
+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(this.data))
+    }
+
+    private class UniqueInfoCollector(private val context: UniqueCheckerContext) :
+        PathAwareControlFlowGraphVisitor<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>() {
+
+        override fun mergeInfo(
+            a: ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
+            b: ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
+            node: CFGNode<*>,
+        ): 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<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 valueParameters =
+                declaration.valueParameters.associate {
+                    Pair<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>(
+                        it.symbol,
+                        setOf(this.context.resolveUniqueAnnotation(it))
+                    )
+                }
+            return dataForNode.transformValues { it.putAll(valueParameters) }
+        }
+
+        override fun visitVariableDeclarationNode(
+            node: VariableDeclarationNode,
+            data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
+        ): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
+            val dataForNode = visitNode(node, data)
+            val lSymbol = node.fir.symbol
+            if (node.fir.initializer == null) {
+                return dataForNode.transformValues { it.put(lSymbol, setOf(UniqueLevel.Shared)) }
+            }
+            val rhsUniqueLevel = data.getUniqueLevel(context.getTopExprPathUnique())
+            return dataForNode.transformValues { it.put(lSymbol, rhsUniqueLevel) }
+        }
+
+        override fun visitVariableAssignmentNode(
+            node: VariableAssignmentNode,
+            data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
+        ): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
+            val dataForNode = visitNode(node, data)
+            val lSymbol = node.fir.lValue.toResolvedCallableSymbol(context.session)
+            val rhsUniqueLevel = data.getUniqueLevel(context.getTopExprPathUnique())
+            return dataForNode.transformValues { it.put(lSymbol as FirVariableSymbol, rhsUniqueLevel) }
+        }
+
+        @OptIn(SymbolInternals::class)
+        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
+            // Here do not call visitNode
+            val params = (node.fir.toResolvedCallableSymbol() as FirFunctionSymbol<*>).fir.valueParameters
+            val argumentAliasOrUnique: MutableList<PathUnique> = mutableListOf()
+            for (i in params.indices) {
+                argumentAliasOrUnique.add(context.uniqueStack.last().popLast())
+            }
+            // FIXME: might not get annotation in this way
+            params.zip(argumentAliasOrUnique).forEach { (param, varUnique) ->
+                val argUniqueLevel = this.context.resolveUniqueAnnotation(param)
+                val passedUnique: Set<UniqueLevel> = data.getUniqueLevel(varUnique)
+                // TODO: more general comparison
+                if (argUniqueLevel == UniqueLevel.Unique && passedUnique.any { it != UniqueLevel.Unique }) {
+                    throw IllegalArgumentException("uniqueness level not match ${param.source.text}")
+                }
+            }
+            return data
+        }
+
+        override fun visitQualifiedAccessNode(
+            node: QualifiedAccessNode,
+            data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
+        ): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
+            val dataForNode = visitNode(node, data)
+            val callee = node.fir.calleeReference.symbol as FirVariableSymbol
+            context.uniqueStack.last().add(Path(callee))
+            return dataForNode
+        }
+
+        @OptIn(SymbolInternals::class)
+        override fun visitFunctionCallExitNode(
+            node: FunctionCallExitNode,
+            data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
+        ): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
+            val uniqueLevel = mutableSetOf<UniqueLevel>()
+            val symbol = node.fir.calleeReference.symbol
+            if (symbol == null) {
+                uniqueLevel.add(UniqueLevel.Shared)
+            } else {
+                uniqueLevel.add(this.context.resolveUniqueAnnotation(symbol.fir))
+            }
+            context.pushExprPathUnique(Level(uniqueLevel))
+
+            // Function call does not change context
+            return data
+        }
+
+        fun PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>.getUniqueLevel(pathUnique: PathUnique?): Set<UniqueLevel> =
+            when (pathUnique) {
+                is Path -> this.getPathUniqueOrShared(pathUnique.symbol)
+                is Level -> pathUnique.level
+                else -> setOf(UniqueLevel.Shared)
+            }
+
+        fun PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>.getPathUniqueOrShared(
+            symbol: FirVariableSymbol<FirVariable>,
+        ): Set<UniqueLevel> {
+            return this[NormalPath]?.get(symbol) ?: setOf(UniqueLevel.Shared)
+        }
+    }
+}

plugins/formal-verification/formver.uniqueness/src/org/jetbrains/kotlin/formver/UniqueCheckerContext.kt

@@ -7,11 +7,35 @@ package org.jetbrains.kotlin.formver
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirVariable
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+
+/**
+ * An enum for either a variable or a unique level
+ * [Path] denotes a possible variable, while [Level] is for a unique level
+ * This type is used for the stack for control-flow visitor
+ */
+sealed class PathUnique
+data class Path(val symbol: FirVariableSymbol<FirVariable>) : PathUnique()
+data class Level(val level: Set<UniqueLevel>) : PathUnique()
 
 interface UniqueCheckerContext {
     val config: PluginConfiguration
     val errorCollector: ErrorCollector
     val session: FirSession
+    val uniqueStack: ArrayDeque<ArrayDeque<PathUnique>>
 
+    /**
+     * Resolve the annotation for borrow and uniqueness
+     */
     fun resolveUniqueAnnotation(declaration: FirDeclaration): UniqueLevel
+
+    /**
+     * Pushes the alias or unique level of the last expression onto the top of last stack.
+     *
+     * @param pathUnique the unique path to be added to the stack
+     */
+    fun pushExprPathUnique(pathUnique: PathUnique)
+
+    fun getTopExprPathUnique(): PathUnique?
 }

plugins/formal-verification/formver.uniqueness/src/org/jetbrains/kotlin/formver/UniqueChecker.kt → plugins/formal-verification/formver.uniqueness/src/org/jetbrains/kotlin/formver/UniqueCheckerData.kt

@@ -12,12 +12,12 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
-class UniqueChecker(
+class UniqueCheckerData(
     override val session: FirSession,
     override val config: PluginConfiguration,
     override val errorCollector: ErrorCollector,
-) :
-    UniqueCheckerContext {
+    override val uniqueStack: ArrayDeque<ArrayDeque<PathUnique>> = ArrayDeque(),
+) : UniqueCheckerContext {
 
     private fun getAnnotationId(name: String): ClassId =
         ClassId(FqName.fromSegments(listOf("org", "jetbrains", "kotlin", "formver", "plugin")), Name.identifier(name))
@@ -31,4 +31,10 @@ class UniqueChecker(
         }
         return UniqueLevel.Shared
     }
+
+    override fun pushExprPathUnique(pathUnique: PathUnique) {
+        this.uniqueStack.last().add(pathUnique)
+    }
+
+    override fun getTopExprPathUnique(): PathUnique? = this.uniqueStack.last().lastOrNull()
 }

plugins/formal-verification/testData/diagnostics/uniqueness/bad_unique/direct_pass_shared_to_unique.fir.diag.txt

@@ -1,4 +1,4 @@
-/direct_pass_shared_to_unique.kt:(212,242): error: uniqueness level not match y
+/direct_pass_shared_to_unique.kt:(212,242): error: uniqueness level not match @Unique x: Int
 ... while checking uniqueness level for fun use_f(y: Int) {
     f(y)
 }

plugins/formal-verification/testData/diagnostics/uniqueness/good_unique/unique_to_unique.kt

@@ -0,0 +1,13 @@
+// UNIQUE_CHECK_ONLY
+
+import org.jetbrains.kotlin.formver.plugin.NeverConvert
+import org.jetbrains.kotlin.formver.plugin.NeverVerify
+import org.jetbrains.kotlin.formver.plugin.Unique
+
+fun f(@Unique x: Int) {
+
+}
+
+fun use_f(@Unique y: Int) {
+    f(y)
+}