Repository: joernio/joern (scala)

Issue:
[rubysrc2cpg] Generate methods under type declarations
**Is your feature request related to a problem? Please describe.**
Methods defined under classes do not generate method nodes under the type decl nodes. See the test under `ast/TypeDeclAstCreationPassTest:21`

Pull request description:
* Added a few foundational tests for type declarations
* Fixed an instance where a null `StatementsContext` would throw an exception when parsing some moderately invalid Ruby
* Solved methods under type declarations by handling class bodies differently to successfully construct methods and their modifiers correctly

Resolves #3005 and partially resolves #3008

Patch:
diff --git a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstCreator.scala b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstCreator.scala
index cc2df9921b11..e2cc3b85a99b 100644
--- a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstCreator.scala
+++ b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstCreator.scala
@@ -1,5 +1,5 @@
 package io.joern.rubysrc2cpg.astcreation
-import io.joern.rubysrc2cpg.parser.RubyParser._
+import io.joern.rubysrc2cpg.parser.RubyParser.*
 import io.joern.rubysrc2cpg.parser.{RubyLexer, RubyParser}
 import io.joern.rubysrc2cpg.passes.Defines
 import io.joern.rubysrc2cpg.utils.PackageContext
@@ -7,18 +7,18 @@ import io.joern.x2cpg.Ast.storeInDiffGraph
 import io.joern.x2cpg.Defines.DynamicCallUnknownFullName
 import io.joern.x2cpg.datastructures.{Global, Scope}
 import io.joern.x2cpg.{Ast, AstCreatorBase, AstNodeBuilder}
-import io.shiftleft.codepropertygraph.generated._
-import io.shiftleft.codepropertygraph.generated.nodes._
+import io.shiftleft.codepropertygraph.generated.*
+import io.shiftleft.codepropertygraph.generated.nodes.*
 import org.antlr.v4.runtime.tree.TerminalNode
 import org.antlr.v4.runtime.{CharStreams, CommonTokenStream, ParserRuleContext, Token}
 import org.slf4j.LoggerFactory
 import overflowdb.BatchedUpdate
 
