Add ModuleBase base class and combine class/module tests

ql/src/codeql_ruby/ast/Expr.qll

@@ -142,6 +142,15 @@ class ExprSequence extends Expr {
   final predicate isEmpty() { this.getNumberOfExpressions() = 0 }
 }
 
+/**
+ * A sequence of expressions representing the body of a method, class, module,
+ * or do-block. That is, any body that may also include rescue/ensure/else
+ * statements.
+ */
+class BodyStatement extends ExprSequence {
+  override BodyStatement::Range range;
+}
+
 /**
  * A scope resolution, typically used to access constants defined in a class or
  * module.

ql/src/codeql_ruby/ast/Method.qll

@@ -18,7 +18,7 @@ class Callable extends Expr, CfgScope {
 }
 
 /** A method. */
-class Method extends Callable, @method {
+class Method extends Callable, BodyStatement, @method {
   final override Method::Range range;
 
   final override string getAPrimaryQlClass() { result = "Method" }
@@ -42,7 +42,7 @@ class Method extends Callable, @method {
 }
 
 /** A singleton method. */
-class SingletonMethod extends Callable, @singleton_method {
+class SingletonMethod extends Callable, BodyStatement, @singleton_method {
   final override SingletonMethod::Range range;
 
   final override string getAPrimaryQlClass() { result = "SingletonMethod" }
@@ -73,7 +73,7 @@ class Block extends AstNode, Callable {
 }
 
 /** A block enclosed within `do` and `end`. */
-class DoBlock extends Block, @do_block {
+class DoBlock extends Block, BodyStatement, @do_block {
   final override DoBlock::Range range;
 
   final override string getAPrimaryQlClass() { result = "DoBlock" }

ql/src/codeql_ruby/ast/Module.qll

@@ -1,6 +1,31 @@
 private import codeql_ruby.AST
 private import internal.Module
 
+/**
+ * The base class for classes, singleton classes, and modules.
+ */
+class ModuleBase extends BodyStatement {
+  override ModuleBase::Range range;
+
+  /** Gets a method defined in this module/class. */
+  Method getAMethod() { result = this.getAnExpr() }
+
+  /** Gets the method named `name` in this module/class, if any. */
+  Method getMethod(string name) { result = this.getAMethod() and result.getName() = name }
+
+  /** Gets a class defined in this module/class. */
+  Class getAClass() { result = this.getAnExpr() }
+
+  /** Gets the class named `name` in this module/class, if any. */
+  Class getClass(string name) { result = this.getAClass() and result.getName() = name }
+
+  /** Gets a module defined in this module/class. */
+  Module getAModule() { result = this.getAnExpr() }
+
+  /** Gets the module named `name` in this module/class, if any. */
+  Module getModule(string name) { result = this.getAModule() and result.getName() = name }
+}
+
 /**
  * A class definition.
  *
@@ -11,7 +36,7 @@ private import internal.Module
  * end
  * ```
  */
-class Class extends ExprSequence {
+class Class extends ModuleBase {
   final override Class::Range range;
 
   final override string getAPrimaryQlClass() { result = "Class" }
@@ -54,17 +79,6 @@ class Class extends ExprSequence {
    */
   final ScopeResolution getNameScopeResolution() { result = range.getNameScopeResolution() }
 
-  /**
-   * Gets a method defined in this class.
-   * ```rb
-   * class Foo
-   *   def bar
-   *   end
-   * end
-   * ```
-   */
-  final Method getAMethod() { result = this.getAnExpr() }
-
   /**
    * Gets the `Expr` used as the superclass in the class definition, if any.
    *
@@ -91,7 +105,7 @@ class Class extends ExprSequence {
  * end
  * ```
  */
-class SingletonClass extends ExprSequence, @singleton_class {
+class SingletonClass extends ModuleBase, @singleton_class {
   final override SingletonClass::Range range;
 
   final override string getAPrimaryQlClass() { result = "Class" }
@@ -108,17 +122,6 @@ class SingletonClass extends ExprSequence, @singleton_class {
    * ```
    */
   final Expr getValue() { result = range.getValue() }
-
-  /**
-   * Gets a method defined in this singleton class.
-   * ```rb
-   * class << foo
-   *   def bar
-   *   end
-   * end
-   * ```
-   */
-  final Method getAMethod() { result = this.getAnExpr() }
 }
 
 /**
@@ -146,7 +149,7 @@ class SingletonClass extends ExprSequence, @singleton_class {
  * end
  * ```
  */
-class Module extends ExprSequence, @module {
+class Module extends ModuleBase, @module {
   final override Module::Range range;
 
   final override string getAPrimaryQlClass() { result = "Module" }
@@ -188,15 +191,4 @@ class Module extends ExprSequence, @module {
    * ```
    */
   final ScopeResolution getNameScopeResolution() { result = range.getNameScopeResolution() }
-
-  /**
-   * Gets a method defined in this module.
-   * ```rb
-   * module Foo
-   *   def bar
-   *   end
-   * end
-   * ```
-   */
-  final Method getAMethod() { result = this.getAnExpr() }
 }

ql/src/codeql_ruby/ast/internal/Expr.qll

@@ -168,6 +168,10 @@ module ExprSequence {
   }
 }
 
+module BodyStatement {
+  abstract class Range extends ExprSequence::Range { }
+}
+
 module ThenExpr {
   class Range extends ExprSequence::Range, @then {
     final override Generated::Then generated;

ql/src/codeql_ruby/ast/internal/Method.qll

@@ -9,7 +9,7 @@ module Callable {
 }
 
 module Method {
-  class Range extends Callable::Range, @method {
+  class Range extends Callable::Range, BodyStatement::Range, @method {
     final override Generated::Method generated;
 
     override Parameter getParameter(int n) { result = generated.getParameters().getChild(n) }
@@ -21,11 +21,13 @@ module Method {
     }
 
     final predicate isSetter() { generated.getName() instanceof Generated::Setter }
+
+    final override Expr getExpr(int i) { result = generated.getChild(i) }
   }
 }
 
 module SingletonMethod {
-  class Range extends Callable::Range, @singleton_method {
+  class Range extends Callable::Range, BodyStatement::Range, @singleton_method {
     final override Generated::SingletonMethod generated;
 
     override Parameter getParameter(int n) { result = generated.getParameters().getChild(n) }
@@ -35,6 +37,8 @@ module SingletonMethod {
       result = generated.getName().(SymbolLiteral).getValueText() or
       result = generated.getName().(Generated::Setter).getName().getValue() + "="
     }
+
+    final override Expr getExpr(int i) { result = generated.getChild(i) }
   }
 }
 
@@ -51,10 +55,12 @@ module Block {
 }
 
 module DoBlock {
-  class Range extends Block::Range, @do_block {
+  class Range extends Block::Range, BodyStatement::Range, @do_block {
     final override Generated::DoBlock generated;
 
     final override Parameter getParameter(int n) { result = generated.getParameters().getChild(n) }
+
+    final override Expr getExpr(int i) { result = generated.getChild(i) }
   }
 }
 

ql/src/codeql_ruby/ast/internal/Module.qll

@@ -2,8 +2,12 @@ private import codeql_ruby.AST
 private import codeql_ruby.ast.internal.Expr
 private import codeql_ruby.ast.internal.TreeSitter
 
+module ModuleBase {
+  abstract class Range extends BodyStatement::Range { }
+}
+
 module Class {
-  class Range extends ExprSequence::Range, @class {
+  class Range extends ModuleBase::Range, @class {
     final override Generated::Class generated;
 
     final override Expr getExpr(int i) { result = generated.getChild(i) }
@@ -20,18 +24,17 @@ module Class {
 }
 
 module SingletonClass {
-  class Range extends ExprSequence::Range, @singleton_class {
+  class Range extends ModuleBase::Range, @singleton_class {
     final override Generated::SingletonClass generated;
 
     final override Expr getExpr(int i) { result = generated.getChild(i) }
 
     final Expr getValue() { result = generated.getValue() }
-
   }
 }
 
 module Module {
-  class Range extends ExprSequence::Range, @module {
+  class Range extends ModuleBase::Range, @module {
     final override Generated::Module generated;
 
     final override Expr getExpr(int n) { result = generated.getChild(n) }