From 24adbb80c7adaa0ee084cd4dac18acd5dc3c4c0c Mon Sep 17 00:00:00 2001
From: Simon Friis Vindum <simonfv@gmail.com>
Date: Wed, 20 Nov 2024 11:50:46 +0100
Subject: [PATCH] Rust: Include self parameters in the CFG

---
 .../controlflow/internal/ControlFlowGraphImpl.qll | 15 +++++++++++++--
 .../test/library-tests/controlflow/Cfg.expected   |  8 ++++++--
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll b/rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll
index b331ddd2d540..ef4f7808fa20 100644
--- a/rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll
@@ -72,10 +72,17 @@ import CfgImpl
 class CallableScopeTree extends StandardTree, PreOrderTree, PostOrderTree, Scope::CallableScope {
   override predicate propagatesAbnormal(AstNode child) { none() }
 
+  private int getNumberOfSelfParams() {
+    if this.getParamList().hasSelfParam() then result = 1 else result = 0
+  }
+
   override AstNode getChildNode(int i) {
-    result = this.getParamList().getParam(i)
+    i = 0 and
+    result = this.getParamList().getSelfParam()
+    or
+    result = this.getParamList().getParam(i - this.getNumberOfSelfParams())
     or
-    i = this.getParamList().getNumberOfParams() and
+    i = this.getParamList().getNumberOfParams() + this.getNumberOfSelfParams() and
     result = this.getBody()
   }
 }
@@ -191,6 +198,10 @@ class NameTree extends LeafTree, Name { }
 
 class NameRefTree extends LeafTree, NameRef { }
 
+class SelfParamTree extends StandardPostOrderTree, SelfParam {
+  override AstNode getChildNode(int i) { i = 0 and result = this.getName() }
+}
+
 class TypeRefTree extends LeafTree instanceof TypeRef { }
 
 /**
diff --git a/rust/ql/test/library-tests/controlflow/Cfg.expected b/rust/ql/test/library-tests/controlflow/Cfg.expected
index 6b2642e65126..8fd71a451b86 100644
--- a/rust/ql/test/library-tests/controlflow/Cfg.expected
+++ b/rust/ql/test/library-tests/controlflow/Cfg.expected
@@ -1051,14 +1051,18 @@ edges
 | test.rs:502:28:504:5 | BlockExpr | test.rs:502:5:504:5 | exit new (normal) |  |
 | test.rs:503:9:503:25 | RecordExpr | test.rs:502:28:504:5 | BlockExpr |  |
 | test.rs:503:23:503:23 | a | test.rs:503:9:503:25 | RecordExpr |  |
-| test.rs:506:5:508:5 | enter negated | test.rs:507:23:507:26 | self |  |
+| test.rs:506:5:508:5 | enter negated | test.rs:506:16:506:19 | self |  |
 | test.rs:506:5:508:5 | exit negated (normal) | test.rs:506:5:508:5 | exit negated |  |
+| test.rs:506:16:506:19 | SelfParam | test.rs:507:23:507:26 | self |  |
+| test.rs:506:16:506:19 | self | test.rs:506:16:506:19 | SelfParam |  |
 | test.rs:506:30:508:5 | BlockExpr | test.rs:506:5:508:5 | exit negated (normal) |  |
 | test.rs:507:9:507:30 | RecordExpr | test.rs:506:30:508:5 | BlockExpr |  |
 | test.rs:507:23:507:26 | self | test.rs:507:23:507:28 | FieldExpr |  |
 | test.rs:507:23:507:28 | FieldExpr | test.rs:507:9:507:30 | RecordExpr |  |
-| test.rs:510:5:512:5 | enter multifly_add | test.rs:510:32:510:32 | a |  |
+| test.rs:510:5:512:5 | enter multifly_add | test.rs:510:26:510:29 | self |  |
 | test.rs:510:5:512:5 | exit multifly_add (normal) | test.rs:510:5:512:5 | exit multifly_add |  |
+| test.rs:510:21:510:29 | SelfParam | test.rs:510:32:510:32 | a |  |
+| test.rs:510:26:510:29 | self | test.rs:510:21:510:29 | SelfParam |  |
 | test.rs:510:32:510:32 | a | test.rs:510:32:510:37 | Param | match |
 | test.rs:510:32:510:37 | Param | test.rs:510:40:510:40 | b |  |
 | test.rs:510:40:510:40 | b | test.rs:510:40:510:45 | Param | match |