diff --git a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForFunctionsCreator.scala b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForFunctionsCreator.scala index 6383c90ed897..a8601327a0f2 100644 --- a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForFunctionsCreator.scala +++ b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForFunctionsCreator.scala @@ -26,10 +26,10 @@ import scala.collection.mutable trait AstForFunctionsCreator(implicit withSchemaValidation: ValidationMode) { this: AstCreator => - /** As expressions may be discarded, we cannot store closure ASTs in the diffgraph at the point of creation. We need - * to only write these at the end. + /** As expressions may be discarded, we cannot store closure ASTs in the diffgraph at the point of creation. So we + * assume every reference to this map means that the closure AST was successfully propagated. */ - protected val closureToRefs = mutable.Map.empty[RubyExpression, Seq[Ast]] + protected val closureToRefs = mutable.Map.empty[RubyExpression, Seq[NewNode]] /** Creates method declaration related structures. * @param node diff --git a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForStatementsCreator.scala b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForStatementsCreator.scala index e08cfa04e06b..33c4c3f96059 100644 --- a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForStatementsCreator.scala +++ b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForStatementsCreator.scala @@ -93,18 +93,18 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t protected def astForDoBlock(block: Block & RubyExpression): Seq[Ast] = { // Create closure structures: [MethodDecl, TypeRef, MethodRef] if (closureToRefs.contains(block)) { - closureToRefs(block) + closureToRefs(block).map(x => Ast(x.copy)) } else { val methodName = nextClosureName() - val methodAstsWithRefs = block.body match { + val methodRefAsts = block.body match { case x: Block => astForMethodDeclaration(x.toMethodDeclaration(methodName, Option(block.parameters)), isClosure = true) case _ => astForMethodDeclaration(block.toMethodDeclaration(methodName, Option(block.parameters)), isClosure = true) } - closureToRefs.put(block, methodAstsWithRefs) - methodAstsWithRefs + closureToRefs.put(block, methodRefAsts.flatMap(_.root)) + methodRefAsts } }