Support tuple deconstruction and follow value flows based on index

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/flows/Dataflow.kt

@@ -63,6 +63,15 @@ class PartialDataflowGranularity(
     val partialTarget: Declaration?
 ) : Granularity
 
+/**
+ * This dataflow granularity denotes that not the "whole" object is flowing from [Dataflow.start] to
+ * [Dataflow.end] but only parts of it. Common examples include tuples or array indices.
+ */
+class IndexedDataflowGranularity(
+    /** The index that is affected by this partial dataflow. */
+    val index: Int
+) : Granularity
+
 /** Creates a new [FullDataflowGranularity]. */
 fun full(): Granularity {
     return FullDataflowGranularity
@@ -80,6 +89,14 @@ fun partial(target: Declaration?): PartialDataflowGranularity {
     return PartialDataflowGranularity(target)
 }
 
+/**
+ * Creates a new [IndexedDataflowGranularity]. The [idx] is the index that is used for the partial
+ * dataflow.
+ */
+fun indexed(idx: Int): IndexedDataflowGranularity {
+    return IndexedDataflowGranularity(idx)
+}
+
 /**
  * This edge class defines a flow of data between [start] and [end]. The flow can have a certain
  * [granularity].

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

@@ -31,6 +31,9 @@ import de.fraunhofer.aisec.cpg.graph.declarations.*
 import de.fraunhofer.aisec.cpg.graph.edges.Edge
 import de.fraunhofer.aisec.cpg.graph.edges.flows.CallingContext
 import de.fraunhofer.aisec.cpg.graph.edges.flows.CallingContextOut
+import de.fraunhofer.aisec.cpg.graph.edges.flows.FullDataflowGranularity
+import de.fraunhofer.aisec.cpg.graph.edges.flows.Granularity
+import de.fraunhofer.aisec.cpg.graph.edges.flows.indexed
 import de.fraunhofer.aisec.cpg.graph.edges.flows.partial
 import de.fraunhofer.aisec.cpg.graph.statements.*
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
@@ -122,18 +125,25 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
                 }
             } else {
                 value.elements.forEach {
+                    var callingContext: CallingContext? = null
+                    var granularity: Granularity = FullDataflowGranularity
+                    edgePropertiesMap[Pair(it, key)]?.let {
+                        callingContext = it.filterIsInstance<CallingContext>().singleOrNull()
+                        granularity =
+                            it.filterIsInstance<Granularity>().singleOrNull()
+                                ?: FullDataflowGranularity
+                    }
+
                     if ((it is VariableDeclaration || it is ParameterDeclaration) && key == it) {
                         // Nothing to do
-                    } else if (
-                        Pair(it, key) in edgePropertiesMap &&
-                            edgePropertiesMap[Pair(it, key)] is CallingContext
-                    ) {
+                    } else if (callingContext != null) {
                         key.prevDFGEdges.addContextSensitive(
                             it,
-                            callingContext = (edgePropertiesMap[Pair(it, key)] as CallingContext),
+                            callingContext = callingContext,
+                            granularity = granularity,
                         )
                     } else {
-                        key.prevDFGEdges += it
+                        key.prevDFGEdges.add(it) { this.granularity = granularity }
                     }
                 }
             }
@@ -147,7 +157,11 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
     protected fun findAndSetProperties(from: Set<Node>, to: Node) {
         edgePropertiesMap
             .filter { it.key.first in from && it.key.second == null }
-            .forEach { edgePropertiesMap[Pair(it.key.first, to)] = it.value }
+            .forEach {
+                edgePropertiesMap
+                    .computeIfAbsent(Pair(it.key.first, to)) { mutableSetOf() }
+                    .addAll(it.value)
+            }
     }
 
     /**
@@ -291,12 +305,31 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
             // The rhs can be anything. The rhs flows to the respective lhs. To identify the
             // correct mapping, we use the "assignments" property which already searches for us.
             currentNode.assignments.forEach { assignment ->
-                // This was the last write to the respective declaration.
-                (assignment.target as? Declaration ?: (assignment.target as? Reference)?.refersTo)
-                    ?.let {
-                        doubleState.declarationsState[it] =
-                            PowersetLattice(identitySetOf(assignment.target as Node))
+                // Sometimes, we have a InitializerListExpression on the lhs which is not good at
+                // all...
+                if (assignment.target is InitializerListExpression) {
+                    assignment.target.initializers.forEachIndexed { idx, initializer ->
+                        (initializer as? Reference)?.let { ref ->
+                            ref.refersTo?.let {
+                                doubleState.declarationsState[it] =
+                                    PowersetLattice(identitySetOf(ref))
+                                doubleState.generalState[ref] =
+                                    PowersetLattice(identitySetOf(assignment.value as Node))
+                                edgePropertiesMap.computeIfAbsent(Pair(assignment.value, ref)) {
+                                    mutableSetOf<Any>()
+                                } += indexed(idx)
+                            }
+                        }
                     }
+                } else {
+                    // This was the last write to the respective declaration.
+                    (assignment.target as? Declaration
+                            ?: (assignment.target as? Reference)?.refersTo)
+                        ?.let {
+                            doubleState.declarationsState[it] =
+                                PowersetLattice(identitySetOf(assignment.target as Node))
+                        }
+                }
             }
         } else if (isIncOrDec(currentNode)) {
             // Increment or decrement => Add the prevWrite of the input to the input. After the
@@ -487,7 +520,9 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
                             }
                         doubleState.declarationsState[arg?.refersTo] =
                             PowersetLattice(identitySetOf(param))
-                        edgePropertiesMap[Pair(param, null)] = CallingContextOut(currentNode)
+                        edgePropertiesMap.computeIfAbsent(Pair(param, null)) {
+                            mutableSetOf<Any>()
+                        } += CallingContextOut(currentNode)
                     }
                 }
             } else {
@@ -514,7 +549,7 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
      * an entry works as follows: The 1st item in the pair is the prevDFG of the 2nd item. If the
      * 2nd item is null, it's obviously not relevant. Ultimately, it will be 2nd -prevDFG-> 1st.
      */
-    val edgePropertiesMap = mutableMapOf<Pair<Node, Node?>, Any>()
+    val edgePropertiesMap = mutableMapOf<Pair<Node, Node?>, MutableSet<Any>>()
 
     /**
      * Checks if the node performs an operation and an assignment at the same time e.g. with the

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.InitializerListExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
 import de.fraunhofer.aisec.cpg.graph.types.InitializerTypePropagation
@@ -200,6 +201,31 @@ class PythonAddDeclarationsPass(ctx: TranslationContext) : ComponentPass(ctx), L
                     assignExpression.declarations += handled
                 }
             }
+            (target as? InitializerListExpression)?.let {
+                it.initializers.forEach { initializer ->
+                    (initializer as? Reference)?.let {
+                        it.access = AccessValues.WRITE
+                        val handled = handleWriteToReference(initializer)
+                        if (handled is Declaration) {
+                            // We cannot assign an initializer here because this will lead to
+                            // duplicate
+                            // DFG edges, but we need to propagate the type information from our
+                            // value
+                            // to the declaration. We therefore add the declaration to the observers
+                            // of
+                            // the value's type, so that it gets informed and can change its own
+                            // type
+                            // accordingly.
+                            assignExpression
+                                .findValue(initializer)
+                                ?.registerTypeObserver(InitializerTypePropagation(handled))
+
+                            // Add it to our assign expression, so that we can find it in the AST
+                            assignExpression.declarations += handled
+                        }
+                    }
+                }
+            }
         }
     }
 

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

@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) 2025, Fraunhofer AISEC. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ *                    $$$$$$\  $$$$$$$\   $$$$$$\
+ *                   $$  __$$\ $$  __$$\ $$  __$$\
+ *                   $$ /  \__|$$ |  $$ |$$ /  \__|
+ *                   $$ |      $$$$$$$  |$$ |$$$$\
+ *                   $$ |      $$  ____/ $$ |\_$$ |
+ *                   $$ |  $$\ $$ |      $$ |  $$ |
+ *                   \$$$$$   |$$ |      \$$$$$   |
+ *                    \______/ \__|       \______/
+ *
+ */
+package de.fraunhofer.aisec.cpg.frontends.python
+
+import de.fraunhofer.aisec.cpg.graph.edges.flows.IndexedDataflowGranularity
+import de.fraunhofer.aisec.cpg.graph.functions
+import de.fraunhofer.aisec.cpg.graph.get
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.AssignExpression
+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.InitializerListExpression
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference
+import de.fraunhofer.aisec.cpg.graph.variables
+import de.fraunhofer.aisec.cpg.test.analyze
+import de.fraunhofer.aisec.cpg.test.assertLocalName
+import java.nio.file.Path
+import kotlin.test.assertEquals
+import kotlin.test.assertIs
+import kotlin.test.assertNotNull
+import kotlin.test.assertTrue
+import org.junit.jupiter.api.Test
+
+class DFGTest {
+    @Test
+    fun testListComprehensions() {
+        val topLevel = Path.of("src", "test", "resources", "python")
+        val result =
+            analyze(listOf(topLevel.resolve("tuple_assign.py").toFile()), topLevel, true) {
+                it.registerLanguage<PythonLanguage>()
+            }
+        assertNotNull(result)
+        val getTuple = result.functions["getTuple"]
+        assertNotNull(getTuple)
+
+        val body = getTuple.body
+        assertIs<Block>(body)
+        val assignment = body.statements[0]
+        assertIs<AssignExpression>(assignment)
+        assertEquals(1, assignment.lhs.size)
+        assertEquals(1, assignment.rhs.size)
+        val lhsTuple = assignment.lhs[0]
+        assertIs<InitializerListExpression>(lhsTuple)
+        assertEquals(2, lhsTuple.initializers.size)
+
+        val cRef = lhsTuple.initializers[0]
+
+        assertIs<Reference>(cRef)
+        assertLocalName("c", cRef)
+        val cRefPrevDFG = cRef.prevDFG.singleOrNull()
+        assertIs<CallExpression>(cRefPrevDFG)
+        assertLocalName("returnTuple", cRefPrevDFG)
+        val cRefPrevDFGGranularity = cRef.prevDFGEdges.single().granularity
+        assertIs<IndexedDataflowGranularity>(cRefPrevDFGGranularity)
+        assertEquals(0, cRefPrevDFGGranularity.index)
+
+        val c = body.variables["c"]
+        assertNotNull(c)
+        assertTrue(c.prevDFG.isEmpty())
+
+        val dRef = lhsTuple.initializers[1]
+        assertIs<Reference>(dRef)
+        assertLocalName("d", dRef)
+        val dRefPrevDFG = dRef.prevDFG.singleOrNull()
+        assertIs<CallExpression>(dRefPrevDFG)
+        assertLocalName("returnTuple", dRefPrevDFG)
+        val dRefPrevDFGGranularity = dRefPrevDFG.prevDFGEdges.single().granularity
+        assertIs<IndexedDataflowGranularity>(dRefPrevDFGGranularity)
+        assertEquals(1, dRefPrevDFGGranularity.index)
+
+        val d = body.variables["d"]
+        assertNotNull(d)
+        assertTrue(d.prevDFG.isEmpty())
+    }
+}

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

@@ -0,0 +1,6 @@
+def returnTuple(a, b):
+    return (a, b)
+
+def getTuple(a, b):
+    (c, d) = returnTuple(a, b)
+    print(str(c) + " " + str(d))