Fix Derivation Across Multiple Passes

Variable nodes should have origin nodes so when expressions are simplified in multiple passes, previous simplifications are preserved

Author: rcosta358

liquidjava-verifier/src/main/java/liquidjava/rj_language/opt/ConstantPropagation.java

@@ -10,6 +10,7 @@
 import liquidjava.rj_language.opt.derivation_node.ValDerivationNode;
 import liquidjava.rj_language.opt.derivation_node.VarDerivationNode;
 
+import java.util.HashMap;
 import java.util.Map;
 
 public class ConstantPropagation {
@@ -19,64 +20,122 @@ public class ConstantPropagation {
      * VariableResolver to extract variable equalities from the expression first. Returns a derivation node representing
      * the propagation steps taken.
      */
-    public static ValDerivationNode propagate(Expression exp, DerivationNode defaultOrigin) {
+    public static ValDerivationNode propagate(Expression exp, ValDerivationNode previousOrigin) {
         Map<String, Expression> substitutions = VariableResolver.resolve(exp);
-        return propagateRecursive(exp, substitutions, defaultOrigin);
+
+        // map of variable origins from the previous derivation tree
+        Map<String, DerivationNode> varOrigins = new HashMap<>();
+        if (previousOrigin != null) {
+            extractVarOrigins(previousOrigin, varOrigins);
+        }
+        return propagateRecursive(exp, substitutions, varOrigins);
     }
 
     /**
      * Recursively performs constant propagation on an expression (e.g. x + y && x == 1 && y == 2 => 1 + 2)
      */
     private static ValDerivationNode propagateRecursive(Expression exp, Map<String, Expression> subs,
-            DerivationNode defaultOrigin) {
+            Map<String, DerivationNode> varOrigins) {
 
         // substitute variable
         if (exp instanceof Var var) {
             String name = var.getName();
             Expression value = subs.get(name);
             // substitution
-            if (value != null)
-                return new ValDerivationNode(value.clone(), new VarDerivationNode(name));
+            if (value != null) {
+                // check if this variable has an origin from a previous pass
+                DerivationNode previousOrigin = varOrigins.get(name);
+                
+                // preserve origin if value came from previous derivation
+                DerivationNode origin = previousOrigin != null ? new VarDerivationNode(name, previousOrigin) : new VarDerivationNode(name);
+                return new ValDerivationNode(value.clone(), origin);
+            }
 
             // no substitution
-            return new ValDerivationNode(var, defaultOrigin);
+            return new ValDerivationNode(var, null);
         }
 
         // lift unary origin
         if (exp instanceof UnaryExpression unary) {
-            ValDerivationNode operand = propagateRecursive(unary.getChildren().get(0), subs, defaultOrigin);
+            ValDerivationNode operand = propagateRecursive(unary.getChildren().get(0), subs, varOrigins);
             UnaryExpression cloned = (UnaryExpression) unary.clone();
             cloned.setChild(0, operand.getValue());
 
-            DerivationNode origin = operand.getOrigin() != null ? new UnaryDerivationNode(operand, cloned.getOp())
-                    : defaultOrigin;
-            return new ValDerivationNode(cloned, origin);
+            return operand.getOrigin() != null
+                    ? new ValDerivationNode(cloned, new UnaryDerivationNode(operand, cloned.getOp()))
+                    : new ValDerivationNode(cloned, null);
         }
 
         // lift binary origin
         if (exp instanceof BinaryExpression binary) {
-            ValDerivationNode left = propagateRecursive(binary.getFirstOperand(), subs, defaultOrigin);
-            ValDerivationNode right = propagateRecursive(binary.getSecondOperand(), subs, defaultOrigin);
+            ValDerivationNode left = propagateRecursive(binary.getFirstOperand(), subs, varOrigins);
+            ValDerivationNode right = propagateRecursive(binary.getSecondOperand(), subs, varOrigins);
             BinaryExpression cloned = (BinaryExpression) binary.clone();
             cloned.setChild(0, left.getValue());
             cloned.setChild(1, right.getValue());
 
-            DerivationNode origin = (left.getOrigin() != null || right.getOrigin() != null)
-                    ? new BinaryDerivationNode(left, right, cloned.getOperator()) : defaultOrigin;
-            return new ValDerivationNode(cloned, origin);
+            return (left.getOrigin() != null || right.getOrigin() != null)
+                    ? new ValDerivationNode(cloned, new BinaryDerivationNode(left, right, cloned.getOperator()))
+                    : new ValDerivationNode(cloned, null);
         }
 
         // recursively propagate children
         if (exp.hasChildren()) {
             Expression propagated = exp.clone();
             for (int i = 0; i < exp.getChildren().size(); i++) {
-                ValDerivationNode child = propagateRecursive(exp.getChildren().get(i), subs, defaultOrigin);
+                ValDerivationNode child = propagateRecursive(exp.getChildren().get(i), subs, varOrigins);
                 propagated.setChild(i, child.getValue());
             }
-            return new ValDerivationNode(propagated, defaultOrigin);
+            return new ValDerivationNode(propagated, null);
         }
 
         // no propagation
-        return new ValDerivationNode(exp, defaultOrigin);
+        return new ValDerivationNode(exp, null);
+    }
+
+
+    /**
+     * Extracts the derivation nodes for variable values from the derivation tree
+     * This is so done so when we find "var == value" in the tree, we store the derivation of the value
+     * So it can be preserved when var is substituted in subsequent passes
+     */
+    private static void extractVarOrigins(ValDerivationNode node, Map<String, DerivationNode> varOrigins) {
+        if (node == null)
+            return;
+
+        Expression value = node.getValue();
+        DerivationNode origin = node.getOrigin();
+
+        // check for equality expressions
+        if (value instanceof BinaryExpression binExp && "==".equals(binExp.getOperator())
+                && origin instanceof BinaryDerivationNode binOrigin) {
+            Expression left = binExp.getFirstOperand();
+            Expression right = binExp.getSecondOperand();
+
+            // extract variable name and value derivation from either side
+            String varName = null;
+            ValDerivationNode valueDerivation = null;
+
+            if (left instanceof Var var && right.isLiteral()) {
+                varName = var.getName();
+                valueDerivation = binOrigin.getRight();
+            } else if (right instanceof Var var && left.isLiteral()) {
+                varName = var.getName();
+                valueDerivation = binOrigin.getLeft();
+            }
+            if (varName != null && valueDerivation != null && valueDerivation.getOrigin() != null) {
+                varOrigins.put(varName, valueDerivation.getOrigin());
+            }
+        }
+
+        // recursively process the origin tree
+        if (origin instanceof BinaryDerivationNode binOrigin) {
+            extractVarOrigins(binOrigin.getLeft(), varOrigins);
+            extractVarOrigins(binOrigin.getRight(), varOrigins);
+        } else if (origin instanceof UnaryDerivationNode unaryOrigin) {
+            extractVarOrigins(unaryOrigin.getOperand(), varOrigins);
+        } else if (origin instanceof ValDerivationNode valOrigin) {
+            extractVarOrigins(valOrigin, varOrigins);
+        }
     }
 }

liquidjava-verifier/src/main/java/liquidjava/rj_language/opt/derivation_node/VarDerivationNode.java

@@ -3,12 +3,23 @@
 public class VarDerivationNode extends DerivationNode {
 
     private final String var;
+    private final DerivationNode origin;
 
     public VarDerivationNode(String var) {
         this.var = var;
+        this.origin = null;
+    }
+
+    public VarDerivationNode(String var, DerivationNode origin) {
+        this.var = var;
+        this.origin = origin;
     }
 
     public String getVar() {
         return var;
     }
+
+    public DerivationNode getOrigin() {
+        return origin;
+    }
 }

liquidjava-verifier/src/test/java/liquidjava/rj_language/opt/ExpressionSimplifierTest.java

@@ -335,35 +335,24 @@ void testFixedPointSimplification() {
 
         // Compare derivation tree structure
 
-        // Origin of y (value 2) - right operand of result
-        ValDerivationNode originY = new ValDerivationNode(new LiteralInt(2), new VarDerivationNode("y"));
-        UnaryDerivationNode originNeg2 = new UnaryDerivationNode(originY, "-");
-        ValDerivationNode rightNode = new ValDerivationNode(new LiteralInt(-2), originNeg2);
+        // Build the derivation chain for the right side:
+        // 6 came from a, 3 came from b
+        ValDerivationNode val6FromA = new ValDerivationNode(new LiteralInt(6), new VarDerivationNode("a"));
+        ValDerivationNode val3FromB = new ValDerivationNode(new LiteralInt(3), new VarDerivationNode("b"));
 
-        // Origin of x - left operand of result
-        // 6 (from a) / 3 (from b) -> 2
-        ValDerivationNode val6 = new ValDerivationNode(new LiteralInt(6), new VarDerivationNode("a"));
-        ValDerivationNode val3 = new ValDerivationNode(new LiteralInt(3), new VarDerivationNode("b"));
-        BinaryDerivationNode divOp = new BinaryDerivationNode(val6, val3, "/");
-        ValDerivationNode val2FromDiv = new ValDerivationNode(new LiteralInt(2), divOp);
-
-        // y == 2 (from y == 6 / 3)
-        ValDerivationNode valYNode = new ValDerivationNode(new Var("y"), null);
-        BinaryDerivationNode eqYOp = new BinaryDerivationNode(valYNode, val2FromDiv, "==");
-        ValDerivationNode yEq2 = new ValDerivationNode(new BinaryExpression(new Var("y"), "==", new LiteralInt(2)),
-                eqYOp);
-
-        // x == -y
-        ValDerivationNode xEqNegY = new ValDerivationNode(
-                new BinaryExpression(new Var("x"), "==", new UnaryExpression("-", new Var("y"))), null);
-
-        // x == -y && y == 2
-        BinaryDerivationNode andOp1 = new BinaryDerivationNode(xEqNegY, yEq2, "&&");
-        ValDerivationNode xEqNegYAndYEq2 = new ValDerivationNode(
-                new BinaryExpression(xEqNegY.getValue(), "&&", yEq2.getValue()), andOp1);
-
-        // Left node x has origin pointing to the previous simplification's tree
-        ValDerivationNode leftNode = new ValDerivationNode(new Var("x"), xEqNegYAndYEq2);
+        // 6 / 3 -> 2
+        BinaryDerivationNode divOrigin = new BinaryDerivationNode(val6FromA, val3FromB, "/");
+
+        // 2 came from y, and y's value came from 6 / 2
+        VarDerivationNode yChainedOrigin = new VarDerivationNode("y", divOrigin);
+        ValDerivationNode val2FromY = new ValDerivationNode(new LiteralInt(2), yChainedOrigin);
+
+        // -2
+        UnaryDerivationNode negOrigin = new UnaryDerivationNode(val2FromY, "-");
+        ValDerivationNode rightNode = new ValDerivationNode(new LiteralInt(-2), negOrigin);
+
+        // Left node x has no origin
+        ValDerivationNode leftNode = new ValDerivationNode(new Var("x"), null);
 
         // Root equality
         BinaryDerivationNode rootOrigin = new BinaryDerivationNode(leftNode, rightNode, "==");