Unified: Fix AST, CFG, and data flow rules for assignments

shared/ssa/codeql/ssa/Ssa.qll

@@ -1419,6 +1419,22 @@ module Make<
      */
     default predicate allowFlowIntoUncertainDef(UncertainWriteDefinition def) { none() }
 
+    /**
+     * Holds if the post-update node corresponding to the given `read` occurs at `bb,i`, meaning
+     * it will propagate to the next use (strictly) after that point in the CFG.
+     *
+     * The default is to use the CFG node associated with the read itself, meaning the post-update
+     * always flows to the next use. The default can however lead to spurious flow in cases like:
+     * ```
+     * x.f = foo(x)
+     * ```
+     * where the post-update node for `x` on the left-hand side flows into the `x` on the right-hand side.
+     *
+     * NOTE: When implementing this predicate, you must ensure that `variableRead` is defined to contain
+     * a synthetic read of this variable at `bb,i`.
+     */
+    default predicate postUpdateCfgNode(Expr read, BasicBlock bb, int i) { read.hasCfgNode(bb, i) }
+
     /** An abstract value that a `Guard` may evaluate to. */
     class GuardValue {
       /** Gets a textual representation of this value. */
@@ -1907,7 +1923,10 @@ module Make<
         isUseStep = false
       )
       or
-      [nodeFrom, nodeFrom.(ExprPostUpdateNode).getPreUpdateNode()].(ReadNode).readsAt(bb, i, v) and
+      nodeFrom.(ReadNode).readsAt(bb, i, v) and
+      isUseStep = true
+      or
+      DfInput::postUpdateCfgNode(nodeFrom.(ExprPostUpdateNode).getExpr(), bb, i) and
       isUseStep = true
     }
 

unified/extractor/ast_types.yml

@@ -23,8 +23,6 @@ supertypes:
     - type_cast_expr
     - type_test_expr
     - if_expr
-    - assign_expr
-    - compound_assign_expr
     - pattern_guard_expr
     - empty_expr
     - block
@@ -159,17 +157,6 @@ named:
   unresolved_operator_sequence:
     element*: expr_or_operator
 
-  # Plain assignment
-  assign_expr:
-    target: expr
-    value: expr
-
-  # Compound assignment
-  compound_assign_expr:
-    target: expr
-    operator: infix_operator
-    value: expr
-
   # A function or method call, such as `f(x)` or `obj.m(x)`.
   #
   # Method calls are represented as a call whose `function` is a `member_access_expr`.

unified/extractor/src/languages/swift/swift.rs

@@ -129,13 +129,6 @@ fn member_chain(
     result
 }
 
