Add Comprehensions of Lists, Sets and Maps and generator expressions

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

@@ -31,6 +31,7 @@ import de.fraunhofer.aisec.cpg.graph.NodeBuilder.log
 import de.fraunhofer.aisec.cpg.graph.edges.flows.ContextSensitiveDataflow
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.AssignExpression
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.CollectionComprehension
 import de.fraunhofer.aisec.cpg.graph.types.ProblemType
 import de.fraunhofer.aisec.cpg.graph.types.Type
 
@@ -525,6 +526,24 @@ fun MetadataProvider.newInitializerListExpression(
     return node
 }
 
+@JvmOverloads
+fun MetadataProvider.newComprehensionExpression(rawNode: Any? = null): ComprehensionExpression {
+    val node = ComprehensionExpression()
+    node.applyMetadata(this, EMPTY_NAME, rawNode, true)
+
+    log(node)
+    return node
+}
+
+@JvmOverloads
+fun MetadataProvider.newCollectionComprehension(rawNode: Any? = null): CollectionComprehension {
+    val node = CollectionComprehension()
+    node.applyMetadata(this, EMPTY_NAME, rawNode, true)
+
+    log(node)
+    return node
+}
+
 /**
  * Creates a new [TypeExpression]. The [MetadataProvider] receiver will be used to fill different
  * meta-data using [Node.applyMetadata]. Calling this extension function outside of Kotlin requires

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

@@ -32,6 +32,7 @@ import de.fraunhofer.aisec.cpg.graph.declarations.*
 import de.fraunhofer.aisec.cpg.graph.scopes.RecordScope
 import de.fraunhofer.aisec.cpg.graph.statements.*
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.CollectionComprehension
 import de.fraunhofer.aisec.cpg.graph.types.FunctionType
 import de.fraunhofer.aisec.cpg.graph.types.Type
 import de.fraunhofer.aisec.cpg.graph.types.UnknownType
@@ -325,6 +326,47 @@ fun LanguageFrontend<*, *>.subscriptExpr(
     return node
 }
 
+context(Holder<out Statement>)
+fun LanguageFrontend<*, *>.listComp(
+    init: (CollectionComprehension.() -> Unit)? = null
+): CollectionComprehension {
+    val node = newCollectionComprehension()
+
+    if (init != null) {
+        init(node)
+    }
+
+    // Only add this to an argument holder if the nearest holder is an argument holder
+    val holder = this@Holder
+    if (holder is StatementHolder) {
+        holder += node
+    }
+    if (holder is ArgumentHolder) {
+        holder += node
+    }
+
+    return node
+}
+
+context(Holder<out Statement>)
+fun LanguageFrontend<*, *>.compExpr(
+    init: (ComprehensionExpression.() -> Unit)? = null
+): ComprehensionExpression {
+    val node = newComprehensionExpression()
+
+    if (init != null) {
+        init(node)
+    }
+
+    // Only add this to an argument holder if the nearest holder is an argument holder
+    val holder = this@Holder
+    if (holder is ArgumentHolder) {
+        holder += node
+    }
+
+    return node
+}
+
 /**
  * Creates a new [DeclarationStatement] in the Fluent Node DSL and adds it to the
  * [StatementHolder.statements] of the nearest enclosing [StatementHolder]. The [init] block can be

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/ast/AstEdge.kt

@@ -48,7 +48,7 @@ fun <NodeType : Node> Node.astEdgesOf(
 }
 
 /**
- * Creates an single optional [AstEdge] starting from this node (wrapped in a [EdgeSingletonList]
+ * Creates a single optional [AstEdge] starting from this node (wrapped in a [EdgeSingletonList]
  * container).
  */
 fun <NodeType : Node> Node.astOptionalEdgeOf(
@@ -64,7 +64,7 @@ fun <NodeType : Node> Node.astOptionalEdgeOf(
 }
 
 /**
- * Creates an single [AstEdge] starting from this node (wrapped in a [EdgeSingletonList] container).
+ * Creates a single [AstEdge] starting from this node (wrapped in a [EdgeSingletonList] container).
  */
 fun <NodeType : Node> Node.astEdgeOf(
     of: NodeType,

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/UnwrappedEdgeList.kt

@@ -70,7 +70,12 @@ class UnwrappedEdgeList<NodeType : Node, EdgeType : Edge<NodeType>>(
     }
 
     override fun subList(fromIndex: Int, toIndex: Int): MutableList<NodeType> {
-        TODO("Not yet implemented")
+        return if (list.outgoing) {
+            list.subList(fromIndex, toIndex).map { it.end }.toMutableList()
+        } else {
+            @Suppress("UNCHECKED_CAST")
+            list.subList(fromIndex, toIndex).map { it.start as NodeType }.toMutableList()
+        }
     }
 
     override fun get(index: Int): NodeType {

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/statements/expressions/CollectionComprehension.kt

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) 2024, 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.graph.statements.expressions
+
+import de.fraunhofer.aisec.cpg.graph.*
+import de.fraunhofer.aisec.cpg.graph.edges.ast.astEdgesOf
+import de.fraunhofer.aisec.cpg.graph.edges.ast.astOptionalEdgeOf
+import de.fraunhofer.aisec.cpg.graph.edges.unwrapping
+import de.fraunhofer.aisec.cpg.graph.statements.Statement
+import java.util.Objects
+import org.apache.commons.lang3.builder.ToStringBuilder
+import org.neo4j.ogm.annotation.Relationship
+
+/**
+ * Represent a list/set/map comprehension or similar expression. It contains four major components:
+ * The statement, the variable, the iterable and a predicate which are combined to something like
+ * `[statement(variable) : variable in iterable if predicate(variable)]`.
+ *
+ * Some languages provide a way to have multiple variables, iterables and predicates. For this
+ * reason, we represent the `variable, iterable and predicate in its own class
+ * [CollectionComprehension.ComprehensionExpression].
+ */
+class CollectionComprehension : Expression(), ArgumentHolder {
+
+    @Relationship("COMPREHENSION_EXPRESSIONS")
+    var comprehensionExpressionEdges = astEdgesOf<ComprehensionExpression>()
+    /**
+     * This field contains one or multiple [ComprehensionExpression]s.
+     *
+     * Note: Instead of having a list here, we could also enforce that the frontend nests the
+     * expressions in a meaningful way (in particular this would help us to satisfy dependencies
+     * between the comprehensions' variables).
+     */
+    var comprehensionExpressions by
+        unwrapping(CollectionComprehension::comprehensionExpressionEdges)
+
+    @Relationship("STATEMENT") var statementEdge = astOptionalEdgeOf<Statement>()
+    /**
+     * This field contains the statement which is applied to each element of the input for which the
+     * predicate returned `true`.
+     */
+    var statement by unwrapping(CollectionComprehension::statementEdge)
+
+    override fun toString() =
+        ToStringBuilder(this, TO_STRING_STYLE)
+            .appendSuper(super.toString())
+            .append("statement", statement)
+            .append("comprehensions", comprehensionExpressions)
+            .toString()
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (other !is CollectionComprehension) return false
+        return super.equals(other) &&
+            statement == other.statement &&
+            comprehensionExpressions == other.comprehensionExpressions
+    }
+
+    override fun hashCode() = Objects.hash(super.hashCode(), statement, comprehensionExpressions)
+
+    override fun addArgument(expression: Expression) {
+        if (this.statement == null) {
+            this.statement = expression
+        } else if (expression is ComprehensionExpression) {
+            this.comprehensionExpressions += expression
+        }
+    }
+
+    override fun replaceArgument(old: Expression, new: Expression): Boolean {
+        if (this.statement == old) {
+            this.statement = new
+            return true
+        }
+        if (new !is ComprehensionExpression) return false
+        var changedSomething = false
+        val newCompExp =
+            this.comprehensionExpressions.map {
+                if (it == old) {
+                    changedSomething = true
+                    new
+                } else it
+            }
+        this.comprehensionExpressions.clear()
+        this.comprehensionExpressions.addAll(newCompExp)
+        return changedSomething
+    }
+
+    override fun hasArgument(expression: Expression): Boolean {
+        return this.statement == expression || expression in this.comprehensionExpressions
+    }
+}

cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/statements/expressions/ComprehensionExpression.kt

@@ -0,0 +1,126 @@
+/*
+ * Copyright (c) 2024, 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.graph.statements.expressions
+
+import de.fraunhofer.aisec.cpg.graph.AccessValues
+import de.fraunhofer.aisec.cpg.graph.ArgumentHolder
+import de.fraunhofer.aisec.cpg.graph.edges.ast.astEdgeOf
+import de.fraunhofer.aisec.cpg.graph.edges.ast.astOptionalEdgeOf
+import de.fraunhofer.aisec.cpg.graph.edges.unwrapping
+import de.fraunhofer.aisec.cpg.graph.statements.Statement
+import java.util.Objects
+import kotlin.collections.plusAssign
+import org.apache.commons.lang3.builder.ToStringBuilder
+import org.neo4j.ogm.annotation.Relationship
+
+/** This class holds the variable, iterable and predicate of the [CollectionComprehension]. */
+class ComprehensionExpression : Expression(), ArgumentHolder {
+    @Relationship("VARIABLE")
+    var variableEdge =
+        astEdgeOf<Statement>(
+            of = ProblemExpression("Missing variableEdge in ${this::class}"),
+            onChanged = { _, new ->
+                val end = new?.end
+                if (end is Reference) {
+                    end.access = AccessValues.WRITE
+                }
+            }
+        )
+
+    /**
+     * This field contains the iteration variable of the comprehension. It can be either a new
+     * variable declaration or a reference (probably to a new variable).
+     */
+    var variable by unwrapping(ComprehensionExpression::variableEdge)
+
+    @Relationship("ITERABLE")
+    var iterableEdge =
+        astEdgeOf<Expression>(ProblemExpression("Missing iterable in ${this::class}"))
+
+    /** This field contains the iteration subject of the loop. */
+    var iterable by unwrapping(ComprehensionExpression::iterableEdge)
+
+    @Relationship("PREDICATE") var predicateEdge = astOptionalEdgeOf<Statement>()
+
+    /**
+     * This field contains the predicate which has to hold to evaluate `statement(variable)` and
+     * include it in the result.
+     */
+    var predicate by unwrapping(ComprehensionExpression::predicateEdge)
+
+    override fun toString() =
+        ToStringBuilder(this, TO_STRING_STYLE)
+            .appendSuper(super.toString())
+            .append("variable", variable)
+            .append("iterable", iterable)
+            .append("predicate", predicate)
+            .toString()
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (other !is ComprehensionExpression) return false
+        return super.equals(other) &&
+            variable == other.variable &&
+            iterable == other.iterable &&
+            predicate == other.predicate
+    }
+
+    override fun hashCode() = Objects.hash(super.hashCode(), variable, iterable, predicate)
+
+    override fun addArgument(expression: Expression) {
+        if (this.variable == null) {
+            this.variable = expression
+        } else if (this.iterable == null) {
+            this.iterable = expression
+        } else {
+            this.predicate = expression
+        }
+    }
+
+    override fun replaceArgument(old: Expression, new: Expression): Boolean {
+        if (this.variable == old) {
+            this.variable = new
+            return true
+        }
+
+        if (this.iterable == old) {
+            this.iterable = new
+            return true
+        }
+
+        if (this.predicate == old) {
+            this.predicate = new
+            return true
+        }
+        return false
+    }
+
+    override fun hasArgument(expression: Expression): Boolean {
+        return this.variable == expression ||
+            this.iterable == expression ||
+            expression == this.predicate
+    }
+}