diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilder.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilder.kt index 260e45f1355..d075ae8a97b 100644 --- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilder.kt +++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilder.kt @@ -99,7 +99,8 @@ fun MetadataProvider.newUnaryOperator( operatorCode: String, postfix: Boolean, prefix: Boolean, - rawNode: Any? = null + rawNode: Any? = null, + init: (UnaryOperator.() -> Unit)? = null, ): UnaryOperator { val node = UnaryOperator() node.applyMetadata(this, operatorCode, rawNode, true) @@ -108,6 +109,10 @@ fun MetadataProvider.newUnaryOperator( node.isPostfix = postfix node.isPrefix = prefix + if (init != null) { + init(node) + } + log(node) return node diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt index cdc7ee7af9a..28385ae55d7 100644 --- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt +++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/NodeBuilder.kt @@ -383,6 +383,21 @@ private fun Node.setCodeAndLocation( this.location = provider.locationOf(rawNode) } +context(ContextProvider) +fun T.withScope(init: T.() -> Unit): T { + val scopeManager = + this@ContextProvider.ctx?.scopeManager + ?: throw TranslationException( + "Trying to create node children without a ContextProvider. This will fail." + ) + + scopeManager.enterScope(this) + init(this) + scopeManager.leaveScope(this) + + return this +} + context(ContextProvider) fun T.withChildren( hasScope: Boolean = false, diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/StatementHolder.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/StatementHolder.kt index 39d9492be57..9f0eaf503e0 100644 --- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/StatementHolder.kt +++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/StatementHolder.kt @@ -25,6 +25,7 @@ */ package de.fraunhofer.aisec.cpg.graph +import de.fraunhofer.aisec.cpg.graph.edge.AstPropertyEdge import de.fraunhofer.aisec.cpg.graph.edge.Properties import de.fraunhofer.aisec.cpg.graph.edge.PropertyEdge import de.fraunhofer.aisec.cpg.graph.edge.PropertyEdge.Companion.unwrap @@ -55,6 +56,7 @@ interface StatementHolder : Holder { return unwrap(statementEdges) } set(value) { + // TODO: we also need to set the ast parent in this case statementEdges = wrap(value, this as Node) } @@ -66,7 +68,7 @@ interface StatementHolder : Holder { * @param s the statement */ fun addStatement(s: Statement) { - val propertyEdge = PropertyEdge((this as Node), s) + val propertyEdge = AstPropertyEdge((this as Node), s) propertyEdge.addProperty(Properties.INDEX, statementEdges.size) statementEdges.add(propertyEdge) } diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edge/PropertyEdge.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edge/PropertyEdge.kt index f9fb9408f77..a771cf3d952 100644 --- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edge/PropertyEdge.kt +++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edge/PropertyEdge.kt @@ -41,6 +41,12 @@ import org.neo4j.ogm.annotation.* import org.neo4j.ogm.annotation.typeconversion.Convert import org.slf4j.LoggerFactory +class AstPropertyEdge : PropertyEdge { + constructor(start: Node, end: T) : super(start, end) { + end.astParent = start + } +} + /** * This class represents an edge between two [Node] objects in a Neo4J graph. It can be used to * store additional information that relate to the relationship between the two nodes that belong to diff --git a/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilderTest.kt b/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilderTest.kt index 1dea80ee771..d499096afb1 100644 --- a/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilderTest.kt +++ b/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilderTest.kt @@ -35,6 +35,7 @@ import de.fraunhofer.aisec.cpg.graph.statements.expressions.Literal import de.fraunhofer.aisec.cpg.graph.statements.expressions.Reference import kotlin.test.Test import kotlin.test.assertEquals +import kotlin.test.assertNotNull import kotlin.test.assertNull import kotlin.test.assertTrue @@ -92,4 +93,21 @@ class ExpressionBuilderTest { assertEquals(binOp, lit2.astParent) } } + + @Test + fun testBlock() { + with(TestLanguageFrontend()) { + val block = + newBlock().withScope { + this += + newUnaryOperator("++", prefix = false, postfix = true) { + input = newReference("a") + } + } + + val binOp = block.statements.firstOrNull() + assertNotNull(binOp) + assertEquals(block, binOp.astParent) + } + } }