Handle InitializerListExpressions as variables of ComprehensionExpressions in the DFG

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

@@ -392,40 +392,7 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
                 state.push(currentNode, it)
             }
         } else if (currentNode is ComprehensionExpression) {
-            val writtenTo =
-                when (val variable = currentNode.variable) {
-                    is DeclarationStatement -> {
-                        if (variable.isSingleDeclaration()) {
-                            variable.singleDeclaration
-                        } else {
-                            log.error(
-                                "Cannot handle multiple declarations in the ComprehensionExpresdsion: Node $currentNode"
-                            )
-                            null
-                        }
-                    }
-                    else -> currentNode.variable
-                }
-            // We wrote something to this variable declaration
-            writtenTo?.let {
-                writtenDeclaration =
-                    when (writtenTo) {
-                        is Declaration -> writtenTo
-                        is Reference -> writtenTo.refersTo
-                        else -> {
-                            log.error(
-                                "The variable of type ${writtenTo.javaClass} is not yet supported in the ComprehensionExpression"
-                            )
-                            null
-                        }
-                    }
-
-                state.push(writtenTo, PowersetLattice(identitySetOf(currentNode.iterable)))
-                // Add the variable declaration (or the reference) to the list of previous
-                // write nodes in this path
-                state.declarationsState[writtenDeclaration] =
-                    PowersetLattice(identitySetOf(writtenTo))
-            }
+            handleComprehensionExpression(currentNode, doubleState)
         } else if (currentNode is ForEachStatement && currentNode.variable != null) {
             // The VariableDeclaration in the ForEachStatement doesn't have an initializer, so
             // the "normal" case won't work. We handle this case separately here...
@@ -538,6 +505,51 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
         return state
     }
 
+    /**
+     * Handles the propagation of data flows to the variables used in a [ComprehensionExpression].
+     * We have a write access to one or multiple [Declaration]s or [Reference]s here. Multiple
+     * values are supported through [InitializerListExpression].
+     */
+    protected fun handleComprehensionExpression(
+        currentNode: ComprehensionExpression,
+        state: DFGPassState<Set<Node>>,
+    ) {
+        val writtenTo =
+            when (val variable = currentNode.variable) {
+                is DeclarationStatement -> {
+                    variable.declarations
+                }
+                is Reference -> listOf(variable)
+                is InitializerListExpression -> variable.initializers
+                else -> {
+                    log.error(
+                        "The type ${variable.javaClass} is not yet supported as ComprehensionExpression::variable"
+                    )
+                    listOf()
+                }
+            }
+        // We wrote something to this variable declaration
+        writtenTo.forEach { writtenToIt ->
+            val writtenDeclaration =
+                when (writtenToIt) {
+                    is Declaration -> writtenToIt
+                    is Reference -> writtenToIt.refersTo
+                    else -> {
+                        log.error(
+                            "The variable of type ${writtenToIt.javaClass} is not yet supported in the ComprehensionExpression"
+                        )
+                        null
+                    }
+                }
+
+            state.push(writtenToIt, PowersetLattice(identitySetOf(currentNode.iterable)))
+            // Add the variable declaration (or the reference) to the list of previous
+            // write nodes in this path
+            state.declarationsState[writtenDeclaration] =
+                PowersetLattice(identitySetOf(writtenToIt))
+        }
+    }
+
     /**
      * We use this map to store additional information on the DFG edges which we cannot keep in the
      * state. This is for example the case to identify if the resulting edge will receive a

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

@@ -396,12 +396,13 @@ class DFGPass(ctx: TranslationContext) : ComponentPass(ctx) {
      */
     protected fun handleInitializerListExpression(node: InitializerListExpression) {
         node.initializers.forEachIndexed { idx, it ->
+            val astParent = node.astParent
             if (
-                node.astParent is AssignExpression &&
-                    node in (node.astParent as AssignExpression).lhs
+                astParent is AssignExpression && node in astParent.lhs ||
+                    astParent is ComprehensionExpression && node == astParent.variable
             ) {
-                // If we're the target of an assignment, the DFG flows from the node to the
-                // initializers.
+                // If we're the target of an assignment or the variable of a comprehension
+                // expression, the DFG flows from the node to the initializers.
                 node.nextDFGEdges.add(it) { granularity = indexed(idx) }
             } else {
                 node.prevDFGEdges.add(it) { granularity = indexed(idx) }

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

@@ -26,11 +26,13 @@
 package de.fraunhofer.aisec.cpg.frontends.python
 
 import de.fraunhofer.aisec.cpg.graph.*
+import de.fraunhofer.aisec.cpg.graph.edges.flows.IndexedDataflowGranularity
 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
@@ -40,6 +42,61 @@ 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 ILE flows to the variables with the indexed granularity
+        assertContains(initializerListExpression.nextDFG, variableK)
+        val granularityTupleToK =
+            initializerListExpression.nextDFGEdges.single { it.end == variableK }.granularity
+        assertIs<IndexedDataflowGranularity>(granularityTupleToK)
+        assertEquals(0, granularityTupleToK.index)
+        assertContains(initializerListExpression.nextDFG, variableV)
+        val granularityTupleToV =
+            initializerListExpression.nextDFGEdges.single { it.end == variableV }.granularity
+        assertIs<IndexedDataflowGranularity>(granularityTupleToV)
+        assertEquals(1, granularityTupleToV.index)
+
+        // Check that the variables flow to their usages
+        assertEquals(setOf<Node>(argK), variableK.nextDFG.toSet())
+        assertEquals(setOf<Node>(argV), variableV.nextDFG.toSet())
+    }
+
     @Test
     fun testListComprehensions() {
         val topLevel = Path.of("src", "test", "resources", "python")

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

@@ -21,4 +21,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]