-import java.io.{File => JFile}
+import java.io.File as JFile
 import scala.collection.immutable.Seq
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
-import scala.jdk.CollectionConverters._
+import scala.jdk.CollectionConverters.*
 
 class AstCreator(
   protected val filename: String,
@@ -1049,11 +1049,33 @@ class AstCreator(
     blockAst(blockNode, asts.toList)
   }
 
+  /** Handles body statements differently from [[astForBodyStatementContext]] by noting that method definitions should
+    * be on the root level and assignments where the LHS starts with @@ should be treated as fields.
+    */
+  def astForClassBody(ctx: BodyStatementContext): Seq[Ast] = {
+    val rootStatements =
+      Option(ctx).map(_.compoundStatement()).map(_.statements()).map(astForStatements).getOrElse(Seq())
+
+    val (methods, blockStmts) =
+      rootStatements
+        .flatMap { ast =>
+          ast.root match
+            case Some(x: NewMethod)                                 => Seq(ast)
+            case Some(x: NewCall) if x.name == Operators.assignment => Seq(ast) :+ astsForClassMembers(ast)
+            case _                                                  => Seq(ast)
+        }
+        .partition(_.root match
+          case Some(_: NewMethod) => true
+          case _                  => false
+        )
+    Seq(blockAst(blockNode(ctx), blockStmts.toList)) ++ methods
+  }
+
   def astForBodyStatementContext(ctx: BodyStatementContext, addReturnNode: Boolean = false): Seq[Ast] = {
     val compoundStatementAsts = astForCompoundStatement(ctx.compoundStatement())
 
     val compoundStatementAstsWithReturn =
-      if (addReturnNode && compoundStatementAsts.size > 0) {
+      if (addReturnNode && compoundStatementAsts.nonEmpty) {
         /*
          * Convert the last statement to a return AST if it is not already a return AST.
          * If it is a return AST leave it untouched.
@@ -1090,7 +1112,7 @@ class AstCreator(
     val rescueAsts = ctx
       .rescueClause()
       .asScala
-      .map(astForRescueClauseContext(_))
+      .map(astForRescueClauseContext)
       .toSeq
 
     if (ctx.elseClause() != null) {
@@ -1165,7 +1187,10 @@ class AstCreator(
       .columnNumber(None)
       .typeFullName(Defines.Any)
 
-    val publicModifier = NewModifier().modifierType(ModifierTypes.PUBLIC)
+    val modifierNode = lastModifier match {
+      case Some(modifier) => NewModifier().modifierType(modifier).code(modifier)
+      case None           => NewModifier().modifierType(ModifierTypes.PUBLIC).code(ModifierTypes.PUBLIC)
+    }
     /*
      * public/private/protected modifiers are in a separate statement
      * TODO find out how they should be used. Need to do this iff it adds any value
@@ -1185,7 +1210,7 @@ class AstCreator(
         paramSeq,
         blockAst(blockNode, astBody.toList),
         methodRetNode,
-        Seq[NewModifier](publicModifier)
+        Seq[NewModifier](modifierNode)
       )
     )
   }
diff --git a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala
index 018046199c47..2ea3c845ee0f 100644
--- a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala
+++ b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala
@@ -4,13 +4,15 @@ import io.joern.rubysrc2cpg.parser.RubyParser.*
 import io.joern.rubysrc2cpg.passes.Defines
 import io.joern.x2cpg.Ast
 import io.shiftleft.codepropertygraph.generated.nodes.NewJumpTarget
-import io.shiftleft.codepropertygraph.generated.{ControlStructureTypes, DispatchTypes, Operators}
+import io.shiftleft.codepropertygraph.generated.{ControlStructureTypes, DispatchTypes, ModifierTypes, Operators}
 import org.antlr.v4.runtime.ParserRuleContext
 
 import scala.jdk.CollectionConverters.CollectionHasAsScala
 
 trait AstForExpressionsCreator { this: AstCreator =>
 
+  protected var lastModifier: Option[String] = None
+
   protected def astForPowerExpression(ctx: PowerExpressionContext): Ast =
     astForBinaryOperatorExpression(ctx, Operators.exponentiation, ctx.expression().asScala)
 
@@ -210,6 +212,9 @@ trait AstForExpressionsCreator { this: AstCreator =>
       Ast(node)
     } else if (methodNames.contains(variableName)) {
       astForCallNode(ctx, variableName)
+    } else if (ModifierTypes.ALL.contains(variableName.toUpperCase)) {
+      lastModifier = Option(variableName.toUpperCase)
+      Ast()
     } else {
       val node = createIdentifierWithScope(ctx, variableName, variableName, Defines.Any, List())
       Ast(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 4e0803b7bf23..0ba1d718696e 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
@@ -81,14 +81,14 @@ trait AstForStatementsCreator {
   }
 
   protected def astForCompoundStatement(ctx: CompoundStatementContext): Seq[Ast] = {
-    val stmtAsts = Option(ctx.statements()).map(astForStatements).getOrElse(Seq())
+    val stmtAsts = Option(ctx).map(_.statements()).map(astForStatements).getOrElse(Seq())
     Seq(blockAst(blockNode(ctx), stmtAsts.toList))
   }
 
   protected def astForStatements(ctx: StatementsContext): Seq[Ast] = {
     Option(ctx) match {
       case Some(ctx) =>
-        Option(ctx.statement()).map(_.asScala).getOrElse(Seq()).flatMap(astForStatement).toSeq
+        Option(ctx).map(_.statement()).map(_.asScala).getOrElse(Seq()).flatMap(astForStatement).toSeq
       case None =>
         Seq()
     }
diff --git a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForTypesCreator.scala b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForTypesCreator.scala
index dfcd9dbc873a..ee53a9e8246b 100644
--- a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForTypesCreator.scala
+++ b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForTypesCreator.scala
@@ -28,21 +28,7 @@ trait AstForTypesCreator { this: AstCreator =>
       classStack.push(className)
       val fullName = classStack.reverse.mkString(":")
 
-      val bodyAst = astForBodyStatementContext(ctx.classDefinition().bodyStatement())
-      val bodyAstSansModifiers = bodyAst
-        .filterNot(ast => {
-          val nodes = ast.nodes
-            .filter(_.isInstanceOf[NewIdentifier])
-
-          if (nodes.size == 1) {
-            val varName = nodes
-              .map(_.asInstanceOf[NewIdentifier].name)
-              .head
-            varName == "public" || varName == "protected" || varName == "private"
-          } else {
-            false
-          }
-        })
+      val bodyAst = astForClassBody(ctx.classDefinition().bodyStatement())
 
       if (classStack.nonEmpty) {
         classStack.pop()
@@ -51,7 +37,7 @@ trait AstForTypesCreator { this: AstCreator =>
       val typeDeclNode = NewTypeDecl()
         .name(className)
         .fullName(fullName)
-      Seq(Ast(typeDeclNode).withChildren(bodyAstSansModifiers))
+      Seq(Ast(typeDeclNode).withChildren(bodyAst))
     } else {
       Seq.empty
     }
@@ -95,8 +81,9 @@ trait AstForTypesCreator { this: AstCreator =>
     }
   }
 
-  def astsForClassMembers(): Seq[Ast] = {
-    ???
+  def astsForClassMembers(ast: Ast): Ast = {
+    // TODO: Handle members
+    Ast()
   }
 
   implicit class ClassDefinitionPrimaryContextExt(val ctx: ClassDefinitionPrimaryContext) {


Test patch:
diff --git a/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/passes/ast/TypeDeclAstCreationPassTest.scala b/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/passes/ast/TypeDeclAstCreationPassTest.scala
index 37e8602dffde..67b92ef9a72a 100644
--- a/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/passes/ast/TypeDeclAstCreationPassTest.scala
+++ b/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/passes/ast/TypeDeclAstCreationPassTest.scala
@@ -1,7 +1,9 @@
 package io.joern.rubysrc2cpg.passes.ast
 
 import io.joern.rubysrc2cpg.testfixtures.RubyCode2CpgFixture
+import io.shiftleft.codepropertygraph.generated.ModifierTypes
 import io.shiftleft.semanticcpg.language._
+
 class TypeDeclAstCreationPassTest extends RubyCode2CpgFixture {
 
   "AST generation for simple classes declarations" should {
@@ -16,6 +18,203 @@ class TypeDeclAstCreationPassTest extends RubyCode2CpgFixture {
       myClass.fullName shouldBe "Test0.rb::program:MyClass"
     }
 
+    "generate methods under type declarations" in {
+      val cpg = code("""
+          |class Vehicle
+          |
+          |   def self.speeding
+          |     "Hello, from a class method"
+          |   end
+          |
+          |   def Vehicle.halting
+          |     "Hello, from another class method"
+          |   end
+          |
+          |   def driving
+          |     "Hello, from an instance method"
+          |   end
+          |
+          |end
+          |""".stripMargin)
+      val Some(vehicle) = cpg.typeDecl.nameExact("Vehicle").headOption: @unchecked
+      vehicle.name shouldBe "Vehicle"
+      vehicle.fullName shouldBe "Test0.rb::program:Vehicle"
+
+      val List(speeding, halting, driving) = vehicle.method.l
+      speeding.name shouldBe "speeding"
+      halting.name shouldBe "halting"
+      driving.name shouldBe "driving"
+
+      speeding.fullName shouldBe "Test0.rb::program:Vehicle:speeding"
+      halting.fullName shouldBe "Test0.rb::program:Vehicle:halting"
+      driving.fullName shouldBe "Test0.rb::program:Vehicle:driving"
+    }
+
+    "generate members for various class members under the respective type declaration" ignore {
+      val cpg = code("""
+          |class Song
+          |  @@plays = 0
+          |  def initialize(name, artist, duration)
+          |    @name     = name
+          |    @artist   = artist
+          |    @duration = duration
+          |  end
+          |end
+          |""".stripMargin)
+      val Some(song) = cpg.typeDecl.nameExact("Song").headOption: @unchecked
+      song.name shouldBe "Song"
+      song.fullName shouldBe "Test0.rb::program:Song"
+
+      val List(plays, name, artist, duration) = song.member.l
+
+      plays.name shouldBe "plays"
+      name.name shouldBe "name"
+      artist.name shouldBe "artist"
+      duration.name shouldBe "duration"
+    }
+
+    "generate members for various class members when using the `attr_reader` and `attr_writer` idioms" ignore {
+      val cpg = code("""
+          |class Song
+          |  attr_reader :name, :artist, :duration
+          |  attr_writer :album
+          |end
+          |""".stripMargin)
+      val Some(song) = cpg.typeDecl.nameExact("Song").headOption: @unchecked
+      song.name shouldBe "Song"
+      song.fullName shouldBe "Test0.rb::program:Song"
+
+      val List(name, artist, duration, album) = song.member.l
+
+      name.name shouldBe "name"
+      artist.name shouldBe "artist"
+      duration.name shouldBe "duration"
+      album.name shouldBe "album"
+    }
+
+    "generate methods with the correct access control modifiers case 1" in {
+      val cpg = code("""
+          |class MyClass
+          |
+          |    def method1    # default is 'public'
+          |      #...
+          |    end
+          |
+          |  protected          # subsequent methods will be 'protected'
+          |
+          |    def method2    # will be 'protected'
+          |      #...
+          |    end
+          |
+          |  private            # subsequent methods will be 'private'
+          |
+          |    def method3    # will be 'private'
+          |      #...
+          |    end
+          |
+          |  public             # subsequent methods will be 'public'
+          |
+          |    def method4    # and this will be 'public'
+          |      #...
+          |    end
+          |end
+          |""".stripMargin)
+      val Some(myClass) = cpg.typeDecl.nameExact("MyClass").headOption: @unchecked
+      myClass.name shouldBe "MyClass"
+      myClass.fullName shouldBe "Test0.rb::program:MyClass"
+
+      val List(m1, m2, m3, m4) = myClass.method.l
+      m1.name shouldBe "method1"
+      m2.name shouldBe "method2"
+      m3.name shouldBe "method3"
+      m4.name shouldBe "method4"
+
+      m1.fullName shouldBe "Test0.rb::program:MyClass:method1"
+      m2.fullName shouldBe "Test0.rb::program:MyClass:method2"
+      m3.fullName shouldBe "Test0.rb::program:MyClass:method3"
+      m4.fullName shouldBe "Test0.rb::program:MyClass:method4"
+
+      m1.modifier.modifierType.l shouldBe List(ModifierTypes.PUBLIC)
+      m2.modifier.modifierType.l shouldBe List(ModifierTypes.PROTECTED)
+      m3.modifier.modifierType.l shouldBe List(ModifierTypes.PRIVATE)
+      m4.modifier.modifierType.l shouldBe List(ModifierTypes.PUBLIC)
+    }
+
+    "generate methods with the correct access control modifiers case 2" ignore {
+      val cpg = code("""
+          |class MyClass
+          |
+          |  def method1
+          |  end
+          |
+          |  def method2
+          |  end
+          |
+          |  def method3
+          |  end
+          |
+          |  def method4
+          |  end
+          |
+          |  public    :method1, :method4
+          |  protected :method2
+          |  private   :method3
+          |end
+          |""".stripMargin)
+      val Some(myClass) = cpg.typeDecl.nameExact("MyClass").headOption: @unchecked
+      myClass.name shouldBe "MyClass"
+      myClass.fullName shouldBe "Test0.rb::program:MyClass"
+
+      val List(m1, m2, m3, m4) = myClass.method.l
+      m1.name shouldBe "method1"
+      m2.name shouldBe "method2"
+      m3.name shouldBe "method3"
+      m4.name shouldBe "method4"
+
+      m1.fullName shouldBe "Test0.rb::program:MyClass:method1"
+      m2.fullName shouldBe "Test0.rb::program:MyClass:method2"
+      m3.fullName shouldBe "Test0.rb::program:MyClass:method3"
+      m4.fullName shouldBe "Test0.rb::program:MyClass:method4"
+
+      m1.modifier.modifierType.l shouldBe List(ModifierTypes.PUBLIC)
+      m2.modifier.modifierType.l shouldBe List(ModifierTypes.PROTECTED)
+      m3.modifier.modifierType.l shouldBe List(ModifierTypes.PRIVATE)
+      m4.modifier.modifierType.l shouldBe List(ModifierTypes.PUBLIC)
+    }
+
+  }
+
+  "Polymorphism in classes" should {
+
+    "correctly contain the inherited base type name in the super type" ignore {
+      val cpg = code("""
+          |class GeeksforGeeks
+          |    def initialize
+          |        puts "This is Superclass"
+          |    end
+          |     
+          |    def super_method
+          |        puts "Method of superclass"
+          |    end
+          |end
+          | 
+          |class Sudo_Placement < GeeksforGeeks
+          |    def initialize
+          |       puts "This is Subclass"
+          |    end
+          |end
+          |""".stripMargin)
+
+      val Some(baseType) = cpg.typeDecl.nameExact("GeeksforGeeks").headOption: @unchecked
+      baseType.name shouldBe "GeeksforGeeks"
+      baseType.fullName shouldBe "Test0.rb::program:GeeksforGeeks"
+
+      val Some(subType) = cpg.typeDecl.nameExact("Sudo_Placement").headOption: @unchecked
+      subType.name shouldBe "Sudo_Placement"
+      subType.fullName shouldBe "Test0.rb::program:Sudo_Placement"
+      subType.inheritsFromTypeFullName shouldBe Seq("Test0.rb::program:GeeksforGeeks")
+    }
+
   }
 
 }