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-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 = objectType("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 = objectType("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 = objectType("dict")
         }
     }
 

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

@@ -37,6 +37,7 @@ 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.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 +68,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 +195,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.

cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandlerTest.kt

@@ -26,20 +26,67 @@
 package de.fraunhofer.aisec.cpg.frontends.python
 
 import de.fraunhofer.aisec.cpg.graph.*
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.AssignExpression
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.BinaryOperator
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.Block
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.CallExpression
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.CollectionComprehension
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.KeyValueExpression
-import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
-import de.fraunhofer.aisec.cpg.test.analyze
-import de.fraunhofer.aisec.cpg.test.assertLiteralValue
-import de.fraunhofer.aisec.cpg.test.assertLocalName
+import de.fraunhofer.aisec.cpg.graph.declarations.FunctionDeclaration
+import de.fraunhofer.aisec.cpg.graph.declarations.ParameterDeclaration
+import de.fraunhofer.aisec.cpg.graph.declarations.VariableDeclaration
+import de.fraunhofer.aisec.cpg.graph.scopes.LocalScope
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
+import de.fraunhofer.aisec.cpg.test.*
 import java.nio.file.Path
 import kotlin.test.*
 
 class ExpressionHandlerTest {
+
+    @Test
+    fun testComprehensionExpressionTuple() {
+        val topLevel = Path.of("src", "test", "resources", "python")
+        val result =
+            analyze(listOf(topLevel.resolve("comprehension.py").toFile()), topLevel, true) {
+                it.registerLanguage<PythonLanguage>()
+            }
+        assertNotNull(result)
+
+        val tupleComp = result.functions["tuple_comp"]
+        assertNotNull(tupleComp)
+
+        val body = tupleComp.body
+        assertIs<Block>(body)
+        val tupleAsVariableAssignment = body.statements[0]
+        assertIs<AssignExpression>(tupleAsVariableAssignment)
+        val tupleAsVariable = tupleAsVariableAssignment.rhs[0]
+        assertIs<CollectionComprehension>(tupleAsVariable)
+        val barCall = tupleAsVariable.statement
+        assertIs<CallExpression>(barCall)
+        assertLocalName("bar", barCall)
+        val argK = barCall.arguments[0]
+        assertIs<Reference>(argK)
+        assertLocalName("k", argK)
+        val argV = barCall.arguments[1]
+        assertIs<Reference>(argV)
+        assertLocalName("v", argV)
+        assertEquals(1, tupleAsVariable.comprehensionExpressions.size)
+        val initializerListExpression = tupleAsVariable.comprehensionExpressions[0].variable
+        assertIs<InitializerListExpression>(initializerListExpression)
+        val variableK = initializerListExpression.initializers[0]
+        assertIs<Reference>(variableK)
+        assertLocalName("k", variableK)
+        val variableV = initializerListExpression.initializers[1]
+        assertIs<Reference>(variableV)
+        assertLocalName("v", variableV)
+
+        // Check that the declarations exist for the variables k and v
+        val declK = variableK.refersTo
+        assertIs<VariableDeclaration>(declK)
+        assertIs<LocalScope>(declK.scope)
+        assertEquals(tupleAsVariable, declK.scope?.astNode)
+        assertRefersTo(argK, declK)
+        val declV = variableV.refersTo
+        assertIs<VariableDeclaration>(declV)
+        assertIs<LocalScope>(declV.scope)
+        assertEquals(tupleAsVariable, declV.scope?.astNode)
+        assertRefersTo(argV, declV)
+    }
+
     @Test
     fun testListComprehensions() {
         val topLevel = Path.of("src", "test", "resources", "python")
@@ -48,23 +95,44 @@ class ExpressionHandlerTest {
                 it.registerLanguage<PythonLanguage>()
             }
         assertNotNull(result)
-        val listComp = result.functions["listComp"]
+        val listComp = result.functions["list_comp"]
         assertNotNull(listComp)
+        val paramX = listComp.parameters[0]
+        assertIs<ParameterDeclaration>(paramX)
+        assertLocalName("x", paramX)
 
         val body = listComp.body
         assertIs<Block>(body)
         val singleWithIfAssignment = body.statements[0]
         assertIs<AssignExpression>(singleWithIfAssignment)
         val singleWithIf = singleWithIfAssignment.rhs[0]
         assertIs<CollectionComprehension>(singleWithIf)
-        assertIs<CallExpression>(singleWithIf.statement)
+        val fooCall = singleWithIf.statement
+        assertIs<CallExpression>(fooCall)
+        val usageI = fooCall.arguments[0]
+        assertIs<Reference>(usageI)
         assertEquals(1, singleWithIf.comprehensionExpressions.size)
-        assertLocalName("i", singleWithIf.comprehensionExpressions[0].variable)
-        assertIs<Reference>(singleWithIf.comprehensionExpressions[0].iterable)
-        assertLocalName("x", singleWithIf.comprehensionExpressions[0].iterable)
+        val variableI = singleWithIf.comprehensionExpressions[0].variable
+        assertIs<Reference>(variableI)
+        assertLocalName("i", variableI)
+        val declI = variableI.refersTo
+        assertIs<VariableDeclaration>(declI)
+        assertEquals(singleWithIf, declI.scope?.astNode)
+        val iterableX = singleWithIf.comprehensionExpressions[0].iterable
+        assertIs<Reference>(iterableX)
+        assertLocalName("x", iterableX)
+        assertRefersTo(iterableX, paramX)
         val ifPredicate = singleWithIf.comprehensionExpressions[0].predicate
         assertIs<BinaryOperator>(ifPredicate)
         assertEquals("==", ifPredicate.operatorCode)
+        assertRefersTo(usageI, declI)
+
+        val fooIOutside = body.statements[4]
+        assertIs<CallExpression>(fooIOutside)
+        val outsideI = fooIOutside.arguments[0]
+        assertIs<Reference>(outsideI)
+        assertLocalName("i", outsideI)
+        assertNotRefersTo(outsideI, declI)
 
         val singleWithoutIfAssignment = body.statements[1]
         assertIs<AssignExpression>(singleWithoutIfAssignment)
@@ -107,7 +175,7 @@ class ExpressionHandlerTest {
                 it.registerLanguage<PythonLanguage>()
             }
         assertNotNull(result)
-        val listComp = result.functions["setComp"]
+        val listComp = result.functions["set_comp"]
         assertNotNull(listComp)
 
         val body = listComp.body as? Block
@@ -165,7 +233,7 @@ class ExpressionHandlerTest {
                 it.registerLanguage<PythonLanguage>()
             }
         assertNotNull(result)
-        val listComp = result.functions["dictComp"]
+        val listComp = result.functions["dict_comp"]
         assertNotNull(listComp)
 
         val body = listComp.body as? Block
@@ -269,6 +337,86 @@ class ExpressionHandlerTest {
         assertNull(singleWithoutIf.comprehensionExpressions[0].predicate)
     }
 
+    @Test
+    /**
+     * This test ensures that variables in a comprehension do not bind to the outer scope. See
+     * [testCompBindingAssignExpr] for exceptions.
+     */
+    fun testCompBinding() {
+        val topLevel = Path.of("src", "test", "resources", "python")
+        val result =
+            analyze(listOf(topLevel.resolve("comprehension.py").toFile()), topLevel, true) {
+                it.registerLanguage<PythonLanguage>()
+            }
+        assertNotNull(result)
+        val compBindingFunc = result.functions["comp_binding"]
+        assertIs<FunctionDeclaration>(compBindingFunc)
+
+        val xDecl = compBindingFunc.variables.firstOrNull()
+        assertIs<VariableDeclaration>(xDecl)
+
+        assertEquals(
+            5,
+            compBindingFunc.variables.size,
+            "Expected 5 variables. One for the \"outside\" x and one for each of the 4 comprehensions.",
+        )
+
+        assertEquals(
+            2,
+            xDecl.usages.size,
+            "Expected two usages: one for the initial assignment and one for the usage in \"print(x)\".",
+        )
+
+        val comprehensions =
+            compBindingFunc.body.statements.filterIsInstance<CollectionComprehension>()
+        assertEquals(4, comprehensions.size, "Expected to find 4 comprehensions.")
+
+        comprehensions.forEach { it.refs.forEach { ref -> assertNotRefersTo(ref, xDecl) } }
+    }
+
+    @Test
+    /**
+     * This test ensures that variables in a comprehension do not bind to the outer scope if they
+     * are used in an `AssignExpr`. See https://peps.python.org/pep-0572/#scope-of-the-target
+     */
+    fun testCompBindingAssignExpr() {
+        val topLevel = Path.of("src", "test", "resources", "python")
+        val result =
+            analyze(listOf(topLevel.resolve("comprehension.py").toFile()), topLevel, true) {
+                it.registerLanguage<PythonLanguage>()
+            }
+        assertNotNull(result)
+
+        val compBindingAssignFunc = result.functions["comp_binding_assign_expr"]
+        assertIs<FunctionDeclaration>(compBindingAssignFunc)
+
+        val xDecl = compBindingAssignFunc.variables.firstOrNull()
+        assertIs<VariableDeclaration>(xDecl)
+
+        assertEquals(
+            2,
+            compBindingAssignFunc.variables.size,
+            "Expected 2 variables. One for the \"outside\" x and one for the \"temp\" inside the comprehension.",
+        )
+
+        assertEquals(
+            3,
+            xDecl.usages.size,
+            "Expected 3 usages: one for the initial assignment, one for the comprehension and one for the usage in \"print(x)\".",
+        )
+
+        val comprehension =
+            compBindingAssignFunc.body.statements.singleOrNull { it is CollectionComprehension }
+        assertNotNull(comprehension)
+        val xRef = comprehension.refs("x").singleOrNull()
+        assertNotNull(xRef)
+        assertRefersTo(xRef, xDecl)
+
+        // TODO: test for other types of comprehensions
+        // TODO: test nested comprehensions -> AssignExpression binds to
+        // containing scope
+    }
+
     @Test
     fun testBoolOps() {
         val topLevel = Path.of("src", "test", "resources", "python")