QL: Add scoped variable nodes

Author: asgerf

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

@@ -45,6 +45,42 @@ Node astNode(AstNode node) {
   result = MkAstNodeNode(node)
 }
 
+/**
+ * A data-flow node representing a variable within a specific scope.
+ */
+class ScopedVariableNode extends Node, MkScopedVariable {
+  private VarDef var;
+  private AstNode scope;
+
+  ScopedVariableNode() { this = MkScopedVariable(var, scope) }
+
+  override string toString() {
+    result = "Variable '" + var.getName() + "' scoped to " + scope.getLocation().getStartLine() + ":" + scope.getLocation().getStartColumn()
+  }
+
+  override Location getLocation() {
+    result = scope.getLocation()
+  }
+
+  /** Gets the variable being refined to a specific scope. */
+  VarDef getVariable() {
+    result = var
+  }
+
+  /** Gets the scope to which this variable has been refined. */
+  AstNode getScope() {
+    result = scope
+  }
+}
+
+/**
+ * Gets the data-flow node corresponding to `var` restricted to `scope`.
+ */
+pragma[inline]
+Node scopedVariable(VarDef var, AstNode scope) {
+  result = MkScopedVariable(var, scope)
+}
+
 /**
  * A data-flow node representing `this` within a class predicate, charpred, or newtype branch.
  */

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

@@ -1,10 +1,15 @@
 private import codeql_ql.ast.Ast
+private import VarScoping
 
 newtype TNode =
   MkAstNodeNode(AstNode node) {
     node instanceof Expr or
     node instanceof VarDef
   } or
+  MkScopedVariable(VarDef var, AstNode scope) {
+    isRefinement(var, _, scope) and
+    not scope = getVarDefScope(var)
+  } or
   MkThisNode(Predicate pred) {
     pred instanceof ClassPredicate or
     pred instanceof CharPred or

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

@@ -0,0 +1,103 @@
+private import codeql_ql.ast.Ast
+private import codeql_ql.ast.internal.AstNodeNumbering
+
+/** Gets the disjunction immediately containing another disjunction `inner`. */
+private Disjunction getOuterDisjunction(Disjunction inner) { result.getAnOperand() = inner }
+
+/**
+ * Get the root of a disjunction tree containing `f`, if any.
+ */
+private Disjunction getRootDisjunction(Disjunction f) {
+  not exists(getOuterDisjunction(result)) and
+  result = getOuterDisjunction(f)
+  or
+  result = getRootDisjunction(getOuterDisjunction(f))
+}
+
+/** Get the root disjunction for `f` if there is one, other gets `f` itself. */
+pragma[inline]
+private AstNode tryGetRootDisjunction(AstNode f) {
+  result = getRootDisjunction(f)
+  or
+  not exists(getRootDisjunction(f)) and
+  result = f
+}
+
+AstNode getADisjunctionOperand(AstNode disjunction) {
+  exists(Disjunction d |
+    result = d.getAnOperand() and
+    // skip intermediate nodes in large disjunctions
+    disjunction = tryGetRootDisjunction(d) and
+    not result instanceof Disjunction
+  )
+  or
+  result = disjunction.(Implication).getAChild()
+  or
+  result = disjunction.(IfFormula).getThenPart()
+  or
+  result = disjunction.(IfFormula).getElsePart()
+  or
+  exists(Forall all |
+    disjunction = all and
+    exists(all.getFormula()) and
+    exists(all.getRange()) and
+    result = [all.getRange(), all.getFormula()]
+  )
+  or
+  result = disjunction.(Set).getAnElement()
+}
+
+/**
+ * A node that acts as a disjunction:
+ * - The root in a tree of `or` operators, or
+ * - An `implies`, `if`, `forall`, or set literal.
+ */
+class DisjunctionOperator extends AstNode {
+  DisjunctionOperator() { exists(getADisjunctionOperand(this)) }
+
+  AstNode getAnOperand() { result = getADisjunctionOperand(this) }
+}
+
+/**
+ * Gets the scope of `var`, such as the predicate or `exists` clause that binds it.
+ */
+AstNode getVarDefScope(VarDef var) {
+  // TODO: not valid for `as` expressions
+  result = var.getParent()
+}
+
+/** A `VarAccess` or disjunct, representing the input to refinement of a variable. */
+class VarAccessOrDisjunct = AstNode;
+
+/**
+ * Walks upwards from an access to `varDef` until encountering either the scope of `varDef`
+ * or a disjunct. When a disjunct is found, the disjunct becomes the new `access`, representing
+ * a refinement we intend to insert there.
+ */
+private AstNode getVarScope(VarDef varDef, VarAccessOrDisjunct access) {
+  access.(VarAccess).getDeclaration() = varDef and
+  result = access
+  or
+  exists(AstNode scope | scope = getVarScope(varDef, access) |
+    not scope = getADisjunctionOperand(_) and
+    not scope = getVarDefScope(varDef) and
+    result = scope.getParent()
+  )
+  or
+  isRefinement(varDef, _, access) and
+  result = tryGetRootDisjunction(access.getParent())
+}
+
+/**
+ * Holds if `inner` should be seen as a refinement of `outer`.
+ *
+ * `outer` is always a disjunct, and `inner` is either a `VarAccess` or another disjunct.
+ */
+predicate isRefinement(VarDef varDef, VarAccessOrDisjunct inner, VarAccessOrDisjunct outer) {
+  getVarScope(varDef, inner) = outer and
+  (
+    outer = getADisjunctionOperand(_)
+    or
+    outer = getVarDefScope(varDef)
+  )
+}