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] 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" }