Rust: Add some manual classes to CfgNodes.qll

Author: hvitved

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

@@ -9,54 +9,185 @@ private import internal.ControlFlowGraphImpl as CfgImpl
 private import internal.CfgNodes
 import Nodes
 
-/** A CFG node that corresponds to an element in the AST. */
-class AstCfgNode extends CfgNode {
-  AstNode node;
+class AstCfgNode = CfgImpl::AstCfgNode;
 
-  AstCfgNode() { node = this.getAstNode() }
+class ExitCfgNode = CfgImpl::ExitNode;
 
-  /** Gets the underlying ast node. */
-  AstNode getAstNode() { result = node }
-}
+/**
+ * An assignment expression, for example
+ *
+ * ```rust
+ * x = y;
+ * ```
+ */
+final class AssignmentExprCfgNode extends BinaryExprCfgNode {
+  AssignmentExpr a;
 
-/** A CFG node that corresponds to a parameter in the AST. */
-class ParamCfgNode extends AstCfgNode {
-  override Param node;
+  AssignmentExprCfgNode() { a = this.getBinaryExpr() }
 
-  /** Gets the underlying parameter. */
-  Param getParam() { result = node }
+  /** Gets the underlying `AssignmentExpr`. */
+  AssignmentExpr getAssignmentExpr() { result = a }
 }
 
-/** A CFG node that corresponds to a parameter in the AST. */
-class PatCfgNode extends AstCfgNode {
-  override Pat node;
+/**
+ * A match expression. For example:
+ * ```rust
+ * match x {
+ *     Option::Some(y) => y,
+ *     Option::None => 0,
+ * }
+ * ```
+ * ```rust
+ * match x {
+ *     Some(y) if y != 0 => 1 / y,
+ *     _ => 0,
+ * }
+ * ```
+ */
+final class MatchExprCfgNode extends Nodes::MatchExprCfgNode {
+  private MatchExpr node;
 
-  /** Gets the underlying pattern. */
-  Pat getPat() { result = node }
-}
+  MatchExprCfgNode() { node = this.getMatchExpr() }
+
+  /**
+   * Gets the pattern of the `i`th match arm, if it exists.
+   */
+  PatCfgNode getArmPat(int i) {
+    any(ChildMapping mapping).hasCfgChild(node, node.getArm(i).getPat(), this, result)
+  }
 
-/** A CFG node that corresponds to an expression in the AST. */
-class ExprCfgNode extends AstCfgNode {
-  override Expr node;
+  /**
+   * Gets the guard of the `i`th match arm, if it exists.
+   */
+  ExprCfgNode getArmGuard(int i) {
+    any(ChildMapping mapping)
+        .hasCfgChild(node, node.getArm(i).getGuard().getCondition(), this, result)
+  }
 
-  /** Gets the underlying expression. */
-  Expr getExpr() { result = node }
+  /**
+   * Gets the expression of the `i`th match arm, if it exists.
+   */
+  ExprCfgNode getArmExpr(int i) {
+    any(ChildMapping mapping).hasCfgChild(node, node.getArm(i).getExpr(), this, result)
+  }
 }
 
-/** A CFG node that corresponds to a call in the AST. */
-class CallExprCfgNode extends ExprCfgNode {
-  override CallExpr node;
+/**
+ * A block expression. For example:
+ * ```rust
+ * {
+ *     let x = 42;
+ * }
+ * ```
+ * ```rust
+ * 'label: {
+ *     let x = 42;
+ *     x
+ * }
+ * ```
+ */
+final class BlockExprCfgNode extends Nodes::BlockExprCfgNode {
+  private BlockExprChildMapping node;
 
-  /** Gets the underlying `CallExpr`. */
-  CallExpr getCallExpr() { result = node }
+  BlockExprCfgNode() { node = this.getAstNode() }
+
+  /**
+   * Gets the tail expression of this block, if it exists.
+   */
+  ExprCfgNode getTailExpr() {
+    any(ChildMapping mapping).hasCfgChild(node, node.getStmtList().getTailExpr(), this, result)
+  }
+}
+
+/**
+ * A break expression. For example:
+ * ```rust
+ * loop {
+ *     if not_ready() {
+ *         break;
+ *      }
+ * }
+ * ```
+ * ```rust
+ * let x = 'label: loop {
+ *     if done() {
+ *         break 'label 42;
+ *     }
+ * };
+ * ```
+ * ```rust
+ * let x = 'label: {
+ *     if exit() {
+ *         break 'label 42;
+ *     }
+ *     0;
+ * };
+ * ```
+ */
+final class BreakExprCfgNode extends Nodes::BreakExprCfgNode {
+  /**
+   * Gets the target of this `break` expression.
+   *
+   * The target is either a `LoopExpr`, a `ForExpr`, a `WhileExpr`, or a
+   * `BlockExpr`.
+   */
+  ExprCfgNode getTarget() {
+    any(ChildMapping mapping)
+        .hasCfgChild(this.getBreakExpr().getTarget(), this.getBreakExpr(), result, this)
+  }
 }
 
-/** A CFG node that corresponds to a call in the AST. */
-class MethodCallExprCfgNode extends ExprCfgNode {
-  override MethodCallExpr node;
+/**
+ * A function or method call expression. See `CallExpr` and `MethodCallExpr` for further details.
+ */
+final class CallExprBaseCfgNode extends Nodes::CallExprBaseCfgNode {
+  private CallExprBaseChildMapping node;
 
-  /** Gets the underlying `MethodCallExpr`. */
-  MethodCallExpr getMethodCallExpr() { result = node }
+  CallExprBaseCfgNode() { node = this.getAstNode() }
+
+  /** Gets the `i`th argument of this call. */
+  ExprCfgNode getArgument(int i) {
+    any(ChildMapping mapping).hasCfgChild(node, node.getArgList().getArg(i), this, result)
+  }
 }
 
-final class ExitCfgNode = CfgImpl::ExitNode;
+/**
+ * A method call expression. For example:
+ * ```rust
+ * x.foo(42);
+ * x.foo::<u32, u64>(42);
+ * ```
+ */
+final class MethodCallExprCfgNode extends CallExprBaseCfgNode, Nodes::MethodCallExprCfgNode { }
+
+/**
+ * A function call expression. For example:
+ * ```rust
+ * foo(42);
+ * foo::<u32, u64>(42);
+ * foo[0](42);
+ * foo(1) = 4;
+ * ```
+ */
+final class CallExprCfgNode extends CallExprBaseCfgNode, Nodes::CallExprCfgNode { }
+
+/**
+ * A record expression. For example:
+ * ```rust
+ * let first = Foo { a: 1, b: 2 };
+ * let second = Foo { a: 2, ..first };
+ * Foo { a: 1, b: 2 }[2] = 10;
+ * Foo { .. } = second;
+ * ```
+ */
+final class RecordExprCfgNode extends Nodes::RecordExprCfgNode {
+  private RecordExprChildMapping node;
+
+  RecordExprCfgNode() { node = this.getRecordExpr() }
+
+  /** Gets the `i`th record expression. */
+  ExprCfgNode getExpr(int i) {
+    any(ChildMapping mapping)
+        .hasCfgChild(node, node.getRecordExprFieldList().getField(i).getExpr(), this, result)
+  }
+}

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

@@ -40,6 +40,38 @@ private module CfgNodesInput implements InputSig<Location> {
 
 import MakeCfgNodes<Location, CfgNodesInput>
 
+class MatchExprChildMapping extends ParentAstNode, MatchExpr {
+  override predicate relevantChild(AstNode child) {
+    child = this.getAnArm().getPat()
+    or
+    child = this.getAnArm().getGuard().getCondition()
+    or
+    child = this.getAnArm().getExpr()
+  }
+}
+
+class BlockExprChildMapping extends ParentAstNode, BlockExpr {
+  override predicate relevantChild(AstNode child) { child = this.getStmtList().getTailExpr() }
+}
+
+class BreakExprTargetChildMapping extends ParentAstNode, Expr {
+  override predicate relevantChild(AstNode child) { child.(BreakExpr).getTarget() = this }
+}
+
+class CallExprBaseChildMapping extends ParentAstNode, CallExprBase {
+  override predicate relevantChild(AstNode child) { child = this.getArgList().getAnArg() }
+}
+
+class RecordExprChildMapping extends ParentAstNode, RecordExpr {
+  override predicate relevantChild(AstNode child) {
+    child = this.getRecordExprFieldList().getAField().getExpr()
+  }
+}
+
+class FormatArgsExprChildMapping extends ParentAstNode, CfgImpl::ExprTrees::FormatArgsExprTree {
+  override predicate relevantChild(AstNode child) { child = this.getChildNode(_) }
+}
+
 private class ChildMappingImpl extends ChildMapping {
   /** Gets a CFG node for `child`, where `child` is a relevant child node of `parent`. */
   private CfgNode getRelevantChildCfgNode(AstNode parent, AstNode child) {