Rust: Handle chained let expressions

Author: hvitved

rust/ql/lib/codeql/rust/controlflow/internal/CfgNodes.qll

@@ -7,8 +7,7 @@ private import codeql.rust.controlflow.CfgNodes
 private import codeql.rust.internal.CachedStages
 
 private predicate isPostOrder(AstNode n) {
-  n instanceof Expr and
-  not n instanceof LetExpr
+  n instanceof Expr
   or
   n instanceof OrPat
   or

rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll

@@ -200,8 +200,7 @@ class TypeReprTree extends LeafTree instanceof TypeRepr { }
 /**
  * Provides `ControlFlowTree`s for expressions.
  *
- * Since expressions construct values, they are modeled in post-order, except for
- * `LetExpr`s.
+ * Since expressions construct values, they are modeled in post-order.
  */
 module ExprTrees {
   class ArrayExprTree extends StandardPostOrderTree, ArrayExpr {
@@ -341,21 +340,15 @@ module ExprTrees {
       child = [super.getCondition(), super.getABranch()]
     }
 
-    private ConditionalCompletion conditionCompletion(Completion c) {
-      if super.getCondition() instanceof LetExpr
-      then result = c.(MatchCompletion)
-      else result = c.(BooleanCompletion)
-    }
-
     override predicate succ(AstNode pred, AstNode succ, Completion c) {
       // Edges from the condition to the branches
       last(super.getCondition(), pred, c) and
       (
-        first(super.getThen(), succ) and this.conditionCompletion(c).succeeded()
+        first(super.getThen(), succ) and c.(ConditionalCompletion).succeeded()
         or
-        first(super.getElse(), succ) and this.conditionCompletion(c).failed()
+        first(super.getElse(), succ) and c.(ConditionalCompletion).failed()
         or
-        not super.hasElse() and succ = this and this.conditionCompletion(c).failed()
+        not super.hasElse() and succ = this and c.(ConditionalCompletion).failed()
       )
       or
       // An edge from the then branch to the last node
@@ -401,10 +394,7 @@ module ExprTrees {
     }
   }
 
-  // `LetExpr` is a pre-order tree such that the pattern itself ends up
-  // dominating successors in the graph in the same way that patterns do in
-  // `match` expressions.
-  class LetExprTree extends StandardPreOrderTree, LetExpr {
+  class LetExprTree extends StandardPostOrderTree, LetExpr {
     override AstNode getChildNode(int i) {
       i = 0 and
       result = this.getScrutinee()
@@ -456,21 +446,15 @@ module ExprTrees {
 
     override predicate first(AstNode node) { first(super.getCondition(), node) }
 
-    private ConditionalCompletion conditionCompletion(Completion c) {
-      if super.getCondition() instanceof LetExpr
-      then result = c.(MatchCompletion)
-      else result = c.(BooleanCompletion)
-    }
-
     override predicate succ(AstNode pred, AstNode succ, Completion c) {
       super.succ(pred, succ, c)
       or
       last(super.getCondition(), pred, c) and
-      this.conditionCompletion(c).succeeded() and
+      c.(ConditionalCompletion).succeeded() and
       first(this.getLoopBody(), succ)
       or
       last(super.getCondition(), pred, c) and
-      this.conditionCompletion(c).failed() and
+      c.(ConditionalCompletion).failed() and
       succ = this
     }
   }

rust/ql/lib/codeql/rust/controlflow/internal/Splitting.qll

@@ -71,13 +71,7 @@ module ConditionalCompletionSplitting {
       child = parent.(LogicalNotExpr).getExpr() and
       childCompletion.getDual() = parentCompletion
       or
-      (
-        childCompletion = parentCompletion
-        or
-        // needed for `let` expressions
-        childCompletion.(MatchCompletion).getValue() =
-          parentCompletion.(BooleanCompletion).getValue()
-      ) and
+      childCompletion = parentCompletion and
       (
         child = parent.(BinaryLogicalOperation).getAnOperand()
         or
@@ -92,6 +86,9 @@ module ConditionalCompletionSplitting {
         or
         child = parent.(PatternTrees::PostOrderPatTree).getPat(_)
       )
+      or
+      child = parent.(LetExpr).getPat() and
+      childCompletion.(MatchCompletion).getValue() = parentCompletion.(BooleanCompletion).getValue()
     }
   }
 

rust/ql/lib/codeql/rust/dataflow/Ssa.qll

@@ -194,9 +194,16 @@ module Ssa {
         ae.getRhs() = value
       )
       or
-      exists(LetStmtCfgNode ls |
-        ls.getPat().(IdentPatCfgNode).getName() = write and
-        ls.getInitializer() = value
+      exists(IdentPatCfgNode pat | pat.getName() = write |
+        exists(LetStmtCfgNode ls |
+          pat = ls.getPat() and
+          ls.getInitializer() = value
+        )
+        or
+        exists(LetExprCfgNode le |
+          pat = le.getPat() and
+          le.getScrutinee() = value
+        )
       )
     }
 

rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll

@@ -241,6 +241,12 @@ module LocalFlow {
       nodeTo.getCfgNode() = s.getPat()
     )
     or
+    // An edge from the right-hand side of a let expression to the left-hand side.
+    exists(LetExprCfgNode e |
+      nodeFrom.getCfgNode() = e.getScrutinee() and
+      nodeTo.getCfgNode() = e.getPat()
+    )
+    or
     exists(IdentPatCfgNode p |
       not p.isRef() and
       nodeFrom.getCfgNode() = p and
@@ -379,6 +385,8 @@ module RustDataFlow implements InputSig<Location> {
   predicate neverSkipInPathGraph(Node node) {
     node.(Node::Node).getCfgNode() = any(LetStmtCfgNode s).getPat()
     or
+    node.(Node::Node).getCfgNode() = any(LetExprCfgNode e).getPat()
+    or
     node.(Node::Node).getCfgNode() = any(AssignmentExprCfgNode a).getLhs()
     or
     exists(MatchExprCfgNode match |
@@ -899,6 +907,12 @@ module VariableCapture {
           v.getPat() = ls.getPat().getPat() and
           ls.getInitializer() = source
         )
+        or
+        exists(LetExprCfgNode le |
+          this = le and
+          v.getPat() = le.getPat().getPat() and
+          le.getScrutinee() = source
+        )
       }
 
       CapturedVariable getVariable() { result = v }

rust/ql/lib/codeql/rust/elements/internal/VariableImpl.qll

@@ -1,6 +1,6 @@
 private import rust
 private import codeql.rust.controlflow.ControlFlowGraph
-private import codeql.rust.elements.internal.generated.ParentChild
+private import codeql.rust.elements.internal.generated.ParentChild as ParentChild
 private import codeql.rust.elements.internal.PathImpl::Impl as PathImpl
 private import codeql.rust.elements.internal.PathExprBaseImpl::Impl as PathExprBaseImpl
 private import codeql.rust.elements.internal.FormatTemplateVariableAccessImpl::Impl as FormatTemplateVariableAccessImpl
@@ -36,6 +36,38 @@ module Impl {
     ClosureBodyScope() { this = any(ClosureExpr ce).getBody() }
   }
 
+  /**
+   * A scope for conditions, which may introduce variables using `let` expressions.
+   *
+   * Such variables are only available in the body guarded by the condition.
+   */
+  class ConditionScope extends VariableScope, Expr {
+    private AstNode parent;
+    private AstNode body;
+
+    ConditionScope() {
+      parent =
+        any(IfExpr ie |
+          this = ie.getCondition() and
+          body = ie.getThen()
+        )
+      or
+      parent =
+        any(WhileExpr we |
+          this = we.getCondition() and
+          body = we.getLoopBody()
+        )
+    }
+
+    /** Gets the parent of this condition. */
+    AstNode getParent() { result = parent }
+
+    /**
+     * Gets the body in which variables introduced in this scope are available.
+     */
+    AstNode getBody() { result = body }
+  }
+
   private Pat getAPatAncestor(Pat p) {
     (p instanceof IdentPat or p instanceof OrPat) and
     exists(Pat p0 | result = p0.getParentPat() |
@@ -152,8 +184,14 @@ module Impl {
     /** Gets the `let` statement that introduces this variable, if any. */
     LetStmt getLetStmt() { this.getPat() = result.getPat() }
 
+    /** Gets the `let` expression that introduces this variable, if any. */
+    LetExpr getLetExpr() { this.getPat() = result.getPat() }
+
     /** Gets the initial value of this variable, if any. */
-    Expr getInitializer() { result = this.getLetStmt().getInitializer() }
+    Expr getInitializer() {
+      result = this.getLetStmt().getInitializer() or
+      result = this.getLetExpr().getScrutinee()
+    }
 
     /** Holds if this variable is captured. */
     predicate isCaptured() { this.getAnAccess().isCapture() }
@@ -193,15 +231,53 @@ module Impl {
     string getName() { result = name_ }
   }
 
+  pragma[nomagic]
+  private Element getImmediateChildAdj(Element e, int preOrd, int index) {
+    result = ParentChild::getImmediateChild(e, index) and
+    preOrd = 0 and
+    not exists(ConditionScope cs |
+      e = cs.getParent() and
+      result = cs.getBody()
+    )
+    or
+    result = e.(ConditionScope).getBody() and
+    preOrd = 1 and
+    index = 0
+  }
+
+  /**
+   * An adjusted version of `ParentChild::getImmediateChild`, which makes the following
+   * two adjustments:
+   *
+   * 1. For conditions like `if cond body`, instead of letting `body` be the second child
+   *    of `if`, we make it the last child of `cond`. This ensures that variables
+   *    introduced in the `cond` scope are available in `body`.
+   *
+   * 2. A similar adjustment is made for `while` loops: the body of the loop is made a
+   *    child of the loop condition instead of the loop itself.
+   */
+  pragma[nomagic]
+  private Element getImmediateChildAdj(Element e, int index) {
+    result =
+      rank[index + 1](Element res, int preOrd, int i |
+        res = getImmediateChildAdj(e, preOrd, i)
+      |
+        res order by preOrd, i
+      )
+  }
+
+  private Element getImmediateParentAdj(Element e) { e = getImmediateChildAdj(result, _) }
+
   private AstNode getAnAncestorInVariableScope(AstNode n) {
     (
       n instanceof Pat or
       n instanceof VariableAccessCand or
       n instanceof LetStmt or
+      n = any(LetExpr le).getScrutinee() or
       n instanceof VariableScope
     ) and
     exists(AstNode n0 |
-      result = getImmediateParent(n0) or
+      result = getImmediateParentAdj(n0) or
       result = n0.(FormatTemplateVariableAccess).getArgument().getParent().getParent()
     |
       n0 = n
@@ -243,31 +319,32 @@ module Impl {
       this instanceof VariableScope or
       this instanceof VariableAccessCand or
       this instanceof LetStmt or
-      getImmediateChild(this, _) instanceof RelevantElement
+      this = any(LetExpr le).getScrutinee() or
+      getImmediateChildAdj(this, _) instanceof RelevantElement
     }
 
     pragma[nomagic]
-    private RelevantElement getChild(int index) { result = getImmediateChild(this, index) }
+    private RelevantElement getChild(int index) { result = getImmediateChildAdj(this, index) }
 
     pragma[nomagic]
-    private RelevantElement getImmediateChildMin(int index) {
+    private RelevantElement getImmediateChildAdjMin(int index) {
       // A child may have multiple positions for different accessors,
       // so always use the first
       result = this.getChild(index) and
       index = min(int i | result = this.getChild(i) | i)
     }
 
     pragma[nomagic]
-    RelevantElement getImmediateChild(int index) {
+    RelevantElement getImmediateChildAdj(int index) {
       result =
-        rank[index + 1](Element res, int i | res = this.getImmediateChildMin(i) | res order by i)
+        rank[index + 1](Element res, int i | res = this.getImmediateChildAdjMin(i) | res order by i)
     }
 
     pragma[nomagic]
     RelevantElement getImmediateLastChild() {
       exists(int last |
-        result = this.getImmediateChild(last) and
-        not exists(this.getImmediateChild(last + 1))
+        result = this.getImmediateChildAdj(last) and
+        not exists(this.getImmediateChildAdj(last + 1))
       )
     }
   }
@@ -288,13 +365,13 @@ module Impl {
     |
       // first child of a previously numbered node
       result = getPreOrderNumbering(scope, parent) + 1 and
-      n = parent.getImmediateChild(0)
+      n = parent.getImmediateChildAdj(0)
       or
       // non-first child of a previously numbered node
       exists(RelevantElement child, int i |
         result = getLastPreOrderNumbering(scope, child) + 1 and
-        child = parent.getImmediateChild(i) and
-        n = parent.getImmediateChild(i + 1)
+        child = parent.getImmediateChildAdj(i) and
+        n = parent.getImmediateChildAdj(i + 1)
       )
     )
   }
@@ -309,7 +386,7 @@ module Impl {
       result = getPreOrderNumbering(scope, leaf) and
       leaf != scope and
       (
-        not exists(leaf.getImmediateChild(_))
+        not exists(leaf.getImmediateChildAdj(_))
         or
         leaf instanceof VariableScope
       )
@@ -331,7 +408,7 @@ module Impl {
   /**
    * Holds if `v` is named `name` and is declared inside variable scope
    * `scope`. The pre-order numbering of the binding site of `v`, amongst
-   * all nodes nester under `scope`, is `ord`.
+   * all nodes nested under `scope`, is `ord`.
    */
   private predicate variableDeclInScope(Variable v, VariableScope scope, string name, int ord) {
     name = v.getText() and
@@ -354,25 +431,20 @@ module Impl {
           ord = getLastPreOrderNumbering(scope, let) + 1
         )
         or
-        exists(IfExpr ie, LetExpr let |
+        exists(LetExpr let, Expr scrutinee |
           let.getPat() = pat and
-          ie.getCondition() = let and
-          scope = ie.getThen() and
-          ord = getPreOrderNumbering(scope, scope)
+          scrutinee = let.getScrutinee() and
+          scope = getEnclosingScope(scrutinee) and
+          // for `let` expressions, variables are bound _after_ the expression, i.e.
+          // not in the RHS
+          ord = getLastPreOrderNumbering(scope, scrutinee) + 1
         )
         or
         exists(ForExpr fe |
           fe.getPat() = pat and
           scope = fe.getLoopBody() and
           ord = getPreOrderNumbering(scope, scope)
         )
-        or
-        exists(WhileExpr we, LetExpr let |
-          let.getPat() = pat and
-          we.getCondition() = let and
-          scope = we.getLoopBody() and
-          ord = getPreOrderNumbering(scope, scope)
-        )
       )
     )
   }
@@ -612,7 +684,7 @@ module Impl {
     or
     exists(Expr mid |
       assignmentExprDescendant(mid) and
-      getImmediateParent(e) = mid and
+      getImmediateParentAdj(e) = mid and
       not mid instanceof DerefExpr and
       not mid instanceof FieldExpr and
       not mid instanceof IndexExpr