QL: Support local flow via unification

Author: asgerf

ql/ql/src/codeql_ql/dataflow/DataFlow.qll

@@ -1,9 +1,15 @@
 private import codeql_ql.ast.Ast
 private import internal.NodesInternal
+private import internal.DataFlowNumbering
+private import internal.LocalFlow as LocalFlow
 
 /**
  * An expression or variable in a formula, including some additional nodes
  * that are not part of the AST.
+ *
+ * Nodes that are locally bound together by equalities are clustered into a "super node",
+ * which can be accessed using `getSuperNode()`. There is usually no reason to use `Node` directly
+ * other than to reason about what kind of node is contained in a super node.
  */
 class Node extends TNode {
   string toString() { none() } // overridden in subclasses
@@ -21,6 +27,12 @@ class Node extends TNode {
    * TODO: select clauses
    */
   Predicate getEnclosingPredicate() { none() } // overridden in subclasses
+
+  /**
+   * Gets the collection of data-flow nodes locally bound by equalities, represented
+   * by a "super node".
+   */
+  SuperNode getSuperNode() { result.getANode() = this }
 }
 
 /**
@@ -209,3 +221,71 @@ pragma[inline]
 Node fieldNode(Predicate pred, FieldDecl fieldDecl) {
   result = MkFieldNode(pred, fieldDecl)
 }
+
+/**
+ * A collection of data-flow nodes in the same predicate, locally bound by equalities.
+ */
+class SuperNode extends LocalFlow::TSuperNode {
+  private int repr;
+
+  SuperNode() { this = LocalFlow::MkSuperNode(repr) }
+
+  /** Gets a data-flow node that is part of this super node. */
+  Node getANode() {
+    LocalFlow::getRepr(result) = repr
+  }
+
+  /** Gets an AST node from any of the nodes in this super node. */
+  AstNode asAstNode() {
+    result = getANode().asAstNode()
+  }
+
+  /**
+   * Gets a single node from this super node.
+   *
+   * The node is arbitrary and the caller should not rely on how the node is chosen.
+   * The node is currently chosen such that:
+   * - An `AstNodeNode` is preferred over other nodes.
+   * - A node occuring earlier is preferred over one occurring later.
+   */
+  Node getArbitraryRepr() {
+    result = min(Node n | n = getANode() | n order by getInternalId(n))
+  }
+
+  /**
+   * Gets the predicate containing all nodes that are part of this super node.
+   */
+  Predicate getEnclosingPredicate() {
+    result = getANode().getEnclosingPredicate()
+  }
+
+  /** Gets a string representation of this super node. */
+  string toString() {
+    exists(int c |
+      c = strictcount(getANode()) and
+      result = "Super node of " + c + " nodes in " + getEnclosingPredicate().getName()
+    )
+  }
+
+  /** Gets the location of an arbitrary node in this super node. */
+  Location getLocation() {
+    result = getArbitraryRepr().getLocation()
+  }
+
+  /** Gets any member call whose receiver is in the same super node. */
+  MemberCall getALocalMemberCall() {
+    superNode(result.getBase()) = this
+  }
+
+  /** Gets any member call whose receiver is in the same super node. */
+  MemberCall getALocalMemberCall(string name) {
+    result = this.getALocalMemberCall() and
+    result.getMemberName() = name
+  }
+}
+
+/** Gets the super node for the given AST node. */
+pragma[inline]
+SuperNode superNode(AstNode node) {
+  result = astNode(node).getSuperNode()
+}

ql/ql/src/codeql_ql/dataflow/internal/DataFlowNumbering.qll

@@ -0,0 +1,51 @@
+private import codeql_ql.ast.Ast
+private import codeql_ql.ast.internal.AstNodeNumbering
+private import NodesInternal
+
+/** An arbitrary total ordering of data-flow nodes. */
+private predicate internalOrderingKey(TNode node, int tag, int field1, int field2) {
+  exists(AstNode ast |
+    node = MkAstNodeNode(ast) and
+    tag = 0 and
+    field1 = getPreOrderId(ast) and
+    field2 = 0
+  )
+  or
+  exists(VarDef var, Formula scope |
+    node = MkScopedVariable(var, scope) and
+    tag = 1 and
+    field1 = getPreOrderId(var) and
+    field2 = getPreOrderId(scope)
+  )
+  or
+  exists(Predicate pred |
+    node = MkThisNode(pred) and
+    tag = 2 and
+    field1 = getPreOrderId(pred) and
+    field2 = 0
+  )
+  or
+  exists(Predicate pred |
+    node = MkResultNode(pred) and
+    tag = 3 and
+    field1 = getPreOrderId(pred) and
+    field2 = 0
+  )
+  or
+  exists(Predicate pred, FieldDecl fieldDecl |
+    node = MkFieldNode(pred, fieldDecl) and
+    tag = 4 and
+    field1 = getPreOrderId(pred) and
+    field2 = getPreOrderId(fieldDecl)
+  )
+}
+
+/** Gets an integer unique to `node`. */
+int getInternalId(TNode node) {
+  node =
+    rank[result](TNode n, int tag, int field1, int field2 |
+      internalOrderingKey(n, tag, field1, field2)
+    |
+      n order by tag, field1, field2
+    )
+}

ql/ql/src/codeql_ql/dataflow/internal/LocalFlow.qll

@@ -0,0 +1,102 @@
+/**
+ * Models local flow edges. Each equivalence class in the local flow relation becomes a super node.
+ */
+
+private import codeql_ql.dataflow.DataFlow
+private import codeql_ql.ast.Ast
+private import codeql_ql.ast.internal.AstNodeNumbering
+private import NodesInternal
+private import VarScoping
+private import DataFlowNumbering
+
+private module Cached {
+  /**
+   * Holds if `x` and `y` are bound by an equality (intra-predicate only).
+   *
+   * This edge has no orientation, and is used to construct the equivalence relation.
+   * Each equivalence class becomes a `SuperNode`.
+   */
+  private predicate localEdge(Node x, Node y) {
+    exists(AstNode a, AstNode b |
+      x = astNode(a) and
+      y = astNode(b)
+    |
+      // x ~ any(x)
+      a = b.(Any).getExpr(0)
+      or
+      // v ~ any(T v)
+      a = b.(Any).getArgument(0)
+      or
+      // x ~ x as VAR
+      a = b.(AsExpr).getInnerExpr()
+      or
+      // x ~ x.(Type)
+      a = b.(InlineCast).getBase()
+      or
+      // x = y ==> x ~ y
+      exists(ComparisonFormula compare |
+        compare.getOperator() = "=" and
+        a = compare.getLeftOperand() and
+        b = compare.getRightOperand()
+      )
+    )
+    or
+    // VarAccess -> ScopedVariable
+    exists(VarDef var, VarAccess access, VarAccessOrDisjunct scope |
+      isRefinement(var, access, scope) and
+      x = astNode(access) and
+      y = scopedVariable(var, scope)
+    )
+    or
+    // VarAccess -> VarDef  (if no refinement exists)
+    exists(VarDef var, VarAccess access |
+      isRefinement(var, access, getVarDefScope(var)) and
+      x = astNode(access) and
+      y = astNode(var)
+    )
+    or
+    // result ~ enclosing 'result' node
+    x = resultNode(y.(AstNodeNode).getAstNode().(ResultAccess).getEnclosingPredicate())
+    or
+    // this ~ enclosing 'this' node
+    x = thisNode(y.(AstNodeNode).getAstNode().(ThisAccess).getEnclosingPredicate())
+    or
+    // f ~ enclosing field node for 'f'
+    exists(FieldAccess access |
+      x = astNode(access) and
+      y = fieldNode(access.getEnclosingPredicate(), access.getDeclaration())
+    )
+    or
+    // field declaration ~ field node in the charpred
+    exists(FieldDecl field, Class cls |
+      cls.getAField() = field and
+      x = astNode(field.getVarDecl()) and
+      y = fieldNode(cls.getCharPred(), field)
+    )
+  }
+
+  /** Like `localEdge` but the parameters are mapped to their internal ID. */
+  private predicate rawLocalEdge(int x, int y) {
+    exists(Node a, Node b |
+      localEdge(a, b) and
+      x = getInternalId(a) and
+      y = getInternalId(b)
+    )
+    or
+    // Ensure a representative is generated for singleton components
+    x = getInternalId(_) and
+    y = x
+  }
+
+  /** Gets the representative for the equivalence class containing the node with ID `x`. */
+  private int getRawRepr(int x) = equivalenceRelation(rawLocalEdge/2)(x, result)
+
+  /** Gets the ID for the equivalence class containing `node`. */
+  cached
+  int getRepr(Node node) { result = getRawRepr(getInternalId(node)) }
+
+  cached
+  newtype TSuperNode = MkSuperNode(int repr) { repr = getRepr(_) }
+}
+
+import Cached