CFA unique checker with borrow

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,167 @@
+/*
+ * 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 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 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))
+                    )
+                }.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 = visitNode(node, data)
+            val lSymbol = node.fir.symbol
+            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>> {
+            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) }
+        }
+
+        @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> = when (varUnique) {
+                    is Path -> data[NormalPath]?.get(varUnique.symbol) ?: setOf(UniqueLevel.Shared)
+                    is Level -> varUnique.level
+                }
+                // TODO: more general comparison
+                if (!passedUnique.lessThanOrEqual(argUniqueLevel)) {
+                    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.uniqueStack.last().add(Level(uniqueLevel))
+
+            // Function call does not change context
+            return data
+        }
+    }
+}

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

@@ -7,11 +7,19 @@ 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
+
+// TODO: find a better name for it
+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>>
 
     fun resolveUniqueAnnotation(declaration: FirDeclaration): UniqueLevel
 }

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,23 +12,32 @@ 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))
 
     private val uniqueId: ClassId
         get() = getAnnotationId("Unique")
+    private val sharedId: ClassId
+        get() = getAnnotationId("Shared")
+    private val borrowedId: ClassId
+        get() = getAnnotationId("Borrowed")
 
     override fun resolveUniqueAnnotation(declaration: FirDeclaration): UniqueLevel {
-        if (declaration.hasAnnotation(uniqueId, session)) {
-            return UniqueLevel.Unique
-        }
-        return UniqueLevel.Shared
+        return if (declaration.hasAnnotation(borrowedId, session) && declaration.hasAnnotation(sharedId, session)) {
+            UniqueLevel.BorrowedShared
+        } else if (declaration.hasAnnotation(borrowedId, session) && declaration.hasAnnotation(uniqueId, session)) {
+            UniqueLevel.BorrowedUnique
+        } else if (declaration.hasAnnotation(sharedId, session)) {
+            UniqueLevel.Shared
+        } else if (declaration.hasAnnotation(uniqueId, session)) {
+            UniqueLevel.Unique
+        } else UniqueLevel.Shared
     }
 }

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

@@ -5,9 +5,20 @@
 
 package org.jetbrains.kotlin.formver
 
-enum class UniqueLevel {
-    Bottom,
-    Unique,
-    Shared,
-    Top,
+enum class UniqueLevel(val level: Int) {
+    Bottom(0),
+    Unique(1),
+    BorrowedUnique(2),
+    Shared(2),
+    BorrowedShared(3),
+    Top(4),
 }
+
+fun UniqueLevel.lessThanOrEqual(other: UniqueLevel): Boolean = this.level < other.level || this == other
+
+fun UniqueLevel.leastUpperBound(other: UniqueLevel): UniqueLevel =
+    UniqueLevel.entries.first { this.lessThanOrEqual(it) && other.lessThanOrEqual(it) }
+
+fun Set<UniqueLevel>.leastUpperBound(): UniqueLevel = this.fold(UniqueLevel.Bottom) { acc, uniqueLevel -> acc.leastUpperBound(uniqueLevel) }
+
+fun Set<UniqueLevel>.lessThanOrEqual(uniqueLevel: UniqueLevel): Boolean = this.leastUpperBound().lessThanOrEqual(uniqueLevel)

plugins/formal-verification/testData/diagnostics/uniqueness/bad_unique/direct_pass_borrowed_unique_to_unique.kt

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

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_fun_call.kt

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