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();