AST: Toplevel and BEGIN/ END blocks

ql/src/codeql_ruby/ast/Module.qll

@@ -27,6 +27,32 @@ class ModuleBase extends BodyStatement {
   Module getModule(string name) { result = this.getAModule() and result.getName() = name }
 }
 
+/**
+ * A Ruby source file.
+ *
+ * ```rb
+ * def main
+ *   puts "hello world!"
+ * end
+ * main
+ * ```
+ */
+class Toplevel extends ModuleBase, @program {
+  final override Toplevel::Range range;
+
+  final override string getAPrimaryQlClass() { result = "Toplevel" }
+
+  /**
+   * Get the `n`th `BEGIN` block.
+   */
+  final StmtSequence getBeginBlock(int n) { result = range.getBeginBlock(n) }
+
+  /**
+   * Get a `BEGIN` block.
+   */
+  final StmtSequence getABeginBlock() { result = getBeginBlock(_) }
+}
+
 /**
  * A class definition.
  *

ql/src/codeql_ruby/ast/Statement.qll

@@ -1,5 +1,6 @@
 private import codeql_ruby.AST
 private import codeql_ruby.CFG
+private import internal.Expr
 private import internal.Statement
 private import internal.Variable
 private import codeql_ruby.controlflow.internal.ControlFlowGraphImpl
@@ -34,6 +35,18 @@ class EmptyStmt extends Stmt, @token_empty_statement {
   final override string getAPrimaryQlClass() { result = "EmptyStmt" }
 }
 
+/**
+ * An `END` block.
+ * ```rb
+ * END{ puts "shutting down" }
+ * ```
+ */
+class EndBlock extends StmtSequence, @end_block {
+  final override EndBlock::Range range;
+
+  final override string getAPrimaryQlClass() { result = "EndBlock" }
+}
+
 /**
  * A statement that may return a value: `return`, `break` and `next`.
  *

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

@@ -20,16 +20,10 @@ module AstNode {
     // an AST node, for example we include the `in` keyword in `for` loops
     // in the CFG, but not the AST
     RemoveWhenFullCoverage() {
-      this instanceof Generated::Program
-      or
       this = any(Generated::Method m).getName()
       or
       this = any(Generated::SingletonMethod m).getName()
       or
-      this instanceof Generated::BeginBlock
-      or
-      this instanceof Generated::EndBlock
-      or
       this = any(Generated::Call c).getMethod() and
       not this instanceof Generated::ScopeResolution
       or

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

@@ -232,6 +232,16 @@ module ParenthesizedExpr {
   }
 }
 
+module BeginBlock {
+  class Range extends StmtSequence::Range, @begin_block {
+    final override Generated::BeginBlock generated;
+
+    final override Stmt getStmt(int n) { result = generated.getChild(n) }
+
+    final override string toString() { result = "BEGIN { ... }" }
+  }
+}
+
 module ThenExpr {
   class Range extends StmtSequence::Range, @then {
     final override Generated::Then generated;

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

@@ -7,6 +7,25 @@ module ModuleBase {
   abstract class Range extends BodyStatement::Range { }
 }
 
+module Toplevel {
+  class Range extends ModuleBase::Range, @program {
+    final override Generated::Program generated;
+
+    Range() { generated.getLocation().getFile().getExtension() != "erb" }
+
+    final override Generated::AstNode getChild(int i) {
+      result = generated.getChild(i) and
+      not result instanceof Generated::BeginBlock
+    }
+
+    final StmtSequence getBeginBlock(int n) {
+      result = rank[n](int i, Generated::BeginBlock b | b = generated.getChild(i) | b order by i)
+    }
+
+    final override string toString() { result = generated.getLocation().getFile().getBaseName() }
+  }
+}
+
 module Class {
   class Range extends ModuleBase::Range, ConstantWriteAccess::Range, @class {
     final override Generated::Class generated;

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

@@ -1,5 +1,6 @@
 private import codeql_ruby.AST
 private import codeql_ruby.ast.internal.AST
+private import codeql_ruby.ast.internal.Expr
 private import codeql_ruby.ast.internal.TreeSitter
 
 module Stmt {
@@ -14,6 +15,16 @@ module EmptyStmt {
   }
 }
 
+module EndBlock {
+  class Range extends StmtSequence::Range, @end_block {
+    final override Generated::EndBlock generated;
+
+    final override Stmt getStmt(int n) { result = generated.getChild(n) }
+
+    final override string toString() { result = "END { ... }" }
+  }
+}
+
 module ReturningStmt {
   abstract class Range extends Stmt::Range {
     abstract Generated::ArgumentList getArgumentList();

ql/test/library-tests/ast/modules/module_base.expected

@@ -1,4 +1,5 @@
 moduleBases
+| classes.rb:2:1:56:3 | classes.rb | Toplevel |
 | classes.rb:3:1:4:3 | Foo | Class |
 | classes.rb:7:1:8:3 | Bar | Class |
 | classes.rb:11:1:12:3 | Baz | Class |
@@ -10,6 +11,7 @@ moduleBases
 | classes.rb:41:1:52:3 | class << ... | Class |
 | classes.rb:55:1:56:3 | MyClassInGlobalScope | Class |
 | modules.rb:1:1:2:3 | Empty | Module |
+| modules.rb:1:1:61:3 | modules.rb | Toplevel |
 | modules.rb:4:1:24:3 | Foo | Module |
 | modules.rb:5:3:14:5 | Bar | Module |
 | modules.rb:6:5:7:7 | ClassInFooBar | Class |
@@ -21,6 +23,12 @@ moduleBases
 | modules.rb:49:3:50:5 | ClassInAnotherDefinitionOfFooBar | Class |
 | modules.rb:60:1:61:3 | MyModuleInGlobalScope | Module |
 moduleBaseClasses
+| classes.rb:2:1:56:3 | classes.rb | classes.rb:3:1:4:3 | Foo |
+| classes.rb:2:1:56:3 | classes.rb | classes.rb:7:1:8:3 | Bar |
+| classes.rb:2:1:56:3 | classes.rb | classes.rb:11:1:12:3 | Baz |
+| classes.rb:2:1:56:3 | classes.rb | classes.rb:16:1:17:3 | MyClass |
+| classes.rb:2:1:56:3 | classes.rb | classes.rb:20:1:37:3 | Wibble |
+| classes.rb:2:1:56:3 | classes.rb | classes.rb:55:1:56:3 | MyClassInGlobalScope |
 | classes.rb:20:1:37:3 | Wibble | classes.rb:32:3:33:5 | ClassInWibble |
 | modules.rb:4:1:24:3 | Foo | modules.rb:19:3:20:5 | ClassInFoo |
 | modules.rb:5:3:14:5 | Bar | modules.rb:6:5:7:7 | ClassInFooBar |
@@ -38,5 +46,12 @@ moduleBaseMethods
 | modules.rb:37:1:46:3 | Bar | modules.rb:41:3:42:5 | method_b |
 | modules.rb:48:1:57:3 | Bar | modules.rb:52:3:53:5 | method_in_another_definition_of_foo_bar |
 moduleBaseModules
+| classes.rb:2:1:56:3 | classes.rb | classes.rb:15:1:15:20 | MyModule |
 | classes.rb:20:1:37:3 | Wibble | classes.rb:35:3:36:5 | ModuleInWibble |
+| modules.rb:1:1:61:3 | modules.rb | modules.rb:1:1:2:3 | Empty |
+| modules.rb:1:1:61:3 | modules.rb | modules.rb:4:1:24:3 | Foo |
+| modules.rb:1:1:61:3 | modules.rb | modules.rb:26:1:35:3 | Foo |
+| modules.rb:1:1:61:3 | modules.rb | modules.rb:37:1:46:3 | Bar |
+| modules.rb:1:1:61:3 | modules.rb | modules.rb:48:1:57:3 | Bar |
+| modules.rb:1:1:61:3 | modules.rb | modules.rb:60:1:61:3 | MyModuleInGlobalScope |
 | modules.rb:4:1:24:3 | Foo | modules.rb:5:3:14:5 | Bar |

ql/test/library-tests/controlflow/graph/Cfg.expected

@@ -1,5 +1,5 @@
 break_ensure.rb:
-#    1| enter AstNode
+#    1| enter break_ensure.rb
 #-----|  -> m1
 
 #    1| enter m1
@@ -15,20 +15,20 @@ break_ensure.rb:
 #-----|  -> elements
 
 case.rb:
-#    1| enter AstNode
+#    1| enter case.rb
 #-----|  -> if_in_case
 
 #    1| enter if_in_case
 #-----|  -> case ...
 
 cfg.rb:
-#    1| enter AstNode
+#    1| enter cfg.rb
 #-----|  -> bar
 
-#   15| enter AstNode
+#   15| enter BEGIN { ... }
 #-----|  -> puts
 
-#   19| enter AstNode
+#   19| enter END { ... }
 #-----|  -> puts
 
 #   25| enter { ... }
@@ -65,7 +65,7 @@ cfg.rb:
 #-----|  -> x
 
 exit.rb:
-#    1| enter AstNode
+#    1| enter exit.rb
 #-----|  -> m1
 
 #    1| enter m1
@@ -75,14 +75,14 @@ exit.rb:
 #-----|  -> x
 
 heredoc.rb:
-#    1| enter AstNode
+#    1| enter heredoc.rb
 #-----|  -> double_heredoc
 
 #    1| enter double_heredoc
 #-----|  -> puts
 
 ifs.rb:
-#    1| enter AstNode
+#    1| enter ifs.rb
 #-----|  -> m1
 
 #    1| enter m1
@@ -107,7 +107,7 @@ ifs.rb:
 #-----|  -> true
 
 loops.rb:
-#    1| enter AstNode
+#    1| enter loops.rb
 #-----|  -> m1
 
 #    1| enter m1
@@ -123,7 +123,7 @@ loops.rb:
 #-----|  -> x
 
 raise.rb:
-#    1| enter AstNode
+#    1| enter raise.rb
 #-----|  -> ExceptionA
 
 #    7| enter m1
@@ -422,7 +422,7 @@ break_ensure.rb:
 #-----|  -> call to puts
 
 #   44| m4
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit break_ensure.rb (normal)
 
 #   44| m4
 #-----|  -> m4
@@ -511,7 +511,7 @@ break_ensure.rb:
 
 case.rb:
 #    1| if_in_case
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit case.rb (normal)
 
 #    1| if_in_case
 #-----|  -> if_in_case
@@ -620,7 +620,7 @@ cfg.rb:
 #-----|  -> StringArray
 
 #   12| call to puts
-#-----|  -> BeginBlock
+#-----|  -> BEGIN { ... }
 
 #   12| puts
 #-----|  -> 4
@@ -628,11 +628,11 @@ cfg.rb:
 #   12| 4
 #-----|  -> call to puts
 
-#   15| BeginBlock
-#-----|  -> EndBlock
+#   15| BEGIN { ... }
+#-----|  -> END { ... }
 
 #   16| call to puts
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit BEGIN { ... } (normal)
 
 #   16| puts
 #-----|  -> hello
@@ -640,11 +640,11 @@ cfg.rb:
 #   16| hello
 #-----|  -> call to puts
 
-#   19| EndBlock
+#   19| END { ... }
 #-----|  -> 41
 
 #   20| call to puts
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit END { ... } (normal)
 
 #   20| puts
 #-----|  -> world
@@ -2039,7 +2039,7 @@ cfg.rb:
 #  188| 42
 
 #  191| call to run_block
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit cfg.rb (normal)
 
 #  191| run_block
 #-----|  -> { ... }
@@ -2101,7 +2101,7 @@ exit.rb:
 #-----|  -> call to puts
 
 #    8| m2
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit exit.rb (normal)
 
 #    8| m2
 #-----|  -> m2
@@ -2142,7 +2142,7 @@ exit.rb:
 
 heredoc.rb:
 #    1| double_heredoc
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit heredoc.rb (normal)
 
 #    1| double_heredoc
 #-----|  -> double_heredoc
@@ -2519,7 +2519,7 @@ ifs.rb:
 #-----|  -> ... == ...
 
 #   40| constant_condition
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit ifs.rb (normal)
 
 #   40| constant_condition
 #-----|  -> constant_condition
@@ -2681,7 +2681,7 @@ loops.rb:
 #-----|  -> call to puts
 
 #   24| m3
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit loops.rb (normal)
 
 #   24| m3
 #-----|  -> m3
@@ -3722,7 +3722,7 @@ raise.rb:
 #-----|  -> exit m13 (normal)
 
 #  154| m14
-#-----|  -> exit AstNode (normal)
+#-----|  -> exit raise.rb (normal)
 
 #  154| m14
 #-----|  -> m14
@@ -3768,7 +3768,7 @@ raise.rb:
 #-----|  -> call to nil?
 
 break_ensure.rb:
-#    1| exit AstNode
+#    1| exit break_ensure.rb
 
 #    1| exit m1
 
@@ -3779,16 +3779,16 @@ break_ensure.rb:
 #   44| exit m4
 
 case.rb:
-#    1| exit AstNode
+#    1| exit case.rb
 
 #    1| exit if_in_case
 
 cfg.rb:
-#    1| exit AstNode
+#    1| exit cfg.rb
 
-#   15| exit AstNode
+#   15| exit BEGIN { ... }
 
-#   19| exit AstNode
+#   19| exit END { ... }
 
 #   25| exit { ... }
 
@@ -3811,19 +3811,19 @@ cfg.rb:
 #  191| exit { ... }
 
 exit.rb:
-#    1| exit AstNode
+#    1| exit exit.rb
 
 #    1| exit m1
 
 #    8| exit m2
 
 heredoc.rb:
-#    1| exit AstNode
+#    1| exit heredoc.rb
 
 #    1| exit double_heredoc
 
 ifs.rb:
-#    1| exit AstNode
+#    1| exit ifs.rb
 
 #    1| exit m1
 
@@ -3840,7 +3840,7 @@ ifs.rb:
 #   40| exit constant_condition
 
 loops.rb:
-#    1| exit AstNode
+#    1| exit loops.rb
 
 #    1| exit m1
 
@@ -3851,7 +3851,7 @@ loops.rb:
 #   25| exit do ... end
 
 raise.rb:
-#    1| exit AstNode
+#    1| exit raise.rb
 
 #    7| exit m1
 
@@ -3884,8 +3884,8 @@ raise.rb:
 #  155| exit { ... }
 
 break_ensure.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit break_ensure.rb (normal)
+#-----|  -> exit break_ensure.rb
 
 #    1| exit m1 (normal)
 #-----|  -> exit m1
@@ -3900,21 +3900,21 @@ break_ensure.rb:
 #-----|  -> exit m4
 
 case.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit case.rb (normal)
+#-----|  -> exit case.rb
 
 #    1| exit if_in_case (normal)
 #-----|  -> exit if_in_case
 
 cfg.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit cfg.rb (normal)
+#-----|  -> exit cfg.rb
 
-#   15| exit AstNode (normal)
-#-----|  -> exit AstNode
+#   15| exit BEGIN { ... } (normal)
+#-----|  -> exit BEGIN { ... }
 
-#   19| exit AstNode (normal)
-#-----|  -> exit AstNode
+#   19| exit END { ... } (normal)
+#-----|  -> exit END { ... }
 
 #   25| exit { ... } (normal)
 #-----|  -> exit { ... }
@@ -3947,8 +3947,8 @@ cfg.rb:
 #-----|  -> exit { ... }
 
 exit.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit exit.rb (normal)
+#-----|  -> exit exit.rb
 
 #    1| exit m1 (abnormal)
 #-----|  -> exit m1
@@ -3963,15 +3963,15 @@ exit.rb:
 #-----|  -> exit m2
 
 heredoc.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit heredoc.rb (normal)
+#-----|  -> exit heredoc.rb
 
 #    1| exit double_heredoc (normal)
 #-----|  -> exit double_heredoc
 
 ifs.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit ifs.rb (normal)
+#-----|  -> exit ifs.rb
 
 #    1| exit m1 (normal)
 #-----|  -> exit m1
@@ -3995,8 +3995,8 @@ ifs.rb:
 #-----|  -> exit constant_condition
 
 loops.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit loops.rb (normal)
+#-----|  -> exit loops.rb
 
 #    1| exit m1 (normal)
 #-----|  -> exit m1
@@ -4011,8 +4011,8 @@ loops.rb:
 #-----|  -> exit do ... end
 
 raise.rb:
-#    1| exit AstNode (normal)
-#-----|  -> exit AstNode
+#    1| exit raise.rb (normal)
+#-----|  -> exit raise.rb
 
 #    7| exit m1 (abnormal)
 #-----|  -> exit m1