From 635b6fb45b825b300f829d26b2b14010ea0d573b Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Wed, 10 Feb 2021 13:28:16 +0100
Subject: [PATCH] AST: lambda and brace block bodies

---
 ql/src/codeql_ruby/ast/Method.qll          |  4 ++--
 ql/src/codeql_ruby/ast/internal/Method.qll | 13 +++++++++++--
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/ql/src/codeql_ruby/ast/Method.qll b/ql/src/codeql_ruby/ast/Method.qll
index b6390e71d06..05245497fa5 100644
--- a/ql/src/codeql_ruby/ast/Method.qll
+++ b/ql/src/codeql_ruby/ast/Method.qll
@@ -56,14 +56,14 @@ class SingletonMethod extends Callable, BodyStatement, @singleton_method {
  * -> (x) { x + 1 }
  * ```
  */
-class Lambda extends Callable, @lambda {
+class Lambda extends Callable, BodyStatement, @lambda {
   final override Lambda::Range range;
 
   final override string getAPrimaryQlClass() { result = "Lambda" }
 }
 
 /** A block. */
-class Block extends Callable {
+class Block extends Callable, ExprSequence {
   override Block::Range range;
 }
 
diff --git a/ql/src/codeql_ruby/ast/internal/Method.qll b/ql/src/codeql_ruby/ast/internal/Method.qll
index 834a959abbf..455a1e9c5c2 100644
--- a/ql/src/codeql_ruby/ast/internal/Method.qll
+++ b/ql/src/codeql_ruby/ast/internal/Method.qll
@@ -48,19 +48,26 @@ module SingletonMethod {
 }
 
 module Lambda {
-  class Range extends Callable::Range, @lambda {
+  class Range extends Callable::Range, BodyStatement::Range, @lambda {
     final override Generated::Lambda generated;
 
     final override Parameter::Range getParameter(int n) {
       result = generated.getParameters().getChild(n)
     }
 
+    final override Expr getExpr(int i) {
+      result = generated.getBody().(Generated::DoBlock).getChild(i) or
+      result = generated.getBody().(Generated::Block).getChild(i)
+    }
+
     final override string toString() { result = "-> { ... }" }
   }
 }
 
 module Block {
-  abstract class Range extends Callable::Range { }
+  abstract class Range extends Callable::Range, ExprSequence::Range {
+    Range() { not generated.getParent() instanceof Generated::Lambda }
+  }
 }
 
 module DoBlock {
@@ -85,6 +92,8 @@ module BraceBlock {
       result = generated.getParameters().getChild(n)
     }
 
+    final override Expr getExpr(int i) { result = generated.getChild(i) }
+
     final override string toString() { result = "{ ... }" }
   }
 }