Create a LocalScope for CollectionComprehensions and generate a Declaration

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

@@ -242,6 +242,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/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

@@ -99,13 +99,25 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
         }
     }
 
+    /**
+     * Works similar to [apply] but before executing [block], it enters the scope for this object
+     * and afterward leaves the scope again.
+     */
+    private inline fun <T : Node> T.applyWithScope(block: T.() -> Unit): T {
+        return this.apply {
+            ctx?.scopeManager?.enterScope(this)
+            block()
+            ctx?.scopeManager?.leaveScope(this)
+        }
+    }
+
     /**
      * Translates a Python
      * [`GeneratorExp`](https://docs.python.org/3/library/ast.html#ast.GeneratorExp) into a
      * [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 +129,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 +141,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 +153,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,77 @@
 package de.fraunhofer.aisec.cpg.frontends.python
 
 import de.fraunhofer.aisec.cpg.graph.*
+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.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.InitializerListExpression
 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.test.assertNotRefersTo
+import de.fraunhofer.aisec.cpg.test.assertRefersTo
 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["tupleComp"]
+        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")
@@ -50,21 +107,42 @@ class ExpressionHandlerTest {
         assertNotNull(result)
         val listComp = result.functions["listComp"]
         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)

cpg-language-python/src/test/resources/python/comprehension.py

@@ -6,6 +6,7 @@ def listComp(x, y):
     b = [foo(i) for i in x]
     c = {foo(i) for i in x if i == 10 if i < 20}
     d = [foo(i) for z in y if z in x for i in z if i == 10 ]
+    foo(i)
 
 def setComp(x, y):
     a = {foo(i) for i in x if i == 10}
@@ -21,4 +22,10 @@ def dictComp(x, y):
 
 def generator(x, y):
     a = (i**2 for i in range(10) if i == 10)
-    b = (i**2 for i in range(10))
+    b = (i**2 for i in range(10))
+
+def bar(k, v):
+    return k+v
+
+def tupleComp(x):
+    a = [bar(k, v) for (k, v) in x]