From 9185a93312e28090a9fee2cee74a89ccb4f30696 Mon Sep 17 00:00:00 2001
From: Nick Rolfe <nickrolfe@github.com>
Date: Thu, 6 May 2021 16:20:50 +0100
Subject: [PATCH] Make SingletonClassDeclarationTree post-order

---
 .../internal/ControlFlowGraphImpl.qll          | 18 ++++++++++++++++--
 .../controlflow/graph/Cfg.expected             |  6 +++---
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
index 104c9692296..c05083afb9d 100644
--- a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -1176,13 +1176,27 @@ module Trees {
     }
   }
 
-  private class SingletonClassTree extends BodyStmtPreOrderTree, SingletonClass {
+  private class SingletonClassTree extends BodyStmtPostOrderTree, SingletonClass {
+    final override predicate first(AstNode first) {
+      this.firstInner(first)
+      or
+      not exists(this.getAChild(_)) and
+      first = this
+    }
+
+    final override predicate succ(AstNode pred, AstNode succ, Completion c) {
+      BodyStmtPostOrderTree.super.succ(pred, succ, c)
+      or
+      succ = this and
+      this.lastInner(pred, c)
+    }
+
     /** Gets the `i`th child in the body of this block. */
     final override AstNode getBodyChild(int i, boolean rescuable) {
       (
         result = this.getValue() and i = 0 and rescuable = false
         or
-        result = BodyStmtPreOrderTree.super.getBodyChild(i - 1, rescuable)
+        result = BodyStmtPostOrderTree.super.getBodyChild(i - 1, rescuable)
       )
     }
   }
diff --git a/ql/test/library-tests/controlflow/graph/Cfg.expected b/ql/test/library-tests/controlflow/graph/Cfg.expected
index 865191c7259..d3731c5e367 100644
--- a/ql/test/library-tests/controlflow/graph/Cfg.expected
+++ b/ql/test/library-tests/controlflow/graph/Cfg.expected
@@ -1599,10 +1599,10 @@ cfg.rb:
 #-----|  -> call to itself
 
 #  141| Constant
-#-----|  -> class << ...
+#-----|  -> Silly
 
 #  143| class << ...
-#-----|  -> Silly
+#-----|  -> silly
 
 #  143| call to itself
 #-----|  -> setter=
@@ -1617,7 +1617,7 @@ cfg.rb:
 #-----|  -> self
 
 #  145| print
-#-----|  -> silly
+#-----|  -> class << ...
 
 #  145| exit print