[ruby] Move lowering of ||= and &&= to AstCreator (#5055)

* Moved lowering for ||= and &&= to AstCreator

* Moved lowering func to AstCreatorHelper trait

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

@@ -1,11 +1,17 @@
 package io.joern.rubysrc2cpg.astcreation
 import io.joern.rubysrc2cpg.astcreation.RubyIntermediateAst.{
   ClassFieldIdentifier,
+  ControlFlowStatement,
   DummyNode,
+  IfExpression,
   InstanceFieldIdentifier,
   MemberAccess,
+  RubyExpression,
   RubyFieldIdentifier,
-  RubyExpression
+  SingleAssignment,
+  StatementList,
+  TextSpan,
+  UnaryExpression
 }
 import io.joern.rubysrc2cpg.datastructures.{BlockScope, FieldDecl}
 import io.joern.rubysrc2cpg.passes.Defines
@@ -146,6 +152,44 @@ trait AstCreatorHelper(implicit withSchemaValidation: ValidationMode) { this: As
     member
   }
 
+  /** Lowers the `||=` and `&&=` assignment operators to the respective `.nil?` checks
+    */
+  def lowerAssignmentOperator(lhs: RubyExpression, rhs: RubyExpression, op: String, span: TextSpan): RubyExpression &
+    ControlFlowStatement = {
+    val condition  = nilCheckCondition(lhs, op, "nil?", span)
+    val thenClause = nilCheckThenClause(lhs, rhs, span)
+    nilCheckIfStatement(condition, thenClause, span)
+  }
+
+  /** Generates the required `.nil?` check condition used in the lowering of `||=` and `&&=`
+    */
+  private def nilCheckCondition(lhs: RubyExpression, op: String, memberName: String, span: TextSpan): RubyExpression = {
+    val memberAccess =
+      MemberAccess(lhs, op = ".", memberName = "nil?")(span.spanStart(s"${lhs.span.text}.nil?"))
+    if op == "||=" then memberAccess
+    else UnaryExpression(op = "!", expression = memberAccess)(span.spanStart(s"!${memberAccess.span.text}"))
+  }
+
+  /** Generates the assignment and the `thenClause` used in the lowering of `||=` and `&&=`
+    */
+  private def nilCheckThenClause(lhs: RubyExpression, rhs: RubyExpression, span: TextSpan): RubyExpression = {
+    StatementList(List(SingleAssignment(lhs, "=", rhs)(span.spanStart(s"${lhs.span.text} = ${rhs.span.text}"))))(
+      span.spanStart(s"${lhs.span.text} = ${rhs.span.text}")
+    )
+  }
+
+  /** Generates the if statement for the lowering of `||=` and `&&=`
+    */
+  private def nilCheckIfStatement(
+    condition: RubyExpression,
+    thenClause: RubyExpression,
+    span: TextSpan
+  ): RubyExpression & ControlFlowStatement = {
+    IfExpression(condition = condition, thenClause = thenClause, elsifClauses = List.empty, elseClause = None)(
+      span.spanStart(s"if ${condition.span.text} then ${thenClause.span.text} end")
+    )
+  }
+
   protected val UnaryOperatorNames: Map[String, String] = Map(
     "!"   -> Operators.logicalNot,
     "not" -> Operators.logicalNot,

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

@@ -11,6 +11,7 @@ import io.joern.rubysrc2cpg.astcreation.RubyIntermediateAst.{
   IfExpression,
   MemberCall,
   NextExpression,
+  OperatorAssignment,
   RescueExpression,
   ReturnExpression,
   RubyExpression,
@@ -25,6 +26,7 @@ import io.joern.rubysrc2cpg.astcreation.RubyIntermediateAst.{
   WhenClause,
   WhileExpression
 }
+import io.joern.rubysrc2cpg.parser.RubyJsonHelpers
 import io.joern.rubysrc2cpg.passes.Defines
 import io.joern.x2cpg.{Ast, ValidationMode}
 import io.shiftleft.codepropertygraph.generated.{ControlStructureTypes, DispatchTypes, Operators}
@@ -39,16 +41,17 @@ import io.shiftleft.codepropertygraph.generated.nodes.{
 trait AstForControlStructuresCreator(implicit withSchemaValidation: ValidationMode) { this: AstCreator =>
 
   protected def astForControlStructureExpression(node: ControlFlowStatement): Ast = node match {
-    case node: WhileExpression   => astForWhileStatement(node)
-    case node: DoWhileExpression => astForDoWhileStatement(node)
-    case node: UntilExpression   => astForUntilStatement(node)
-    case node: CaseExpression    => blockAst(NewBlock(), astsForCaseExpression(node).toList)
-    case node: IfExpression      => astForIfExpression(node)
-    case node: UnlessExpression  => astForUnlessStatement(node)
-    case node: ForExpression     => astForForExpression(node)
-    case node: RescueExpression  => astForRescueExpression(node)
-    case node: NextExpression    => astForNextExpression(node)
-    case node: BreakExpression   => astForBreakExpression(node)
+    case node: WhileExpression    => astForWhileStatement(node)
+    case node: DoWhileExpression  => astForDoWhileStatement(node)
+    case node: UntilExpression    => astForUntilStatement(node)
+    case node: CaseExpression     => blockAst(NewBlock(), astsForCaseExpression(node).toList)
+    case node: IfExpression       => astForIfExpression(node)
+    case node: UnlessExpression   => astForUnlessStatement(node)
+    case node: ForExpression      => astForForExpression(node)
+    case node: RescueExpression   => astForRescueExpression(node)
+    case node: NextExpression     => astForNextExpression(node)
+    case node: BreakExpression    => astForBreakExpression(node)
+    case node: OperatorAssignment => astForOperatorAssignmentExpression(node)
   }
 
   private def astForWhileStatement(node: WhileExpression): Ast = {
@@ -295,4 +298,9 @@ trait AstForControlStructuresCreator(implicit withSchemaValidation: ValidationMo
     astsForStatement(generatedNode)
   }
 
+  private def astForOperatorAssignmentExpression(node: OperatorAssignment): Ast = {
+    val loweredAssignment = lowerAssignmentOperator(node.lhs, node.rhs, node.op, node.span)
+    astForControlStructureExpression(loweredAssignment)
+  }
+
 }

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

@@ -2,6 +2,7 @@ package io.joern.rubysrc2cpg.astcreation
 
 import io.joern.rubysrc2cpg.astcreation.RubyIntermediateAst.{Unknown, Block as RubyBlock, *}
 import io.joern.rubysrc2cpg.datastructures.BlockScope
+import io.joern.rubysrc2cpg.parser.RubyJsonHelpers
 import io.joern.rubysrc2cpg.passes.Defines
 import io.joern.rubysrc2cpg.passes.GlobalTypes
 import io.joern.rubysrc2cpg.passes.Defines.{RubyOperators, getBuiltInType}
@@ -482,7 +483,14 @@ trait AstForExpressionsCreator(implicit withSchemaValidation: ValidationMode) {
                     elseAssignNil
                   )
 
-                astForExpression(transform(cfNode))
+                cfNode match {
+                  case x @ OperatorAssignment(lhs, op, rhs) =>
+                    val loweredNode = lowerAssignmentOperator(lhs, rhs, op, x.span)
+                    astForExpression(transform(loweredNode))
+                  case x =>
+                    astForExpression(transform(cfNode))
+                }
+
               case _ =>
                 // The if the LHS defines a new variable, put the local variable into scope
                 val lhsAst = node.lhs match {

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

@@ -2,6 +2,7 @@ package io.joern.rubysrc2cpg.astcreation
 
 import io.joern.rubysrc2cpg.astcreation.RubyIntermediateAst.{RubyStatement, *}
 import io.joern.rubysrc2cpg.datastructures.BlockScope
+import io.joern.rubysrc2cpg.parser.RubyJsonHelpers
 import io.joern.rubysrc2cpg.passes.Defines
 import io.joern.rubysrc2cpg.passes.Defines.getBuiltInType
 import io.joern.x2cpg.{Ast, ValidationMode}
@@ -14,6 +15,7 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
     baseAstCache.clear() // A safe approximation on where to reset the cache
     node match {
       case node: IfExpression               => astForIfStatement(node)
+      case node: OperatorAssignment         => astForOperatorAssignment(node)
       case node: CaseExpression             => astsForCaseExpression(node)
       case node: StatementList              => astForStatementList(node) :: Nil
       case node: ReturnExpression           => astForReturnExpression(node) :: Nil
@@ -48,6 +50,11 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
     }
   }
 
+  private def astForOperatorAssignment(node: OperatorAssignment): Seq[Ast] = {
+    val loweredAssignment = lowerAssignmentOperator(node.lhs, node.rhs, node.op, node.span)
+    astsForStatement(loweredAssignment)
+  }
+
   private def astForJsonIfStatement(node: IfExpression): Seq[Ast] = {
     val conditionAst = astForExpression(node.condition)
     val thenAst      = astForThenClause(node.thenClause)
@@ -177,7 +184,14 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
       case expr: ControlFlowStatement =>
         def transform(e: RubyExpression & ControlFlowStatement): RubyExpression =
           transformLastRubyNodeInControlFlowExpressionBody(e, returnLastNode(_, transform), elseReturnNil)
-        astsForStatement(transform(expr))
+
+        expr match {
+          case x @ OperatorAssignment(lhs, op, rhs) =>
+            val loweredAssignment = lowerAssignmentOperator(lhs, rhs, op, x.span)
+            astsForStatement(transform(loweredAssignment))
+          case x =>
+            astsForStatement(transform(expr))
+        }
       case node: MemberCallWithBlock => returnAstForRubyCall(node)
       case node: SimpleCallWithBlock => returnAstForRubyCall(node)
       case _: (LiteralExpr | BinaryExpression | UnaryExpression | SimpleIdentifier | SelfIdentifier | IndexAccess |
@@ -208,9 +222,7 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
         val nilReturnLiteral = StaticLiteral(Defines.NilClass)(nilReturnSpan)
         stmts.map(astForExpression) ++ astsForImplicitReturnStatement(nilReturnLiteral)
       case node =>
-        logger.warn(
-          s"Implicit return here not supported yet: ${node.text} (${node.getClass.getSimpleName}), only generating statement"
-        )
+        logger.warn(s" not supported yet: ${node.text} (${node.getClass.getSimpleName}), only generating statement")
         astsForStatement(node).toList
   }
 
@@ -311,6 +323,9 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
       case WhileExpression(condition, body)   => WhileExpression(condition, transform(body))(node.span)
       case DoWhileExpression(condition, body) => DoWhileExpression(condition, transform(body))(node.span)
       case UntilExpression(condition, body)   => UntilExpression(condition, transform(body))(node.span)
+      case OperatorAssignment(lhs, op, rhs) =>
+        val loweredNode = lowerAssignmentOperator(lhs, rhs, op, node.span)
+        transformLastRubyNodeInControlFlowExpressionBody(loweredNode, transform, defaultElseBranch)
       case IfExpression(condition, thenClause, elsifClauses, elseClause) =>
         IfExpression(
           condition,

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

@@ -210,6 +210,11 @@ object RubyIntermediateAst {
     def assignments: List[SingleAssignment]
   }
 
+  final case class OperatorAssignment(lhs: RubyExpression, op: String, rhs: RubyExpression)(span: TextSpan)
+      extends RubyExpression(span)
+      with RubyStatement
+      with ControlFlowStatement
+
   final case class DefaultMultipleAssignment(assignments: List[SingleAssignment])(span: TextSpan)
       extends RubyExpression(span)
       with MultipleAssignment

joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/parser/RubyJsonHelpers.scala

@@ -5,6 +5,7 @@ import io.joern.rubysrc2cpg.astcreation.RubyIntermediateAst.{
   AllowedTypeDeclarationChild,
   ArrayLiteral,
   ClassFieldIdentifier,
+  ControlFlowStatement,
   DefaultMultipleAssignment,
   IfExpression,
   MemberAccess,
@@ -167,43 +168,6 @@ object RubyJsonHelpers {
     }
   }
 
-  /** Lowers the `||=` and `&&=` assignment operators to the respective `.nil?` checks
-    */
-  def lowerAssignmentOperator(lhs: RubyExpression, rhs: RubyExpression, op: String, span: TextSpan): RubyExpression = {
-    val condition  = nilCheckCondition(lhs, op, "nil?", span)
-    val thenClause = nilCheckThenClause(lhs, rhs, span)
-    nilCheckIfStatement(condition, thenClause, span)
-  }
-
-  /** Generates the required `.nil?` check condition used in the lowering of `||=` and `&&=`
-    */
-  private def nilCheckCondition(lhs: RubyExpression, op: String, memberName: String, span: TextSpan): RubyExpression = {
-    val memberAccess =
-      MemberAccess(lhs, op = ".", memberName = "nil?")(span.spanStart(s"${lhs.span.text}.nil?"))
-    if op == "||=" then memberAccess
-    else UnaryExpression(op = "!", expression = memberAccess)(span.spanStart(s"!${memberAccess.span.text}"))
-  }
-
-  /** Generates the assignment and the `thenClause` used in the lowering of `||=` and `&&=`
-    */
-  private def nilCheckThenClause(lhs: RubyExpression, rhs: RubyExpression, span: TextSpan): RubyExpression = {
-    StatementList(List(SingleAssignment(lhs, "=", rhs)(span.spanStart(s"${lhs.span.text} = ${rhs.span.text}"))))(
-      span.spanStart(s"${lhs.span.text} = ${rhs.span.text}")
-    )
-  }
-
-  /** Generates the if statement for the lowering of `||=` and `&&=`
-    */
-  private def nilCheckIfStatement(
-    condition: RubyExpression,
-    thenClause: RubyExpression,
-    span: TextSpan
-  ): RubyExpression = {
-    IfExpression(condition = condition, thenClause = thenClause, elsifClauses = List.empty, elseClause = None)(
-      span.spanStart(s"if ${condition.span.text} then ${thenClause.span.text} end")
-    )
-  }
-
   def lowerMultipleAssignment(
     obj: ujson.Obj,
     lhsNodes: List[RubyExpression],

joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/parser/RubyJsonToNodeCreator.scala

@@ -208,7 +208,7 @@ class RubyJsonToNodeCreator(
       case x                         => x
     }
     val rhs = visit(obj(ParserKeys.Rhs))
-    lowerAssignmentOperator(lhs, rhs, "&&=", obj.toTextSpan)
+    OperatorAssignment(lhs, "&&=", rhs)(obj.toTextSpan)
   }
 
   private def visitArg(obj: Obj): RubyExpression = MandatoryParameter(obj(ParserKeys.Value).str)(obj.toTextSpan)
@@ -732,7 +732,7 @@ class RubyJsonToNodeCreator(
       case x                         => x
     }
     val rhs = visit(obj(ParserKeys.Rhs))
-    lowerAssignmentOperator(lhs, rhs, "||=", obj.toTextSpan)
+    OperatorAssignment(lhs, "||=", rhs)(obj.toTextSpan)
   }
 
   private def visitPair(obj: Obj): RubyExpression = {