From f2effce78631f7819b5995580322bcbe071b1faf Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Tue, 15 Dec 2020 19:01:29 +0100
Subject: [PATCH] CFG: improve handling of block and lambda

---
 .../codeql_ruby/controlflow/ControlFlowGraph.qll | 16 ++++++++++++++++
 .../internal/ControlFlowGraphImpl.qll            |  8 +++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll b/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll
index ac2eb82653b..2c8fd98b53b 100644
--- a/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll
+++ b/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll
@@ -73,10 +73,26 @@ private module CfgScope {
 
   private class BlockScope extends Range, Generated::Block {
     final override string getName() { result = "block" }
+
+    final override predicate entry(Generated::AstNode first) {
+      first(this.(Trees::BlockTree).getFirstChildNode(), first)
+    }
+
+    final override predicate exit(Generated::AstNode last, Completion c) {
+      last(this.(Trees::BlockTree).getLastChildNode(), last, c)
+    }
   }
 
   private class LambdaScope extends Range, Generated::Lambda {
     final override string getName() { result = "lambda" }
+
+    final override predicate entry(Generated::AstNode first) {
+      first(this.(Trees::LambdaTree).getFirstChildNode(), first)
+    }
+
+    final override predicate exit(Generated::AstNode last, Completion c) {
+      last(this.(Trees::LambdaTree).getLastChildNode(), last, c)
+    }
   }
 }
 
diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
index bf8df829838..1d3dd316e8c 100644
--- a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -295,7 +295,7 @@ module Trees {
     }
   }
 
-  private class BlockTree extends StandardPreOrderTree, Block {
+  class BlockTree extends StandardNode, PreOrderTree, PostOrderTree, Block {
     final override AstNode getChildNode(int i) {
       result = this.getParameters() and i = 0
       or
@@ -439,7 +439,9 @@ module Trees {
     override predicate isHidden() { any() }
   }
 
-  private class DoBlockTree extends RescueEnsureBlockTree, PreOrderTree, PostOrderTree, DoBlock {
+  private class DoBlockTree extends RescueEnsureBlockTree, PostOrderTree, DoBlock {
+    final override predicate first(AstNode first) { first = this }
+
     final override AstNode getChildNode(int i, boolean rescuable) {
       result = this.getParameters() and i = 0 and rescuable = false
       or
@@ -687,7 +689,7 @@ module Trees {
     final override AstNode getDefaultValue() { result = this.getValue() }
   }
 
-  private class LambdaTree extends StandardPreOrderTree, Lambda {
+  class LambdaTree extends StandardNode, PreOrderTree, PostOrderTree, Lambda {
     final override AstNode getChildNode(int i) {
       result = this.getParameters() and i = 0
       or