Skip to content

Commit

Permalink
track assignement for declaration from CFA
Browse files Browse the repository at this point in the history
  • Loading branch information
Eric-Song-Nop committed Jul 30, 2024
1 parent 2b820fb commit 6d0e46e
Show file tree
Hide file tree
Showing 5 changed files with 139 additions and 122 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ 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) {
Expand All @@ -22,12 +21,13 @@ class UniqueDeclarationChecker(private val session: FirSession, private val conf
if (!config.checkUniqueness) return
val errorCollector = ErrorCollector()
try {
val cfaChecker = UniqueCFA(UniqueCheckerData(session, config, errorCollector))
val uniqueCheckerContext = UniqueCheckerData(session, config, errorCollector)

val cfaChecker = UniqueCFA(uniqueCheckerContext)
declaration.controlFlowGraphReference?.controlFlowGraph?.let { cfaChecker.analyze(it, reporter, context) }
} catch (e: Exception) {
errorCollector.addErrorInfo("... while checking uniqueness level for ${declaration.source.text}")
val errorMsg = errorCollector.formatErrorWithInfos(e.message ?: "No message provided")
reporter.reportOn(declaration.source, PluginErrors.UNIQUENESS_VIOLATION, errorMsg, context)
val error = errorCollector.formatErrorWithInfos(e.message ?: "No message provided")
reporter.reportOn(declaration.source, PluginErrors.UNIQUENESS_VIOLATION, error, context)
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,164 +5,100 @@

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) {
sealed class LastOrPath {
data object LastNode : LastOrPath()
data class Path(val path: FirVariableSymbol<FirVariable>) : LastOrPath()
}

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

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

private class UniqueInfoCollector(private val data: UniqueCheckerData) :
PathAwareControlFlowGraphVisitor<LastOrPath, Set<UniqueLevel>>() {
override fun mergeInfo(
a: ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
b: ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
a: ControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
b: ControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
node: CFGNode<*>,
): ControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
): ControlFlowInfo<LastOrPath, 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
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 visitFunctionEnterNode(
node: FunctionEnterNode,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val dataForNode = visitNode(node, data)
data: PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
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.context.uniqueId, this.context.session)) {
Pair<LastOrPath, Set<UniqueLevel>>(
LastOrPath.Path(it.symbol),
if (it.hasAnnotation(this.data.uniqueId, this.data.session)) {
setOf(UniqueLevel.Unique)
} else {
setOf(UniqueLevel.Shared)
}
)
}.toMap()
return dataForNode.transformValues { it.putAll(valueParameters) }
}.toPersistentMap()
return super.visitFunctionEnterNode(node, persistentMapOf(Pair(NormalPath, valueParameters)))
}

override fun visitVariableDeclarationNode(
node: VariableDeclarationNode,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
val dataForNode = visitNode(node, data)
data: PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
val dataForNode = super.visitVariableDeclarationNode(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) }
}

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 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
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
}

override fun visitQualifiedAccessNode(
node: QualifiedAccessNode,
data: PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>>,
): PathAwareControlFlowInfo<FirVariableSymbol<FirVariable>, Set<UniqueLevel>> {
data: PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>>
): PathAwareControlFlowInfo<LastOrPath, Set<UniqueLevel>> {
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)
if (node.fir.calleeReference is FirResolvedNamedReference) {
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) }
}
context.uniqueStack.last().add(Level(uniqueLevel))

// Function call does not change context
return data
return super.visitQualifiedAccessNode(node, data)
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
/*
* 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.fir.FirElement
import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
import org.jetbrains.kotlin.fir.expressions.FirBlock
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.expressions.arguments
import org.jetbrains.kotlin.fir.expressions.toResolvedCallableSymbol
import org.jetbrains.kotlin.fir.symbols.SymbolInternals
import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
import org.jetbrains.kotlin.fir.visitors.FirVisitor
import org.jetbrains.kotlin.text

/**
* The UniqueCheckVisitor class is responsible for visiting expression and return its unique level
*
* @see FirVisitor
* @see UniqueLevel
*
* @property uniqueLevel The uniqueness level to be of the expression
* @property uniqueCheckerContext The context object that contains information about the uniqueness annotations and resolution.
*/
object UniqueCheckVisitor : FirVisitor<UniqueLevel, UniqueCheckerContext>() {
override fun visitElement(element: FirElement, data: UniqueCheckerContext): UniqueLevel = UniqueLevel.Shared

override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: UniqueCheckerContext): UniqueLevel {
simpleFunction.body?.accept(this, data)
// Function definition don't have to return a unique level
return UniqueLevel.Shared
}

override fun visitBlock(block: FirBlock, data: UniqueCheckerContext): UniqueLevel {
block.statements.forEach { statement ->
when (statement) {
is FirFunctionCall -> {
statement.accept(this, data)
}
}
}
return UniqueLevel.Shared
}

@OptIn(SymbolInternals::class)
override fun visitFunctionCall(functionCall: FirFunctionCall, data: UniqueCheckerContext): UniqueLevel {
// To keep is simple, assume a functionCall always return Shared for now
val symbol = functionCall.toResolvedCallableSymbol()
val params = (symbol as FirFunctionSymbol<*>).fir.valueParameters
val requiredUniqueLevels = params.map { data.resolveUniqueAnnotation(it) }
// Skip merge of context for now
val arguments = functionCall.arguments
arguments.forEachIndexed { index, argument ->
val requiredUnique = requiredUniqueLevels[index]
if (requiredUnique == UniqueLevel.Unique && visitExpression(argument, data) == UniqueLevel.Shared) {
throw IllegalArgumentException("uniqueness level not match ${argument.source.text}")
}
}

val callee = functionCall.toResolvedCallableSymbol()?.fir as FirSimpleFunction
return data.resolveUniqueAnnotation(callee)
}
}

object UniquenessCheckExceptionWrapper : FirVisitor<UniqueLevel, UniqueCheckerContext>() {
override fun visitElement(element: FirElement, data: UniqueCheckerContext): UniqueLevel {
try {
return element.accept(UniqueCheckVisitor, data)
} catch (e: Exception) {
data.errorCollector.addErrorInfo("... while checking uniqueness level for ${element.source.text}")
throw e
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,16 @@ 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
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
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 uniqueStack:ArrayDeque<ArrayDeque<PathUnique>>

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.symbols.impl.FirPropertySymbol
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
Expand All @@ -16,7 +17,7 @@ class UniqueCheckerData(
override val session: FirSession,
override val config: PluginConfiguration,
override val errorCollector: ErrorCollector,
override val uniqueStack: ArrayDeque<ArrayDeque<PathUnique>> = ArrayDeque(),
public val pathUniqueTable: MutableMap<FirPropertySymbol, Set<UniqueLevel>> = mutableMapOf(),
) : UniqueCheckerContext {

private fun getAnnotationId(name: String): ClassId =
Expand All @@ -31,4 +32,12 @@ 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 6d0e46e

Please sign in to comment.