-/// Compound-assignment operator spellings (`+=`, `<<=`, ...). Used to tell a
-/// compound assignment from an ordinary binary application, both of which
-/// arrive as a `binaryOperator`-based `infixOperatorExpr`.
-const COMPOUND_ASSIGN_OPS: &[&str] = &[
-    "+=", "-=", "*=", "/=", "%=", "<<=", ">>=", "&=", "|=", "^=", "&+=", "&-=", "&*=",
-];
-
 fn translation_rules() -> Vec<Rule<SwiftContext>> {
     vec![
         // ---- Top-level ----
@@ -243,29 +236,21 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
         // operator leaf. Used by `infixOperatorExpr` (folded) and `sequenceExpr`
         // (unresolved).
         rule!((binaryOperatorExpr operator: @op) => (infix_operator #{op})),
-        // Compound assignment (`x += y`) vs. an ordinary binary application
-        // (`a + b`): both are `binaryOperator`-based `infixOperatorExpr`s,
-        // distinguishable only by the operator's spelling. The query engine
-        // can't match on token text, so a small Rust block reads the spelling
-        // and routes to `compound_assign_expr` or `binary_expr`. The operator
-        // is captured raw (`@@op`) to read its spelling.
+        // A `binaryOperator`-based `infixOperatorExpr` represents both ordinary
+        // binary applications (`a + b`) and compound assignments (`x += y`).
+        // Both have the same target AST shape; the QL library distinguishes
+        // assignments by the operator spelling.
         rule!(
             (infixOperatorExpr leftOperand: @l operator: (binaryOperatorExpr) @@op rightOperand: @r)
             =>
-            expr {
-                if COMPOUND_ASSIGN_OPS.contains(&ctx.source_text(op).as_str()) {
-                    tree!((compound_assign_expr target: {l} operator: (infix_operator #{op}) value: {r}))
-                } else {
-                    tree!((binary_expr left: {l} operator: (infix_operator #{op}) right: {r}))
-                }
-            }
+            (binary_expr left: {l} operator: (infix_operator #{op}) right: {r})
         ),
-        // Plain assignment (`x = y`). In a folded chain the `=` is an
-        // `assignmentExpr` node (distinct from other operators), matched by kind.
+        // Plain assignment (`x = y`). In a folded chain the `=` is represented
+        // by an `assignmentExpr` node rather than a `binaryOperatorExpr`.
         rule!(
-            (infixOperatorExpr leftOperand: @l operator: (assignmentExpr) rightOperand: @r)
+            (infixOperatorExpr leftOperand: @l operator: (assignmentExpr) @op rightOperand: @r)
             =>
-            (assign_expr target: {l} value: {r})
+            (binary_expr left: {l} operator: (infix_operator #{op}) right: {r})
         ),
         // In an unresolved `sequenceExpr` (below) the operator positions are not
         // only `binaryOperatorExpr`s: a plain assignment (`=`), an `as`/`is` cast

unified/extractor/tests/corpus/swift/control-flow/if-case-let-with-shadowing-in-condition-value.output

@@ -151,7 +151,7 @@ top_level
           then:
             block
               stmt:
-                compound_assign_expr
-                  target: identifier "y"
+                binary_expr
+                  left: identifier "y"
                   operator: infix_operator "+="
-                  value: int_literal "1"
+                  right: int_literal "1"

unified/extractor/tests/corpus/swift/expressions/array-type-constructor.output

@@ -167,6 +167,7 @@ top_level
                       body:
                         block
                           stmt:
-                            assign_expr
-                              target: identifier "count"
-                              value: int_literal "0"
+                            binary_expr
+                              left: identifier "count"
+                              operator: infix_operator "="
+                              right: int_literal "0"

unified/extractor/tests/corpus/swift/functions/function-with-inout-parameter.output

@@ -71,7 +71,7 @@ top_level
           body:
             block
               stmt:
-                compound_assign_expr
-                  target: identifier "x"
+                binary_expr
+                  left: identifier "x"
                   operator: infix_operator "+="
-                  value: int_literal "1"
+                  right: int_literal "1"

unified/extractor/tests/corpus/swift/loops/repeat-while-loop.output

@@ -51,10 +51,10 @@ top_level
           body:
             block
               stmt:
-                compound_assign_expr
-                  target: identifier "x"
+                binary_expr
+                  left: identifier "x"
                   operator: infix_operator "-="
-                  value: int_literal "1"
+                  right: int_literal "1"
           condition:
             binary_expr
               left: identifier "x"

unified/extractor/tests/corpus/swift/loops/while-loop.output

@@ -57,7 +57,7 @@ top_level
           body:
             block
               stmt:
-                compound_assign_expr
-                  target: identifier "x"
+                binary_expr
+                  left: identifier "x"
                   operator: infix_operator "-="
-                  value: int_literal "1"
+                  right: int_literal "1"

unified/extractor/tests/corpus/swift/types/class-with-initializer.output

@@ -106,9 +106,10 @@ top_level
               body:
                 block
                   stmt:
-                    assign_expr
-                      target:
+                    binary_expr
+                      left:
                         member_access_expr
                           base: identifier "self"
                           member_name_node: identifier "x"
-                      value: identifier "x"
+                      operator: infix_operator "="
+                      right: identifier "x"

unified/extractor/tests/corpus/swift/types/class-with-method.output

@@ -90,7 +90,7 @@ top_level
               body:
                 block
                   stmt:
-                    compound_assign_expr
-                      target: identifier "n"
+                    binary_expr
+                      left: identifier "n"
                       operator: infix_operator "+="
-                      value: int_literal "1"
+                      right: int_literal "1"

unified/extractor/tests/corpus/swift/types/property-with-getter-and-setter.output

@@ -136,6 +136,7 @@ top_level
               body:
                 block
                   stmt:
-                    assign_expr
-                      target: identifier "_v"
-                      value: identifier "newValue"
+                    binary_expr
+                      left: identifier "_v"
+                      operator: infix_operator "="
+                      right: identifier "newValue"

unified/extractor/tests/corpus/swift/variables/assignment.output

@@ -24,6 +24,7 @@ top_level
   body:
     block
       stmt:
-        assign_expr
-          target: identifier "x"
-          value: int_literal "1"
+        binary_expr
+          left: identifier "x"
+          operator: infix_operator "="
+          right: int_literal "1"

unified/extractor/tests/corpus/swift/variables/compound-assignment.output

@@ -24,7 +24,7 @@ top_level
   body:
     block
       stmt:
-        compound_assign_expr
-          target: identifier "x"
+        binary_expr
+          left: identifier "x"
           operator: infix_operator "+="
-          value: int_literal "1"
+          right: int_literal "1"

unified/ql/lib/codeql/unified/internal/Ast.qll

@@ -183,23 +183,6 @@ module Unified {
     final override F::AstNode getAFieldOrChild() { unified_array_literal_element(this, _, result) }
   }
 
-  /** A class representing `assign_expr` nodes. */
-  class AssignExpr extends @unified_assign_expr, F::Expr {
-    /** Gets the name of the primary QL class for this element. */
-    final override string getAPrimaryQlClass() { result = "AssignExpr" }
-
-    /** Gets the node corresponding to the field `target`. */
-    final F::Expr getTarget() { unified_assign_expr_def(this, result, _) }
-
-    /** Gets the node corresponding to the field `value`. */
-    final F::Expr getValue() { unified_assign_expr_def(this, _, result) }
-
-    /** Gets a field or child node of this node. */
-    final override F::AstNode getAFieldOrChild() {
-      unified_assign_expr_def(this, result, _) or unified_assign_expr_def(this, _, result)
-    }
-  }
-
   /** A class representing `associated_type_declaration` nodes. */
   class AssociatedTypeDeclaration extends @unified_associated_type_declaration, F::Member {
     /** Gets the name of the primary QL class for this element. */
@@ -468,28 +451,6 @@ module Unified {
     }
   }
 
-  /** A class representing `compound_assign_expr` nodes. */
-  class CompoundAssignExpr extends @unified_compound_assign_expr, F::Expr {
-    /** Gets the name of the primary QL class for this element. */
-    final override string getAPrimaryQlClass() { result = "CompoundAssignExpr" }
-
-    /** Gets the node corresponding to the field `operator`. */
-    final F::InfixOperator getOperator() { unified_compound_assign_expr_def(this, result, _, _) }
-
-    /** Gets the node corresponding to the field `target`. */
-    final F::Expr getTarget() { unified_compound_assign_expr_def(this, _, result, _) }
-
-    /** Gets the node corresponding to the field `value`. */
-    final F::Expr getValue() { unified_compound_assign_expr_def(this, _, _, result) }
-
-    /** Gets a field or child node of this node. */
-    final override F::AstNode getAFieldOrChild() {
-      unified_compound_assign_expr_def(this, result, _, _) or
-      unified_compound_assign_expr_def(this, _, result, _) or
-      unified_compound_assign_expr_def(this, _, _, result)
-    }
-  }
-
   /** A class representing `conditional_pattern` nodes. */
   class ConditionalPattern extends @unified_conditional_pattern, F::Expr {
     /** Gets the name of the primary QL class for this element. */
@@ -1576,10 +1537,6 @@ module Unified {
       or
       result = node.(ArrayLiteral).getElement(i) and name = "getElement"
       or
-      result = node.(AssignExpr).getTarget() and i = -1 and name = "getTarget"
-      or
-      result = node.(AssignExpr).getValue() and i = -1 and name = "getValue"
-      or
       result = node.(AssociatedTypeDeclaration).getBound() and i = -1 and name = "getBound"
       or
       result = node.(AssociatedTypeDeclaration).getModifier(i) and name = "getModifier"
@@ -1634,12 +1591,6 @@ module Unified {
       or
       result = node.(ClassLikeDeclaration).getTypeParameter(i) and name = "getTypeParameter"
       or
-      result = node.(CompoundAssignExpr).getOperator() and i = -1 and name = "getOperator"
-      or
-      result = node.(CompoundAssignExpr).getTarget() and i = -1 and name = "getTarget"
-      or
-      result = node.(CompoundAssignExpr).getValue() and i = -1 and name = "getValue"
-      or
       result = node.(ConditionalPattern).getCondition() and i = -1 and name = "getCondition"
       or
       result = node.(ConditionalPattern).getModifier(i) and name = "getModifier"
@@ -1879,8 +1830,6 @@ module UnifiedFinal {
 
   final class ArrayLiteral = F::ArrayLiteral;
 
-  final class AssignExpr = F::AssignExpr;
-
   final class AssociatedTypeDeclaration = F::AssociatedTypeDeclaration;
 
   final class BaseType = F::BaseType;
@@ -1907,8 +1856,6 @@ module UnifiedFinal {
 
   final class ClassLikeDeclaration = F::ClassLikeDeclaration;
 
-  final class CompoundAssignExpr = F::CompoundAssignExpr;
-
   final class ConditionalPattern = F::ConditionalPattern;
 
   final class ConstructorDeclaration = F::ConstructorDeclaration;

unified/ql/lib/codeql/unified/internal/AstExtra.qll

@@ -6,6 +6,38 @@ private import unified
 private import codeql.unified.internal.NameBindingPlugin
 
 module Public {
+  /** An assignment, possibly a compound assignment. */
+  abstract class AssignmentImpl extends BinaryExpr {
+    /** Gets the target of this assignment. */
+    abstract Expr getTarget();
+
+    /** Gets the value assigned by this assignment. */
+    abstract Expr getValue();
+  }
+
+  final class Assignment = AssignmentImpl;
+
+  /** A plain assignment expression. */
+  final class AssignExpr extends BinaryExpr, AssignmentImpl {
+    AssignExpr() { this.getOperator().getValue() = "=" }
+
+    override Expr getTarget() { result = this.getLeft() }
+
+    override Expr getValue() { result = this.getRight() }
+  }
+
+  /** A compound assignment expression. */
+  final class CompoundAssignExpr extends BinaryExpr, AssignmentImpl {
+    CompoundAssignExpr() {
+      this.getOperator().getValue() =
+        ["+=", "-=", "*=", "/=", "%=", "<<=", ">>=", "&=", "|=", "^=", "&+=", "&-=", "&*="]
+    }
+
+    override Expr getTarget() { result = this.getLeft() }
+
+    override Expr getValue() { result = this.getRight() }
+  }
+
   /** A short-circuiting logical AND expression. */
   class LogicalAndExpr extends BinaryExpr {
     LogicalAndExpr() { this.getOperator().getValue() = "&&" }

unified/ql/lib/codeql/unified/internal/ControlFlowGraph.qll

@@ -195,20 +195,23 @@ private module Ast implements AstSig<Location> {
 
   class LogicalNotExpr = U::LogicalNotExpr;
 
-  // TODO
-  class Assignment extends BinaryExpr {
-    Assignment() { none() }
-  }
+  class Assignment extends BinaryExpr, U::Assignment { }
 
-  class AssignExpr extends Assignment { }
+  class AssignExpr extends Assignment, U::AssignExpr { }
 
-  class CompoundAssignment extends Assignment { }
+  class CompoundAssignment extends Assignment, U::CompoundAssignExpr { }
 
-  class AssignLogicalAndExpr extends CompoundAssignment { }
+  class AssignLogicalAndExpr extends CompoundAssignment {
+    AssignLogicalAndExpr() { this.getOperator().getValue() = "&&=" }
+  }
 
-  class AssignLogicalOrExpr extends CompoundAssignment { }
+  class AssignLogicalOrExpr extends CompoundAssignment {
+    AssignLogicalOrExpr() { this.getOperator().getValue() = "||=" }
+  }
 
-  class AssignNullCoalescingExpr extends CompoundAssignment { }
+  class AssignNullCoalescingExpr extends CompoundAssignment {
+    AssignNullCoalescingExpr() { this.getOperator().getValue() = "??=" }
+  }
 
   class BooleanLiteral extends U::BooleanLiteral {
     boolean getValue() { result.toString() = super.getValue() }