Create a LocalScope for CollectionComprehensions and generate a Declaration

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/ScopeManager.kt

@@ -240,6 +240,7 @@ class ScopeManager : ScopeProvider {
                     is TryStatement,
                     is IfStatement,
                     is CatchClause,
+                    is CollectionComprehension,
                     is Block -> LocalScope(nodeToScope)
                     is FunctionDeclaration -> FunctionScope(nodeToScope)
                     is RecordDeclaration -> RecordScope(nodeToScope)

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/Node.kt

@@ -369,3 +369,15 @@ abstract class Node :
         const val EMPTY_NAME = ""
     }
 }
+
+/**
+ * Works similar to [apply] but before executing [block], it enters the scope for this object and
+ * afterward leaves the scope again.
+ */
+inline fun <reified T : Node> T.applyWithScope(block: T.() -> Unit): T {
+    return this.apply {
+        ctx?.scopeManager?.enterScope(this)
+        block()
+        ctx?.scopeManager?.leaveScope(this)
+    }
+}

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/TypeBuilder.kt

@@ -101,7 +101,7 @@ fun LanguageProvider.objectType(
 ): Type {
     // First, we check, whether this is a built-in type, to avoid necessary allocations of simple
     // types
-    val builtIn = language?.getSimpleTypeOf(name.toString())
+    val builtIn = language.getSimpleTypeOf(name.toString())
     if (builtIn != null) {
         return builtIn
     }

cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/helpers/ExtensionsTest.kt

@@ -28,18 +28,27 @@ package de.fraunhofer.aisec.cpg.helpers
 import de.fraunhofer.aisec.cpg.GraphExamples.Companion.testFrontend
 import de.fraunhofer.aisec.cpg.TranslationConfiguration
 import de.fraunhofer.aisec.cpg.frontends.TestLanguage
+import de.fraunhofer.aisec.cpg.graph.*
+import de.fraunhofer.aisec.cpg.graph.Name
+import de.fraunhofer.aisec.cpg.graph.applyWithScope
 import de.fraunhofer.aisec.cpg.graph.builder.body
 import de.fraunhofer.aisec.cpg.graph.builder.declare
 import de.fraunhofer.aisec.cpg.graph.builder.function
 import de.fraunhofer.aisec.cpg.graph.builder.problemDecl
 import de.fraunhofer.aisec.cpg.graph.builder.translationResult
 import de.fraunhofer.aisec.cpg.graph.builder.translationUnit
+import de.fraunhofer.aisec.cpg.graph.declarations.VariableDeclaration
 import de.fraunhofer.aisec.cpg.graph.problems
+import de.fraunhofer.aisec.cpg.graph.statements.DeclarationStatement
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.CollectionComprehension
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.ProblemExpression
+import de.fraunhofer.aisec.cpg.graph.variables
 import de.fraunhofer.aisec.cpg.test.BaseTest
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertIs
 import kotlin.test.assertNotNull
+import kotlin.test.assertNull
 
 internal class ExtensionsTest : BaseTest() {
     val problemDeclText = "This is a problem declaration."
@@ -75,4 +84,19 @@ internal class ExtensionsTest : BaseTest() {
             "Failed to find the problem expression.",
         )
     }
+
+    @Test
+    fun testApplyWithScopeWithoutCtxAndScopeManager() {
+        val collectionComprehension =
+            CollectionComprehension().applyWithScope {
+                val varA = VariableDeclaration()
+                varA.name = Name("a")
+                val declarationStatement = DeclarationStatement()
+                declarationStatement.addDeclaration(varA)
+                this.statement = declarationStatement
+            }
+        val varA = collectionComprehension.variables["a"]
+        assertIs<VariableDeclaration>(varA)
+        assertNull(varA.scope)
+    }
 }

cpg-core/src/testFixtures/kotlin/de/fraunhofer/aisec/cpg/test/TestUtils.kt

@@ -223,18 +223,18 @@ fun compareLineFromLocationIfExists(n: Node, startLine: Boolean, toCompare: Int)
 }
 
 /** Asserts, that the expression given in [expression] refers to the expected declaration [b]. */
