From 64cba18c41d7dcbf46b584e5e1882169e9b9dacb Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Fri, 12 Feb 2021 13:15:49 +0100
Subject: [PATCH 1/4] AST: add Self class

---
 ql/src/codeql_ruby/ast/Expr.qll                | 12 ++++++++++++
 ql/src/codeql_ruby/ast/internal/AST.qll        |  2 --
 ql/src/codeql_ruby/ast/internal/Expr.qll       |  8 ++++++++
 ql/test/library-tests/ast/calls/calls.expected |  2 +-
 ql/test/library-tests/ast/calls/calls.rb       |  4 ++--
 5 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/ql/src/codeql_ruby/ast/Expr.qll b/ql/src/codeql_ruby/ast/Expr.qll
index 3073501e477..e48f091dfd1 100644
--- a/ql/src/codeql_ruby/ast/Expr.qll
+++ b/ql/src/codeql_ruby/ast/Expr.qll
@@ -12,6 +12,18 @@ class Expr extends Stmt {
   Expr() { this = range }
 }
 
+/**
+ * A reference to the current object. For example:
+ * - `self == other`
+ * - `self.method_name`
+ * - `def self.method_name ... end`
+ */
+class Self extends Expr, @token_self {
+  override Self::Range range;
+
+  final override string getAPrimaryQlClass() { result = "Self" }
+}
+
 /**
  * A literal.
  *
diff --git a/ql/src/codeql_ruby/ast/internal/AST.qll b/ql/src/codeql_ruby/ast/internal/AST.qll
index d92810b8e2f..38df294bf98 100644
--- a/ql/src/codeql_ruby/ast/internal/AST.qll
+++ b/ql/src/codeql_ruby/ast/internal/AST.qll
@@ -47,8 +47,6 @@ module AstNode {
       or
       this instanceof Generated::BareString
       or
-      this instanceof Generated::Self
-      or
       this instanceof Generated::Float
       or
       this instanceof Generated::Superclass
diff --git a/ql/src/codeql_ruby/ast/internal/Expr.qll b/ql/src/codeql_ruby/ast/internal/Expr.qll
index b61d54835ee..7b848910580 100644
--- a/ql/src/codeql_ruby/ast/internal/Expr.qll
+++ b/ql/src/codeql_ruby/ast/internal/Expr.qll
@@ -6,6 +6,14 @@ module Expr {
   abstract class Range extends Stmt::Range { }
 }
 
+module Self {
+  class Range extends Expr::Range, @token_self {
+    final override Generated::Self generated;
+
+    final override string toString() { result = "self" }
+  }
+}
+
 module Literal {
   abstract class Range extends Expr::Range {
     abstract string getValueText();
diff --git a/ql/test/library-tests/ast/calls/calls.expected b/ql/test/library-tests/ast/calls/calls.expected
index 893e3fb4e99..813b19f1ba8 100644
--- a/ql/test/library-tests/ast/calls/calls.expected
+++ b/ql/test/library-tests/ast/calls/calls.expected
@@ -74,7 +74,6 @@ callsWithNoReceiverArgumentsOrBlock
 | calls.rb:286:5:286:9 | call to super | super |
 | calls.rb:287:5:287:11 | call to super | super |
 | calls.rb:303:5:303:7 | call to foo | foo |
-| calls.rb:304:5:304:14 | call to super | super |
 | calls.rb:305:5:305:9 | call to super | super |
 callsWithArguments
 | calls.rb:14:1:14:11 | call to foo | foo | 0 | calls.rb:14:5:14:5 | 0 |
@@ -174,6 +173,7 @@ callsWithReceiver
 | calls.rb:275:7:275:12 | call to bar | calls.rb:275:7:275:7 | X |
 | calls.rb:279:11:279:16 | call to bar | calls.rb:279:11:279:11 | X |
 | calls.rb:303:5:303:13 | call to super | calls.rb:303:5:303:7 | call to foo |
+| calls.rb:304:5:304:14 | call to super | calls.rb:304:5:304:8 | self |
 | calls.rb:305:5:305:15 | call to super | calls.rb:305:5:305:9 | call to super |
 callsWithBlock
 | calls.rb:17:1:17:17 | call to foo | calls.rb:17:5:17:17 | { ... } |
diff --git a/ql/test/library-tests/ast/calls/calls.rb b/ql/test/library-tests/ast/calls/calls.rb
index c38896ccf47..fa44b5fad0b 100644
--- a/ql/test/library-tests/ast/calls/calls.rb
+++ b/ql/test/library-tests/ast/calls/calls.rb
@@ -301,7 +301,7 @@ end
 class AnotherClass
   def another_method
     foo.super
-    self.super # TODO: this shows up as a call without a receiver, but that should be fixed once we handle `self` expressions
+    self.super
     super.super # we expect the receiver to be a SuperCall, while the outer call should not (it's just a regular Call)
   end
-end
\ No newline at end of file
+end

From 015b581f57af7a39949003231bbb6f882eecdeb8 Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Fri, 12 Feb 2021 14:29:34 +0100
Subject: [PATCH 2/4] AST: add redo, retry, empty-statement

---
 ql/src/codeql_ruby/ast/Statement.qll          | 33 +++++++++++++++++++
 ql/src/codeql_ruby/ast/internal/AST.qll       |  4 ---
 ql/src/codeql_ruby/ast/internal/Statement.qll | 24 ++++++++++++++
 .../controlflow/graph/Cfg.expected            |  8 ++---
 4 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/ql/src/codeql_ruby/ast/Statement.qll b/ql/src/codeql_ruby/ast/Statement.qll
index d9c284a416d..34e9481acd1 100644
--- a/ql/src/codeql_ruby/ast/Statement.qll
+++ b/ql/src/codeql_ruby/ast/Statement.qll
@@ -25,6 +25,15 @@ class Stmt extends AstNode {
   Callable getEnclosingCallable() { result = this.getCfgScope() }
 }
 
+/**
+ * An empty statement (`;`).
+ */
+class EmptyStmt extends Stmt, @token_empty_statement {
+  final override EmptyStmt::Range range;
+
+  final override string getAPrimaryQlClass() { result = "EmptyStmt" }
+}
+
 /**
  * A statement that may return a value: `return`, `break` and `next`.
  *
@@ -82,3 +91,27 @@ class NextStmt extends ReturningStmt, @next {
 
   final override string getAPrimaryQlClass() { result = "NextStmt" }
 }
+
+/**
+ * A `redo` statement.
+ * ```rb
+ * redo
+ * ```
+ */
+class RedoStmt extends Stmt, @redo {
+  final override RedoStmt::Range range;
+
+  final override string getAPrimaryQlClass() { result = "RedoStmt" }
+}
+
+/**
+ * A `retry` statement.
+ * ```rb
+ * retry
+ * ```
+ */
+class RetryStmt extends Stmt, @retry {
+  final override RetryStmt::Range range;
+
+  final override string getAPrimaryQlClass() { result = "RetryStmt" }
+}
diff --git a/ql/src/codeql_ruby/ast/internal/AST.qll b/ql/src/codeql_ruby/ast/internal/AST.qll
index 38df294bf98..d39febed38d 100644
--- a/ql/src/codeql_ruby/ast/internal/AST.qll
+++ b/ql/src/codeql_ruby/ast/internal/AST.qll
@@ -51,10 +51,6 @@ module AstNode {
       or
       this instanceof Generated::Superclass
       or
-      this instanceof Generated::EmptyStatement
-      or
-      this instanceof Generated::Redo
-      or
       this instanceof Generated::Hash
       or
       this instanceof Generated::Array
diff --git a/ql/src/codeql_ruby/ast/internal/Statement.qll b/ql/src/codeql_ruby/ast/internal/Statement.qll
index b4fa88ce159..d9a5c9b294b 100644
--- a/ql/src/codeql_ruby/ast/internal/Statement.qll
+++ b/ql/src/codeql_ruby/ast/internal/Statement.qll
@@ -6,6 +6,14 @@ module Stmt {
   abstract class Range extends AstNode::Range { }
 }
 
+module EmptyStmt {
+  class Range extends Stmt::Range, @token_empty_statement {
+    final override Generated::EmptyStatement generated;
+
+    final override string toString() { result = ";" }
+  }
+}
+
 module ReturningStmt {
   abstract class Range extends Stmt::Range {
     abstract Generated::ArgumentList getArgumentList();
@@ -51,3 +59,19 @@ module NextStmt {
     final override Generated::ArgumentList getArgumentList() { result = generated.getChild() }
   }
 }
+
+module RedoStmt {
+  class Range extends Stmt::Range, @redo {
+    final override Generated::Redo generated;
+
+    final override string toString() { result = "redo" }
+  }
+}
+
+module RetryStmt {
+  class Range extends Stmt::Range, @retry {
+    final override Generated::Retry generated;
+
+    final override string toString() { result = "retry" }
+  }
+}
diff --git a/ql/test/library-tests/controlflow/graph/Cfg.expected b/ql/test/library-tests/controlflow/graph/Cfg.expected
index ced2fa6d3cb..610ff396d78 100644
--- a/ql/test/library-tests/controlflow/graph/Cfg.expected
+++ b/ql/test/library-tests/controlflow/graph/Cfg.expected
@@ -1973,7 +1973,7 @@ cfg.rb:
 
 #  181| ... == ...
 #-----| false -> if ...
-#-----| true -> Redo
+#-----| true -> redo
 
 #  181| x
 #-----|  -> 5
@@ -1981,7 +1981,7 @@ cfg.rb:
 #  181| 5
 #-----|  -> ... == ...
 
-#  181| Redo
+#  181| redo
 #-----| redo -> x
 
 #  182| call to puts
@@ -2651,7 +2651,7 @@ loops.rb:
 
 #   16| ... > ...
 #-----| false -> elsif ...
-#-----| true -> Redo
+#-----| true -> redo
 
 #   16| x
 #-----|  -> 10
@@ -2659,7 +2659,7 @@ loops.rb:
 #   16| 10
 #-----|  -> ... > ...
 
-#   17| Redo
+#   17| redo
 #-----| redo -> puts
 
 #   19| call to puts

From 874ac121d9997d409ab990c2d8c49f4e224e76e1 Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Fri, 12 Feb 2021 15:13:43 +0100
Subject: [PATCH 3/4] AST: Toplevel and BEGIN/ END blocks

---
 ql/src/codeql_ruby/ast/Module.qll             |  26 +++++
 ql/src/codeql_ruby/ast/Statement.qll          |  13 +++
 ql/src/codeql_ruby/ast/internal/AST.qll       |   6 -
 ql/src/codeql_ruby/ast/internal/Expr.qll      |  10 ++
 ql/src/codeql_ruby/ast/internal/Module.qll    |  19 +++
 ql/src/codeql_ruby/ast/internal/Statement.qll |  11 ++
 .../ast/modules/module_base.expected          |  15 +++
 .../controlflow/graph/Cfg.expected            | 108 +++++++++---------
 8 files changed, 148 insertions(+), 60 deletions(-)

diff --git a/ql/src/codeql_ruby/ast/Module.qll b/ql/src/codeql_ruby/ast/Module.qll
index bb267c8118a..53ef95bf3f2 100644
--- a/ql/src/codeql_ruby/ast/Module.qll
+++ b/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.
  *
diff --git a/ql/src/codeql_ruby/ast/Statement.qll b/ql/src/codeql_ruby/ast/Statement.qll
index 34e9481acd1..b5b2bbdbfef 100644
--- a/ql/src/codeql_ruby/ast/Statement.qll
+++ b/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`.
  *
diff --git a/ql/src/codeql_ruby/ast/internal/AST.qll b/ql/src/codeql_ruby/ast/internal/AST.qll
index d39febed38d..618aea0313a 100644
--- a/ql/src/codeql_ruby/ast/internal/AST.qll
+++ b/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
diff --git a/ql/src/codeql_ruby/ast/internal/Expr.qll b/ql/src/codeql_ruby/ast/internal/Expr.qll
index 7b848910580..9a7dfe22758 100644
--- a/ql/src/codeql_ruby/ast/internal/Expr.qll
+++ b/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;
diff --git a/ql/src/codeql_ruby/ast/internal/Module.qll b/ql/src/codeql_ruby/ast/internal/Module.qll
index b0b2ae67022..3897435bf4b 100644
--- a/ql/src/codeql_ruby/ast/internal/Module.qll
+++ b/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;
diff --git a/ql/src/codeql_ruby/ast/internal/Statement.qll b/ql/src/codeql_ruby/ast/internal/Statement.qll
index d9a5c9b294b..5d94ddc0684 100644
--- a/ql/src/codeql_ruby/ast/internal/Statement.qll
+++ b/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();
diff --git a/ql/test/library-tests/ast/modules/module_base.expected b/ql/test/library-tests/ast/modules/module_base.expected
index 11a2769c39c..6a790d4cfcd 100644
--- a/ql/test/library-tests/ast/modules/module_base.expected
+++ b/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 |
diff --git a/ql/test/library-tests/controlflow/graph/Cfg.expected b/ql/test/library-tests/controlflow/graph/Cfg.expected
index 610ff396d78..98762e036bd 100644
--- a/ql/test/library-tests/controlflow/graph/Cfg.expected
+++ b/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

From c0c155361f9ddfaa30ac3ef0f8ce72cf502edc67 Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Fri, 12 Feb 2021 17:50:00 +0100
Subject: [PATCH 4/4] Address comments

---
 ql/src/codeql_ruby/ast/Module.qll                  |  8 ++++----
 ql/src/codeql_ruby/ast/Statement.qll               | 14 +++++++++++++-
 .../library-tests/ast/modules/module_base.expected |  1 +
 .../library-tests/ast/modules/toplevel.expected    |  8 ++++++++
 ql/test/library-tests/ast/modules/toplevel.ql      |  9 +++++++++
 ql/test/library-tests/ast/modules/toplevel.rb      |  5 +++++
 6 files changed, 40 insertions(+), 5 deletions(-)
 create mode 100644 ql/test/library-tests/ast/modules/toplevel.expected
 create mode 100644 ql/test/library-tests/ast/modules/toplevel.ql
 create mode 100644 ql/test/library-tests/ast/modules/toplevel.rb

diff --git a/ql/src/codeql_ruby/ast/Module.qll b/ql/src/codeql_ruby/ast/Module.qll
index 53ef95bf3f2..bc5b105331b 100644
--- a/ql/src/codeql_ruby/ast/Module.qll
+++ b/ql/src/codeql_ruby/ast/Module.qll
@@ -43,14 +43,14 @@ class Toplevel extends ModuleBase, @program {
   final override string getAPrimaryQlClass() { result = "Toplevel" }
 
   /**
-   * Get the `n`th `BEGIN` block.
+   * Gets the `n`th `BEGIN` block.
    */
-  final StmtSequence getBeginBlock(int n) { result = range.getBeginBlock(n) }
+  final BeginBlock getBeginBlock(int n) { result = range.getBeginBlock(n) }
 
   /**
-   * Get a `BEGIN` block.
+   * Gets a `BEGIN` block.
    */
-  final StmtSequence getABeginBlock() { result = getBeginBlock(_) }
+  final BeginBlock getABeginBlock() { result = getBeginBlock(_) }
 }
 
 /**
diff --git a/ql/src/codeql_ruby/ast/Statement.qll b/ql/src/codeql_ruby/ast/Statement.qll
index b5b2bbdbfef..e7af8a952a0 100644
--- a/ql/src/codeql_ruby/ast/Statement.qll
+++ b/ql/src/codeql_ruby/ast/Statement.qll
@@ -35,10 +35,22 @@ class EmptyStmt extends Stmt, @token_empty_statement {
   final override string getAPrimaryQlClass() { result = "EmptyStmt" }
 }
 
+/**
+ * An `BEGIN` block.
+ * ```rb
+ * BEGIN { puts "starting ..." }
+ * ```
+ */
+class BeginBlock extends StmtSequence, @begin_block {
+  final override BeginBlock::Range range;
+
+  final override string getAPrimaryQlClass() { result = "BeginBlock" }
+}
+
 /**
  * An `END` block.
  * ```rb
- * END{ puts "shutting down" }
+ * END { puts "shutting down" }
  * ```
  */
 class EndBlock extends StmtSequence, @end_block {
diff --git a/ql/test/library-tests/ast/modules/module_base.expected b/ql/test/library-tests/ast/modules/module_base.expected
index 6a790d4cfcd..1e26e369d5f 100644
--- a/ql/test/library-tests/ast/modules/module_base.expected
+++ b/ql/test/library-tests/ast/modules/module_base.expected
@@ -22,6 +22,7 @@ moduleBases
 | modules.rb:48:1:57:3 | Bar | Module |
 | modules.rb:49:3:50:5 | ClassInAnotherDefinitionOfFooBar | Class |
 | modules.rb:60:1:61:3 | MyModuleInGlobalScope | Module |
+| toplevel.rb:1:1:5:23 | toplevel.rb | Toplevel |
 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 |
diff --git a/ql/test/library-tests/ast/modules/toplevel.expected b/ql/test/library-tests/ast/modules/toplevel.expected
new file mode 100644
index 00000000000..0eca47938a3
--- /dev/null
+++ b/ql/test/library-tests/ast/modules/toplevel.expected
@@ -0,0 +1,8 @@
+toplevel
+| classes.rb:2:1:56:3 | classes.rb | Toplevel |
+| modules.rb:1:1:61:3 | modules.rb | Toplevel |
+| toplevel.rb:1:1:5:23 | toplevel.rb | Toplevel |
+beginBlocks
+| toplevel.rb:1:1:5:23 | toplevel.rb | 1 | toplevel.rb:5:1:5:22 | BEGIN { ... } |
+endBlocks
+| toplevel.rb:1:1:5:23 | toplevel.rb | toplevel.rb:3:1:3:18 | END { ... } |
diff --git a/ql/test/library-tests/ast/modules/toplevel.ql b/ql/test/library-tests/ast/modules/toplevel.ql
new file mode 100644
index 00000000000..90afb185ff6
--- /dev/null
+++ b/ql/test/library-tests/ast/modules/toplevel.ql
@@ -0,0 +1,9 @@
+import ruby
+
+query predicate toplevel(Toplevel m, string pClass) { pClass = m.getAPrimaryQlClass() }
+
+query predicate beginBlocks(Toplevel m, int i, BeginBlock b) { b = m.getBeginBlock(i) }
+
+query predicate endBlocks(Toplevel m, EndBlock b) {
+  b.getLocation().getFile() = m.getLocation().getFile()
+}
diff --git a/ql/test/library-tests/ast/modules/toplevel.rb b/ql/test/library-tests/ast/modules/toplevel.rb
new file mode 100644
index 00000000000..46105392531
--- /dev/null
+++ b/ql/test/library-tests/ast/modules/toplevel.rb
@@ -0,0 +1,5 @@
+puts "world"
+
+END { puts "!!!" }
+
+BEGIN { puts "hello" }