Add begin-rescue-else-ensure to intermediate ast

joernio · Dec 19, 2023 · 0c8f76f · 0c8f76f
1 parent 216c88a
commit 0c8f76f
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
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
@@ -22,6 +22,12 @@ object RubyIntermediateAst {
     def text: String = span.text
   }
 
+  implicit class RubyNodeHelper(node: RubyNode) {
+    def asStatementList = node match
+      case stmtList: StatementList => stmtList
+      case _                       => StatementList(List(node))(node.span)
+  }
+
   final case class Unknown()(span: TextSpan) extends RubyNode(span)
 
   final case class StatementList(statements: List[RubyNode])(span: TextSpan) extends RubyNode(span) {
@@ -70,6 +76,12 @@ object RubyIntermediateAst {
     span: TextSpan
   ) extends RubyNode(span)
 
+  final case class RescueExpression(body: RubyNode, rescueClauses: List[RubyNode], elseClause: Option[RubyNode], ensureClause: Option[RubyNode])(span: TextSpan) extends RubyNode(span)
+
+  final case class RescueClause(exceptionClassList: Option[RubyNode], assignment: Option[RubyNode], thenClause: RubyNode)(span: TextSpan) extends RubyNode(span)
+
+  final case class EnsureClause(thenClause: RubyNode)(span: TextSpan) extends RubyNode(span)
+
   final case class WhileExpression(condition: RubyNode, body: RubyNode)(span: TextSpan) extends RubyNode(span)
 
   final case class UntilExpression(condition: RubyNode, body: RubyNode)(span: TextSpan) extends RubyNode(span)

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
@@ -7,6 +7,9 @@ import io.joern.rubysrc2cpg.passes.Defines.getBuiltInType
 import org.antlr.v4.runtime.tree.{ErrorNode, ParseTree, TerminalNode}
 
 import scala.jdk.CollectionConverters.*
+import io.joern.rubysrc2cpg.parser.RubyParser.RescueClauseContext
+import io.joern.rubysrc2cpg.parser.RubyParser.EnsureClauseContext
+import io.joern.rubysrc2cpg.parser.RubyParser.ExceptionClassListContext
 
 /** Converts an ANTLR Ruby Parse Tree into the intermediate Ruby AST.
   */
@@ -540,10 +543,32 @@ class RubyNodeCreator extends RubyParserBaseVisitor[RubyNode] {
     if (ctx.rescueClause().isEmpty && Option(ctx.elseClause()).isEmpty && Option(ctx.ensureClause()).isEmpty) {
       visit(ctx.compoundStatement())
     } else {
-      Unknown()(ctx.toTextSpan)
+      RescueExpression(
+        visit(ctx.compoundStatement()), 
+        Option(ctx.rescueClause.asScala).fold(List())(_.map(visit).toList), 
+        Option(ctx.elseClause).map(visit), 
+        Option(ctx.ensureClause).map(visit)
+      )(ctx.toTextSpan)
     }
   }
 
+  override def visitExceptionClassList(ctx: ExceptionClassListContext): RubyNode = {
+    // Requires implementing multiple rhs with splatting
+    Unknown()(ctx.toTextSpan)
+  }
+
+  override def visitRescueClause(ctx: RescueClauseContext): RubyNode = {
+    RescueClause(
+      Option(ctx.exceptionClassList).map(visit), 
+      Option(ctx.exceptionVariableAssignment).map(visit), 
+      visit(ctx.thenClause)
+    )(ctx.toTextSpan)
+  }
+
+  override def visitEnsureClause(ctx: EnsureClauseContext): RubyNode = {
+    EnsureClause(visit(ctx.compoundStatement()))(ctx.toTextSpan)
+  }
+
   override def visitAssociationKey(ctx: RubyParser.AssociationKeyContext): RubyNode = {
     if (Option(ctx.operatorExpression()).isDefined) {
       visit(ctx.operatorExpression())