-fun assertRefersTo(expression: Expression?, b: Declaration?) {
+fun assertRefersTo(expression: Expression?, b: Declaration?, message: String? = null) {
     if (expression is Reference) {
-        assertEquals(b, (expression as Reference?)?.refersTo)
+        assertEquals(b, (expression as Reference?)?.refersTo, message)
     } else {
         fail("not a reference")
     }
 }
 
 /** Asserts, that the expression given in [expression] does not refer to the declaration [b]. */
-fun assertNotRefersTo(expression: Expression?, b: Declaration?) {
+fun assertNotRefersTo(expression: Expression?, b: Declaration?, message: String? = null) {
     if (expression is Reference) {
-        assertNotEquals(b, (expression as Reference?)?.refersTo)
+        assertNotEquals(b, (expression as Reference?)?.refersTo, message)
     } else {
         fail("not a reference")
     }

cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt

@@ -105,7 +105,7 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
      * [CollectionComprehension].
      */
     private fun handleGeneratorExp(node: Python.AST.GeneratorExp): CollectionComprehension {
-        return newCollectionComprehension(rawNode = node).apply {
+        return newCollectionComprehension(rawNode = node).applyWithScope {
             statement = handle(node.elt)
             comprehensionExpressions += node.generators.map { handleComprehension(it, node) }
             type = objectType("Generator")
@@ -117,10 +117,10 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
      * into a [CollectionComprehension].
      */
     private fun handleListComprehension(node: Python.AST.ListComp): CollectionComprehension {
-        return newCollectionComprehension(rawNode = node).apply {
+        return newCollectionComprehension(rawNode = node).applyWithScope {
             statement = handle(node.elt)
             comprehensionExpressions += node.generators.map { handleComprehension(it, node) }
-            type = objectType("list") // TODO: Replace this once we have dedicated types
+            type = primitiveType("list")
         }
     }
 
@@ -129,10 +129,10 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
      * a [CollectionComprehension].
      */
     private fun handleSetComprehension(node: Python.AST.SetComp): CollectionComprehension {
-        return newCollectionComprehension(rawNode = node).apply {
+        return newCollectionComprehension(rawNode = node).applyWithScope {
             this.statement = handle(node.elt)
             this.comprehensionExpressions += node.generators.map { handleComprehension(it, node) }
-            this.type = objectType("set") // TODO: Replace this once we have dedicated types
+            this.type = primitiveType("set")
         }
     }
 
@@ -141,15 +141,15 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
      * into a [CollectionComprehension].
      */
     private fun handleDictComprehension(node: Python.AST.DictComp): CollectionComprehension {
-        return newCollectionComprehension(rawNode = node).apply {
+        return newCollectionComprehension(rawNode = node).applyWithScope {
             this.statement =
                 newKeyValueExpression(
                     key = handle(node.key),
                     value = handle(node.value),
                     rawNode = node,
                 )
             this.comprehensionExpressions += node.generators.map { handleComprehension(it, node) }
-            this.type = objectType("dict") // TODO: Replace this once we have dedicated types
+            this.type = primitiveType("dict")
         }
     }
 

cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/PythonAddDeclarationsPass.kt

@@ -37,6 +37,8 @@ import de.fraunhofer.aisec.cpg.graph.declarations.VariableDeclaration
 import de.fraunhofer.aisec.cpg.graph.scopes.RecordScope
 import de.fraunhofer.aisec.cpg.graph.statements.ForEachStatement
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.AssignExpression
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.CollectionComprehension
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.ComprehensionExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.InitializerListExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
@@ -67,15 +69,15 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx), L
     }
 
     /**
-     * This function checks for each [AssignExpression] whether there is already a matching variable
-     * or not. New variables can be one of:
+     * This function checks for each [AssignExpression], [ComprehensionExpression] and
+     * [ForEachStatement] whether there is already a matching variable or not. New variables can be
+     * one of:
      * - [FieldDeclaration] if we are currently in a record
      * - [VariableDeclaration] otherwise
-     *
-     * TODO: loops
      */
     private fun handle(node: Node?) {
         when (node) {
+            is ComprehensionExpression -> handleComprehensionExpression(node)
             is AssignExpression -> handleAssignExpression(node)
             is ForEachStatement -> handleForEach(node)
             else -> {
@@ -194,6 +196,27 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx), L
                 this.base.name == scopeManager.currentMethod?.receiver?.name
         }
 
+    /**
+     * Generates a new [VariableDeclaration] for [Reference] (and those included in a
+     * [InitializerListExpression]) in the [ComprehensionExpression.variable].
+     */
+    private fun handleComprehensionExpression(comprehensionExpression: ComprehensionExpression) {
+        when (val variable = comprehensionExpression.variable) {
+            is Reference -> {
+                variable.access = AccessValues.WRITE
+                handleWriteToReference(variable)
+            }
+            is InitializerListExpression -> {
+                variable.initializers.forEach {
+                    (it as? Reference)?.let { ref ->
+                        ref.access = AccessValues.WRITE
+                        handleWriteToReference(ref)
+                    }
+                }
+            }
+        }
+    }
+
     /**
      * Generates a new [VariableDeclaration] if [target] is a [Reference] and there is no existing
      * declaration yet.
@@ -223,6 +246,15 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx), L
     }
 
     private fun handleAssignExpression(assignExpression: AssignExpression) {
+        val parentCollectionComprehensions = ArrayDeque<CollectionComprehension>()
+        var parentCollectionComprehension =
+            assignExpression.firstParentOrNull<CollectionComprehension>()
+        while (assignExpression.operatorCode == ":=" && parentCollectionComprehension != null) {
+            scopeManager.leaveScope(parentCollectionComprehension)
+            parentCollectionComprehensions.addLast(parentCollectionComprehension)
+            parentCollectionComprehension =
+                parentCollectionComprehension.firstParentOrNull<CollectionComprehension>()
+        }
         for (target in assignExpression.lhs) {
             handleAssignmentToTarget(assignExpression, target, setAccessValue = false)
             // If the lhs is an InitializerListExpression, we have to handle the individual elements
@@ -233,6 +265,11 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx), L
                 }
             }
         }
+        while (
+            assignExpression.operatorCode == ":=" && parentCollectionComprehensions.isNotEmpty()
+        ) {
+            scopeManager.enterScope(parentCollectionComprehensions.removeLast())
+        }
     }
 
     // New variables can also be declared as `variable` in a [ForEachStatement]