CFG: fix BEGIN and END blocks

aibaars · aibaars · commit 81c907a87ace · 2020-12-17T10:04:01.000+01:00
diff --git a/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll b/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll
@@ -23,10 +23,26 @@ private module CfgScope {
 
   private class BeginBlockScope extends Range, Generated::BeginBlock {
     final override string getName() { result = "BEGIN block" }
+
+    final override predicate entry(Generated::AstNode first) {
+      first(this.(Trees::BeginBlockTree).getFirstChildNode(), first)
+    }
+
+    final override predicate exit(Generated::AstNode last, Completion c) {
+      last(this.(Trees::BeginBlockTree).getLastChildNode(), last, c)
+    }
   }
 
   private class EndBlockScope extends Range, Generated::EndBlock {
     final override string getName() { result = "END block" }
+
+    final override predicate entry(Generated::AstNode first) {
+      first(this.(Trees::EndBlockTree).getFirstChildNode(), first)
+    }
+
+    final override predicate exit(Generated::AstNode last, Completion c) {
+      last(this.(Trees::EndBlockTree).getLastChildNode(), last, c)
+    }
   }
 
   private class MethodScope extends Range, Generated::AstNode {
diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -274,7 +274,7 @@ module Trees {
     override predicate isHidden() { any() }
   }
 
-  private class BeginBlockTree extends StandardPreOrderTree, BeginBlock {
+  class BeginBlockTree extends PreOrderTree, PostOrderTree, StandardNode, BeginBlock {
     final override AstNode getChildNode(int i) { result = this.getChild(i) }
   }
 
@@ -456,7 +456,7 @@ module Trees {
 
   private class EmptyStatementTree extends LeafTree, EmptyStatement { }
 
-  private class EndBlockTree extends StandardPreOrderTree, EndBlock {
+  class EndBlockTree extends StandardNode, PreOrderTree, PostOrderTree, EndBlock {
     final override AstNode getChildNode(int i) { result = this.getChild(i) }
   }
 
diff --git a/ql/test/library-tests/controlflow/graph/Cfg.expected b/ql/test/library-tests/controlflow/graph/Cfg.expected
@@ -26,10 +26,10 @@ cfg.rb:
 #-----|  -> bar
 
 #   15| enter BEGIN block
-#-----|  -> BeginBlock
+#-----|  -> String
 
 #   19| enter END block
-#-----|  -> EndBlock
+#-----|  -> String
 
 #   25| enter block
 #-----|  -> x
@@ -608,10 +608,9 @@ cfg.rb:
 #-----|  -> puts
 
 #   15| BeginBlock
-#-----|  -> String
+#-----|  -> EndBlock
 
 #   16| Call
-#-----|  -> EndBlock
 #-----|  -> exit BEGIN block (normal)
 
 #   16| puts
@@ -621,10 +620,9 @@ cfg.rb:
 #-----|  -> puts
 
 #   19| EndBlock
-#-----|  -> String
+#-----|  -> 41
 
 #   20| Call
-#-----|  -> 41
 #-----|  -> exit END block (normal)
 
 #   20| puts

Original file line number	Diff line number	Diff line change
`@@ -274,7 +274,7 @@ module Trees {`
`274`	`274`	`override predicate isHidden() { any() }`
`275`	`275`	`}`
`276`	`276`
`277`		`- private class BeginBlockTree extends StandardPreOrderTree, BeginBlock {`
	`277`	`+ class BeginBlockTree extends PreOrderTree, PostOrderTree, StandardNode, BeginBlock {`
`278`	`278`	`final override AstNode getChildNode(int i) { result = this.getChild(i) }`
`279`	`279`	`}`
`280`	`280`
`@@ -456,7 +456,7 @@ module Trees {`
`456`	`456`
`457`	`457`	`private class EmptyStatementTree extends LeafTree, EmptyStatement { }`
`458`	`458`
`459`		`- private class EndBlockTree extends StandardPreOrderTree, EndBlock {`
	`459`	`+ class EndBlockTree extends StandardNode, PreOrderTree, PostOrderTree, EndBlock {`
`460`	`460`	`final override AstNode getChildNode(int i) { result = this.getChild(i) }`
`461`	`461`	`}`
`462`	`462`