[ruby] Re-implemented "Ignore "Throwaway" AST Structures (#4982)" (#4985

)

* Revert "Revert "[ruby] Ignore "Throwaway" AST Structures (#4982)" (#4983)"

This reverts commit 464480d.

* [ruby] Re-implemented "Ignore "Throwaway" AST Structures (#4982)"
This correctly prevents re-use of nodes that are already being used elsewhere by ensuring deep copies.

joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala

@@ -157,18 +157,18 @@ trait AstForExpressionsCreator(implicit withSchemaValidation: ValidationMode) {
   /** Attempts to extract a type from the base of a member call.
     */
   protected def typeFromCallTarget(baseNode: RubyExpression): Option[String] = {
-    scope.lookupVariable(baseNode.text) match {
-      // fixme: This should be under type recovery logic
-      case Some(decl: NewLocal) if decl.typeFullName != Defines.Any             => Option(decl.typeFullName)
-      case Some(decl: NewMethodParameterIn) if decl.typeFullName != Defines.Any => Option(decl.typeFullName)
-      case Some(decl: NewLocal) if decl.dynamicTypeHintFullName.nonEmpty => decl.dynamicTypeHintFullName.headOption
-      case Some(decl: NewMethodParameterIn) if decl.dynamicTypeHintFullName.nonEmpty =>
-        decl.dynamicTypeHintFullName.headOption
+    baseNode match {
+      case literal: LiteralExpr => Option(literal.typeFullName)
       case _ =>
-        astForExpression(baseNode).nodes
-          .flatMap(_.properties.get(PropertyNames.TYPE_FULL_NAME).map(_.toString))
-          .filterNot(_ == XDefines.Any)
-          .headOption
+        scope.lookupVariable(baseNode.text) match {
+          // fixme: This should be under type recovery logic
+          case Some(decl: NewLocal) if decl.typeFullName != Defines.Any             => Option(decl.typeFullName)
+          case Some(decl: NewMethodParameterIn) if decl.typeFullName != Defines.Any => Option(decl.typeFullName)
+          case Some(decl: NewLocal) if decl.dynamicTypeHintFullName.nonEmpty => decl.dynamicTypeHintFullName.headOption
+          case Some(decl: NewMethodParameterIn) if decl.dynamicTypeHintFullName.nonEmpty =>
+            decl.dynamicTypeHintFullName.headOption
+          case _ => None
+        }
     }
   }
 
@@ -296,7 +296,10 @@ trait AstForExpressionsCreator(implicit withSchemaValidation: ValidationMode) {
     val createAssignmentToTmp = !baseAstCache.contains(target)
     val tmpName = baseAstCache
       .updateWith(target) {
-        case Some(tmpName) => Option(tmpName)
+        case Some(tmpName) =>
+          // TODO: Type ref nodes are automatically committed on creation, so if we have found a suitable cached AST,
+          //  we want to clean this creation up.
+          Option(tmpName)
         case None =>
           val tmpName     = this.tmpGen.fresh
           val tmpGenLocal = NewLocal().name(tmpName).code(tmpName).typeFullName(Defines.Any)
@@ -872,16 +875,11 @@ trait AstForExpressionsCreator(implicit withSchemaValidation: ValidationMode) {
   }
 
   private def astForMemberCallWithoutBlock(node: SimpleCall, memberAccess: MemberAccess): Ast = {
-    val receiverAst = astForFieldAccess(memberAccess)
-    val methodName  = memberAccess.memberName
-    // TODO: Type recovery should potentially resolve this
-    val methodFullName = typeFromCallTarget(memberAccess.target)
-      .map(x => s"$x.$methodName")
-      .getOrElse(XDefines.DynamicCallUnknownFullName)
-    val argumentAsts = node.arguments.map(astForMethodCallArgument)
-    val call =
-      callNode(node, code(node), methodName, XDefines.DynamicCallUnknownFullName, DispatchTypes.DYNAMIC_DISPATCH)
-        .possibleTypes(IndexedSeq(methodFullName))
+    val receiverAst    = astForFieldAccess(memberAccess)
+    val methodName     = memberAccess.memberName
+    val methodFullName = XDefines.DynamicCallUnknownFullName
+    val argumentAsts   = node.arguments.map(astForMethodCallArgument)
+    val call           = callNode(node, code(node), methodName, methodFullName, DispatchTypes.DYNAMIC_DISPATCH)
 
     callAst(call, argumentAsts, Some(receiverAst))
   }

joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForFunctionsCreator.scala

@@ -26,6 +26,11 @@ 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. So we
+    * assume every reference to this map means that the closure AST was successfully propagated.
+    */
+  protected val closureToRefs = mutable.Map.empty[RubyExpression, Seq[NewNode]]
+
   /** Creates method declaration related structures.
     * @param node
     *   the node to create the AST structure from.
@@ -194,8 +199,11 @@ trait AstForFunctionsCreator(implicit withSchemaValidation: ValidationMode) { th
       case _                                              => false
     })
 
-    val methodRefOption = refs.flatMap(_.nodes).collectFirst { case x: NewTypeRef => x }
+    val typeRefOption = refs.flatMap(_.nodes).collectFirst { case x: NewTypeRef => x }
 
+    val astChildren  = mutable.Buffer.empty[NewNode]
+    val refEdges     = mutable.Buffer.empty[(NewNode, NewNode)]
+    val captureEdges = mutable.Buffer.empty[(NewNode, NewNode)]
     capturedLocalNodes
       .collect {
         case local: NewLocal =>
@@ -216,14 +224,17 @@ trait AstForFunctionsCreator(implicit withSchemaValidation: ValidationMode) { th
         )
 
         // Create new local node for lambda, with corresponding REF edges to identifiers and closure binding
-        capturedBlockAst.root.foreach(rootBlock => diffGraph.addEdge(rootBlock, capturingLocal, EdgeTypes.AST))
-        capturedIdentifiers.filter(_.name == name).foreach(i => diffGraph.addEdge(i, capturingLocal, EdgeTypes.REF))
-        diffGraph.addEdge(closureBinding, capturedLocal, EdgeTypes.REF)
+        val _refEdges =
+          capturedIdentifiers.filter(_.name == name).map(i => i -> capturingLocal) :+ (closureBinding, capturedLocal)
 
-        methodRefOption.foreach(methodRef => diffGraph.addEdge(methodRef, closureBinding, EdgeTypes.CAPTURE))
+        astChildren.addOne(capturingLocal)
+        refEdges.addAll(_refEdges.toList)
+        captureEdges.addAll(typeRefOption.map(typeRef => typeRef -> closureBinding).toList)
       }
 
-    capturedBlockAst
+    val astWithAstChildren = astChildren.foldLeft(capturedBlockAst) { case (ast, child) => ast.withChild(Ast(child)) }
+    val astWithRefEdges = refEdges.foldLeft(astWithAstChildren) { case (ast, (src, dst)) => ast.withRefEdge(src, dst) }
+    captureEdges.foldLeft(astWithRefEdges) { case (ast, (src, dst)) => ast.withCaptureEdge(src, dst) }
   }
 
   /** Creates the bindings between the method and its types. This is useful for resolving function pointers and imports.

joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForStatementsCreator.scala

@@ -92,16 +92,20 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
 
   protected def astForDoBlock(block: Block & RubyExpression): Seq[Ast] = {
     // Create closure structures: [MethodDecl, TypeRef, MethodRef]
-    val methodName = nextClosureName()
+    if (closureToRefs.contains(block)) {
+      closureToRefs(block).map(x => Ast(x.copy))
+    } else {
+      val methodName = nextClosureName()
 
-    val methodAstsWithRefs = 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)
+      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, methodRefAsts.flatMap(_.root))
+      methodRefAsts
     }
-
-    methodAstsWithRefs
   }
 
   protected def astForReturnExpression(node: ReturnExpression): Ast = {

joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/DoBlockTests.scala

@@ -401,7 +401,7 @@ class DoBlockTests extends RubyCode2CpgFixture {
                      |""".stripMargin)
 
     inside(cpg.local.l) {
-      case jfsOutsideLocal :: hashInsideLocal :: jfsCapturedLocal :: tmp0 :: tmp1 :: Nil =>
+      case jfsOutsideLocal :: hashInsideLocal :: tmp0 :: jfsCapturedLocal :: tmp1 :: Nil =>
         jfsOutsideLocal.closureBindingId shouldBe None
         hashInsideLocal.closureBindingId shouldBe None
         jfsCapturedLocal.closureBindingId shouldBe Some("Test0.rb:<main>.get_pto_schedule.jfs")
@@ -412,7 +412,7 @@ class DoBlockTests extends RubyCode2CpgFixture {
     }
 
     inside(cpg.method.isLambda.local.l) {
-      case hashLocal :: jfsLocal :: _ :: Nil =>
+      case hashLocal :: _ :: jfsLocal :: Nil =>
         hashLocal.closureBindingId shouldBe None
         jfsLocal.closureBindingId shouldBe Some("Test0.rb:<main>.get_pto_schedule.jfs")
       case xs => fail(s"Expected 3 locals in lambda, got ${xs.code.mkString(",")}")

joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/MethodTests.scala

@@ -984,4 +984,21 @@ class MethodTests extends RubyCode2CpgFixture {
       }
     }
   }
+
+  "lambdas as arguments to a long chained call" should {
+    val cpg = code("""
+        |def foo(xs, total_ys, hex_values)
+        |  xs.map.with_index { |f, i| [f / total_ys, hex_values[i]] }                       # 1
+        |    .sort_by { |r| -r[0] }                                                         # 2
+        |    .reject { |r| r[1].size == 8 && r[1].end_with?('00') }                         # 3
+        |    .map { |r| Foo::Bar::Baz.new(*r[1][0..5].scan(/../).map { |c| c.to_i(16) }) }  # 4 & 5
+        |    .slice(0, quantity)
+        |  end
+        |""".stripMargin)
+
+    "not write lambda nodes that are already assigned to some temp variable" in {
+      cpg.typeRef.typeFullName(".*Proc").size shouldBe 5
+      cpg.typeRef.whereNot(_.astParent).size shouldBe 0
+    }
+  }
 }

joern-cli/frontends/x2cpg/src/main/scala/io/joern/x2cpg/Ast.scala

@@ -49,6 +49,10 @@ object Ast {
     ast.bindsEdges.foreach { edge =>
       diffGraph.addEdge(edge.src, edge.dst, EdgeTypes.BINDS)
     }
+
+    ast.captureEdges.foreach { edge =>
+      diffGraph.addEdge(edge.src, edge.dst, EdgeTypes.CAPTURE)
+    }
   }
 
   def neighbourValidation(src: NewNode, dst: NewNode, edge: String)(implicit
@@ -92,7 +96,8 @@ case class Ast(
   refEdges: collection.Seq[AstEdge] = Vector.empty,
   bindsEdges: collection.Seq[AstEdge] = Vector.empty,
   receiverEdges: collection.Seq[AstEdge] = Vector.empty,
-  argEdges: collection.Seq[AstEdge] = Vector.empty
+  argEdges: collection.Seq[AstEdge] = Vector.empty,
+  captureEdges: collection.Seq[AstEdge] = Vector.empty
 )(implicit withSchemaValidation: ValidationMode = ValidationMode.Disabled) {
 
   def root: Option[NewNode] = nodes.headOption
@@ -114,7 +119,8 @@ case class Ast(
       argEdges = argEdges ++ other.argEdges,
       receiverEdges = receiverEdges ++ other.receiverEdges,
       refEdges = refEdges ++ other.refEdges,
-      bindsEdges = bindsEdges ++ other.bindsEdges
+      bindsEdges = bindsEdges ++ other.bindsEdges,
+      captureEdges = captureEdges ++ other.captureEdges
     )
   }
 
@@ -126,7 +132,8 @@ case class Ast(
       argEdges = argEdges ++ other.argEdges,
       receiverEdges = receiverEdges ++ other.receiverEdges,
       refEdges = refEdges ++ other.refEdges,
-      bindsEdges = bindsEdges ++ other.bindsEdges
+      bindsEdges = bindsEdges ++ other.bindsEdges,
+      captureEdges = captureEdges ++ other.captureEdges
     )
   }
 
@@ -217,6 +224,16 @@ case class Ast(
     this.copy(receiverEdges = receiverEdges ++ dsts.map(AstEdge(src, _)))
   }
 
+  def withCaptureEdge(src: NewNode, dst: NewNode): Ast = {
+    Ast.neighbourValidation(src, dst, EdgeTypes.CAPTURE)
+    this.copy(captureEdges = captureEdges ++ List(AstEdge(src, dst)))
+  }
+
+  def withCaptureEdges(src: NewNode, dsts: Seq[NewNode]): Ast = {
+    dsts.foreach(dst => Ast.neighbourValidation(src, dst, EdgeTypes.CAPTURE))
+    this.copy(captureEdges = captureEdges ++ dsts.map(AstEdge(src, _)))
+  }
+
   /** Returns a deep copy of the sub tree rooted in `node`. If `order` is set, then the `order` and `argumentIndex`
     * fields of the new root node are set to `order`. If `replacementNode` is set, then this replaces `node` in the new
     * copy.
@@ -250,14 +267,16 @@ case class Ast(
     val newRefEdges       = refEdges.filter(_.src == node).map(x => AstEdge(newNode, newIfExists(x.dst)))
     val newBindsEdges     = bindsEdges.filter(_.src == node).map(x => AstEdge(newNode, newIfExists(x.dst)))
     val newReceiverEdges  = receiverEdges.filter(_.src == node).map(x => AstEdge(newNode, newIfExists(x.dst)))
+    val newCaptureEdges   = captureEdges.filter(_.src == node).map(x => AstEdge(newNode, newIfExists(x.dst)))
 
     Ast(newNode)
       .copy(
         argEdges = newArgEdges,
         conditionEdges = newConditionEdges,
         refEdges = newRefEdges,
         bindsEdges = newBindsEdges,
-        receiverEdges = newReceiverEdges
+        receiverEdges = newReceiverEdges,
+        captureEdges = newCaptureEdges
       )
       .withChildren(newChildren)
   }