[ruby] Begin-End-Until Loop (#4523)

Handling `Begin-End-Until` as a lowered `do-while` loop. Resolves #4507
joernio · May 2, 2024 · 994bb69 · 994bb69
1 parent 6bab41b
commit 994bb69
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 2 deletions.
diff --git a/...rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForStatementsCreator.scala b/...rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForStatementsCreator.scala
@@ -12,6 +12,7 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
 
   protected def astsForStatement(node: RubyNode): Seq[Ast] = node match
     case node: WhileExpression            => astForWhileStatement(node) :: Nil
+    case node: DoWhileExpression          => astForDoWhileStatement(node) :: Nil
     case node: UntilExpression            => astForUntilStatement(node) :: Nil
     case node: IfExpression               => astForIfStatement(node) :: Nil
     case node: UnlessExpression           => astForUnlessStatement(node) :: Nil
@@ -36,6 +37,12 @@ trait AstForStatementsCreator(implicit withSchemaValidation: ValidationMode) { t
     whileAst(Some(conditionAst), bodyAsts, Option(code(node)), line(node), column(node))
   }
 
+  private def astForDoWhileStatement(node: DoWhileExpression): Ast = {
+    val conditionAst = astForExpression(node.condition)
+    val bodyAsts     = astsForStatement(node.body)
+    doWhileAst(Some(conditionAst), bodyAsts, Option(code(node)), line(node), column(node))
+  }
+
   // `until T do B` is lowered as `while !T do B`
   private def astForUntilStatement(node: UntilExpression): Ast = {
     val notCondition = astForExpression(UnaryExpression("!", node.condition)(node.condition.span))

diff --git a/...nds/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/RubyIntermediateAst.scala b/...nds/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/RubyIntermediateAst.scala
@@ -169,6 +169,10 @@ object RubyIntermediateAst {
       extends RubyNode(span)
       with ControlFlowExpression
 
+  final case class DoWhileExpression(condition: RubyNode, body: RubyNode)(span: TextSpan)
+      extends RubyNode(span)
+      with ControlFlowExpression
+
   final case class UntilExpression(condition: RubyNode, body: RubyNode)(span: TextSpan)
       extends RubyNode(span)
       with ControlFlowExpression

diff --git a/...li/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/parser/RubyNodeCreator.scala b/...li/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/parser/RubyNodeCreator.scala
@@ -68,6 +68,10 @@ class RubyNodeCreator extends RubyParserBaseVisitor[RubyNode] {
     UntilExpression(condition, body)(ctx.toTextSpan)
   }
 
+  override def visitBeginEndExpression(ctx: RubyParser.BeginEndExpressionContext): RubyNode = {
+    visit(ctx.bodyStatement())
+  }
+
   override def visitIfExpression(ctx: RubyParser.IfExpressionContext): RubyNode = {
     val condition = visit(ctx.commandOrPrimaryValue())
     val thenBody  = visit(ctx.thenClause())
@@ -120,7 +124,12 @@ class RubyNodeCreator extends RubyParserBaseVisitor[RubyNode] {
         val condition = visit(ctx.expressionOrCommand())
         val body      = visit(ctx.statement())
         WhileExpression(condition, body)(ctx.toTextSpan)
+      case "until" =>
+        val condition = visit(ctx.expressionOrCommand())
+        val body      = visit(ctx.statement())
+        DoWhileExpression(condition, body)(ctx.toTextSpan)
       case _ =>
+        logger.warn(s"Unhandled modifier statement ${ctx.getClass}")
         Unknown()(ctx.toTextSpan)
   }
 

diff --git a/...ends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/dataflow/ControlStructureTests.scala b/...ends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/dataflow/ControlStructureTests.scala
@@ -55,7 +55,7 @@ class ControlStructureTests extends RubyCode2CpgFixture(withPostProcessing = tru
   }
 
   // Works in deprecated - does not parse in new frontend
-  "flow in through until modifier" ignore {
+  "flow in through until modifier" in {
     val cpg = code("""
                      |i = 0
                      |num = 5
@@ -67,7 +67,7 @@ class ControlStructureTests extends RubyCode2CpgFixture(withPostProcessing = tru
 
     val src  = cpg.identifier.name("i").l
     val sink = cpg.call.name("puts").l
-    sink.reachableByFlows(src).l.size shouldBe 3
+    sink.reachableByFlows(src).size shouldBe 3
   }
 
   "flow through for loop" in {

diff --git a/...ends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/ControlStructureTests.scala b/...ends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/ControlStructureTests.scala
@@ -28,6 +28,30 @@ class ControlStructureTests extends RubyCode2CpgFixture {
     assignment.lineNumber shouldBe Some(4)
   }
 
+  "begin-end-until should be lowered as a do-while loop" in {
+
+    val cpg = code("""
+        |i = 0
+        |num = 5
+        |begin
+        |  num = i + 3
+        |end until i < num
+        |puts num
+        |""".stripMargin)
+
+    val List(whileNode)  = cpg.doBlock.l
+    val List(whileCond)  = whileNode.condition.isCall.l
+    val List(assignment) = whileNode.astChildren.isBlock.assignment.l
+
+    whileCond.methodFullName shouldBe Operators.lessThan
+    whileCond.code shouldBe "i < num"
+    whileCond.lineNumber shouldBe Some(6)
+
+    assignment.code shouldBe "num = i + 3"
+    assignment.lineNumber shouldBe Some(5)
+
+  }
+
   "`until-end` statement is represented by a negated `WHILE` CONTROL_STRUCTURE node" in {
     val cpg = code("""
         |x = 1