Chapter 8: Lazy Phis, Break, Continue, and Evaluator

Co-authored-by: Dibyendu Majumdar <mobile@majumdar.org.uk>
Co-authored-by: XmiliaH <45106915+XmiliaH@users.noreply.github.com>
Co-authored-by: Thalia Archibald <thalia@archibald.dev>
Co-authored-by: tannal <tannal2409@gmail.com>
Co-authored-by: hadha <helsgithub15@gmail.com>
Co-authored-by: Tamás Hadházy <85063808+Hels15@users.noreply.github.com>
Co-authored-by: Robert Obkircher <robert.obkircher@gmail.com>
Co-authored-by: Darrick Wiebe <darrickw@untether.ai>

10 files changed, 313 insertions, 42 deletions

diff --git a/src/main/java/com/seaofnodes/simple/GraphEvaluator.java b/src/main/java/com/seaofnodes/simple/GraphEvaluator.java
new file mode 100644
index 00000000..f5694a78
--- /dev/null
+++ b/src/main/java/com/seaofnodes/simple/GraphEvaluator.java
@@ -0,0 +1,181 @@
+package com.seaofnodes.simple;
+
+import com.seaofnodes.simple.node.*;
+import com.seaofnodes.simple.type.TypeInteger;
+
+import java.util.Arrays;
+import java.util.BitSet;
+import java.util.HashMap;
+
+public class GraphEvaluator {
+
+    public enum ResultType { VALUE, FALLTHROUGH, TIMEOUT }
+
+    public record Result(ResultType type, long value) {}
+
+    /**
+     * Find the start node from some node in the graph or null if there is no start node
+     */
+    private static StartNode findStart(BitSet visit, Node node) {
+        if (node == null) return null;
+        if (node instanceof StartNode start) return start;
+        if( visit.get(node._nid) ) return null;
+        visit.set(node._nid);
+        for( Node def : node._inputs ) {
+            var res = findStart(visit,def);
+            if( res != null ) return res;
+        }
+        for( Node use : node._outputs ) {
+            var res = findStart(visit,use);
+            if( res != null ) return res;
+        }
+        return null;
+    }
+
+    /**
+     * Find the control output from a control node
+     */
+    private static Node findControl(Node control) {
+        for( Node use : control._outputs ) {
+            if (use.isCFG()) return use;
+        }
+        return null;
+    }
+
+    /**
+     * Find the projection for a node
+     */
+    private static ProjNode findProjection(Node node, int idx) {
+        for( Node use : node._outputs ) {
+            if (use instanceof ProjNode proj && proj._idx == idx) return proj;
+        }
+        return null;
+    }
+
+    /**
+     * Cache values for phi and parameter projection nodes.
+     */
+    private final HashMap<Node, Long> cacheValues = new HashMap<>();
+    /**
+     * Cache for loop phis as they can depend on itself or other loop phis
+     */
+    private long[] loopPhiCache = new long[16];
+
+    private GraphEvaluator() {
+    }
+
+    private long div(DivNode div) {
+        long in2 = getValue(div.in(2));
+        return in2 == 0 ? 0 : getValue(div.in(1)) / in2;
+    }
+
+    /**
+     * Calculate the value of a node
+     */
+    private long getValue(Node node) {
+        var cache = cacheValues.get(node);
+        if (cache != null) return cache;
+        return switch (node) {
+            case ConstantNode cons  -> ((TypeInteger)cons.compute()).value();
+            case AddNode      add   -> getValue(add.in(1)) + getValue(add.in(2));
+            case BoolNode.EQ  eq    -> getValue(eq.in(1)) == getValue(eq.in(2)) ? 1 : 0;
+            case BoolNode.LE  le    -> getValue(le.in(1)) <= getValue(le.in(2)) ? 1 : 0;
+            case BoolNode.LT  lt    -> getValue(lt.in(1)) <  getValue(lt.in(2)) ? 1 : 0;
+            case DivNode      div   -> div(div);
+            case MinusNode    minus -> -getValue(minus.in(1));
+            case MulNode      mul   -> getValue(mul.in(1)) * getValue(mul.in(2));
+            case NotNode      not   -> getValue(not.in(1)) == 0 ? 1 : 0;
+            case SubNode      sub   -> getValue(sub.in(1)) - getValue(sub.in(2));
+            default                 -> throw Utils.TODO("Unexpected node " + node);
+        };
+    }
+
+    /**
+     * Special case of latchPhis when phis can depend on phis of the same region.
+     */
+    private void latchLoopPhis(RegionNode region, Node prev) {
+        int idx = Utils.find(region._inputs, prev);
+        assert idx > 0;
+        int i = 0;
+        for( Node use : region._outputs ) {
+            if (use instanceof PhiNode phi) {
+                var value = getValue(phi.in(idx));
+                if (i == loopPhiCache.length) loopPhiCache = Arrays.copyOf(loopPhiCache, loopPhiCache.length*2);
+                loopPhiCache[i++] = value;
+            }
+        }
+        i = 0;
+        for( Node use : region._outputs ) {
+            if (use instanceof PhiNode phi) cacheValues.put(phi, loopPhiCache[i++]);
+        }
+    }
+
+    /**
+     * Calculate the values of phis of the region and caches the values. The phis are not allowed to depend on other phis of the region.
+     */
+    private void latchPhis(RegionNode region, Node prev) {
+        int idx = Utils.find(region._inputs, prev);
+        assert idx > 0;
+        for( Node use : region._outputs ) {
+            if (use instanceof PhiNode phi) {
+                var value = getValue(phi.in(idx));
+                cacheValues.put(phi, value);
+            }
+        }
+    }
+
+    /**
+     * Run the graph until either a return is found or the number of loop iterations are done.
+     */
+    private Result evaluate(StartNode start, long parameter, int loops) {
+        var parameter1 = findProjection(start, 1);
+        if (parameter1 != null) cacheValues.put(parameter1, parameter);
+        Node control = findProjection(start, 0);
+        Node prev = start;
+        while (control != null) {
+            Node next;
+            switch (control) {
+                case RegionNode region -> {
+                    if (region instanceof LoopNode && region.in(1) != prev) {
+                        if (loops--<=0) return new Result(ResultType.TIMEOUT, 0);
+                        latchLoopPhis(region, prev);
+                    } else {
+                        latchPhis(region, prev);
+                    }
+                    next = findControl(region);
+                }
+                case IfNode cond -> next = findProjection(cond, getValue(cond.in(1)) != 0 ? 0 : 1);
+                case ReturnNode ret -> {
+                    return new Result(ResultType.VALUE, getValue(ret.in(1)));
+                }
+                case ProjNode ignored -> next = findControl(control);
+                default -> throw Utils.TODO("Unexpected control node " + control);
+            }
+            prev = control;
+            control = next;
+        }
+        return new Result(ResultType.FALLTHROUGH, 0);
+    }
+
+    public static long evaluate(Node graph) {
+        return evaluate(graph, 0);
+    }
+
+    public static long evaluate(Node graph, long parameter) {
+        return evaluate(graph, parameter, 1000);
+    }
+
+    public static long evaluate(Node graph, long parameter, int loops) {
+        var res = evaluateWithResult(graph, parameter, loops);
+        if (res.type==ResultType.TIMEOUT) throw new RuntimeException("Timeout");
+        return res.value;
+    }
+
+    public static Result evaluateWithResult(Node graph, long parameter, int loops) {
+        var start = findStart(new BitSet(), graph);
+        if (start == null) return new Result(ResultType.TIMEOUT, 0);
+        var evaluator = new GraphEvaluator();
+        return evaluator.evaluate(start, parameter, loops);
+    }
+
+}
diff --git a/src/main/java/com/seaofnodes/simple/GraphVisualizer.java b/src/main/java/com/seaofnodes/simple/GraphVisualizer.java
index 07f1af21..d75e5bea 100644
--- a/src/main/java/com/seaofnodes/simple/GraphVisualizer.java
+++ b/src/main/java/com/seaofnodes/simple/GraphVisualizer.java
@@ -26,7 +26,7 @@ public class GraphVisualizer {
         // nodes in the graph.
         Collection<Node> all = findAll(parser);
         StringBuilder sb = new StringBuilder();
-        sb.append("digraph chapter07 {\n");
+        sb.append("digraph chapter08 {\n");
         sb.append("/*\n");
         sb.append(parser.src());
         sb.append("\n*/\n");
@@ -212,6 +212,8 @@ public class GraphVisualizer {
             for( String name : syms.keySet() ) {
                 int idx = syms.get(name);
                 Node def = scope.in(idx);
+                while( def instanceof ScopeNode lazy )
+                    def = lazy.in(idx);
                 if( def==null ) continue;
                 sb.append("\t")
                   .append(scopeName).append(":")
diff --git a/src/main/java/com/seaofnodes/simple/Parser.java b/src/main/java/com/seaofnodes/simple/Parser.java
index cfb4deb4..69796828 100644
--- a/src/main/java/com/seaofnodes/simple/Parser.java
+++ b/src/main/java/com/seaofnodes/simple/Parser.java
@@ -46,6 +46,8 @@ public class Parser {
      * List of keywords disallowed as identifiers
      */
     private final HashSet<String> KEYWORDS = new HashSet<>(){{
+            add("break");
+            add("continue");
             add("else");
             add("false");
             add("if");
@@ -62,10 +64,14 @@ public class Parser {
      */
     public final Stack<ScopeNode> _xScopes = new Stack<>();
 
+    ScopeNode _continueScope;
+    ScopeNode _breakScope;
+
     public Parser(String source, TypeInteger arg) {
         Node.reset();
         _lexer = new Lexer(source);
         _scope = new ScopeNode();
+        _continueScope = _breakScope = null;
         START = new StartNode(new Type[]{ Type.CONTROL, arg });
         STOP = new StopNode();
     }
@@ -136,6 +142,8 @@ public class Parser {
         else if (match ("{"  )) return require(parseBlock(),"}");
         else if (matchx("if" )) return parseIf();
         else if (matchx("while")) return parseWhile();
+        else if (matchx("break")) return parseBreak();
+        else if (matchx("continue")) return parseContinue();
         else if (matchx("#showGraph")) return require(showGraph(),";");
         else if (matchx(";")) return null; // Empty statement
         else return parseExpressionStatement();
@@ -150,6 +158,10 @@ public class Parser {
      * @return a {@link Node}, never {@code null}
      */
     private Node parseWhile() {
+
+        var savedContinueScope = _continueScope;
+        var savedBreakScope    = _breakScope;
+
         require("(");
 
         // Loop region has two control inputs, the first is the entry
@@ -179,36 +191,71 @@ public class Parser {
         ifNode.unkeep();
         Node ifF = new ProjNode(ifNode, 1, "False").peephole();
 
-        // Clone the body Scope to create the exit Scope
-        // which accounts for any side effects in the predicate
-        // The exit Scope will be the final scope after the loop,
-        // And its control input is the False branch of the loop predicate
-        // Note that body Scope is still our current scope
-        var exit = _scope.dup();
-        _xScopes.push(exit);
-        exit.ctrl(ifF);
+        // Clone the body Scope to create the break/exit Scope which accounts for any
+        // side effects in the predicate.  The break/exit Scope will be the final
+        // scope after the loop, and its control input is the False branch of
+        // the loop predicate.  Note that body Scope is still our current scope.
+        ctrl(ifF);
+        _xScopes.push(_breakScope = _scope.dup());
+
+        // No continues yet
+        _continueScope = null;
 
         // Parse the true side, which corresponds to loop body
         // Our current scope is the body Scope
         ctrl(ifT);              // set ctrl token to ifTrue projection
         parseStatement();       // Parse loop body
 
-        // The true branch loops back, so whatever is current control (_scope.ctrl) gets
-        // added to head loop as input. endLoop() updates the head scope,
-        // and goes through all the phis that were created earlier. For each
-        // phi, it sets the second input to the corresponding input from the back edge.
-        // If the phi is redundant, it is replaced by its sole input.
+        // Merge the loop bottom into other continue statements
+        if (_continueScope != null) {
+            _continueScope = jumpTo(_continueScope);
+            _scope.kill();
+            _scope = _continueScope;
+        }
+
+        // The true branch loops back, so whatever is current _scope.ctrl gets
+        // added to head loop as input.  endLoop() updates the head scope, and
+        // goes through all the phis that were created earlier.  For each phi,
+        // it sets the second input to the corresponding input from the back
+        // edge.  If the phi is redundant, it is replaced by its sole input.
+        var exit = _breakScope;
         head.endLoop(_scope, exit);
         head.unkeep().kill();
 
         _xScopes.pop();       // Cleanup
         _xScopes.pop();       // Cleanup
 
+        _continueScope = savedContinueScope;
+        _breakScope = savedBreakScope;
+
         // At exit the false control is the current control, and
         // the scope is the exit scope after the exit test.
-        return (_scope = exit);
+        return _scope = exit;
+    }
+
+    private ScopeNode jumpTo(ScopeNode toScope) {
+        ScopeNode cur = _scope.dup();
+        ctrl(new ConstantNode(Type.XCONTROL).peephole()); // Kill current scope
+        // Prune nested lexical scopes that have depth > than the loop head
+        // We use _breakScope as a proxy for the loop head scope to obtain the depth
+        while( cur._scopes.size() > _breakScope._scopes.size() )
+            cur.pop();
+        // If this is a continue then first time the target is null
+        // So we just use the pruned current scope as the base for the
+        // continue
+        if (toScope == null)
+            return cur;
+        // toScope is either the break scope, or a scope that was created here
+        assert toScope._scopes.size() <= _breakScope._scopes.size();
+        toScope.mergeScopes(cur);
+        return toScope;
     }
 
+    private void checkLoopActive() { if (_breakScope == null) throw Parser.error("No active loop for a break or continue"); }
+
+    private Node parseBreak   () { checkLoopActive(); return (   _breakScope = require(jumpTo(    _breakScope ),";"));  }
+    private Node parseContinue() { checkLoopActive(); return (_continueScope = require(jumpTo( _continueScope ),";"));  }
+
     /**
      * Parses a statement
      *
diff --git a/src/main/java/com/seaofnodes/simple/node/AddNode.java b/src/main/java/com/seaofnodes/simple/node/AddNode.java
index 8f3e5a98..84a368fd 100644
--- a/src/main/java/com/seaofnodes/simple/node/AddNode.java
+++ b/src/main/java/com/seaofnodes/simple/node/AddNode.java
@@ -70,6 +70,11 @@ public class AddNode extends Node {
 
         // Now we only see (add add non)
 
+        // Dead data cycle; comes about from dead infinite loops.  Do nothing,
+        // the loop will peep as dead after a bit.
+        if( lhs.in(1) == lhs )
+            return null;
+
         // Do we have (x + con1) + con2?
         // Replace with (x + (con1+con2) which then fold the constants
         if( lhs.in(2)._type.isConstant() && t2.isConstant() )
diff --git a/src/main/java/com/seaofnodes/simple/node/IfNode.java b/src/main/java/com/seaofnodes/simple/node/IfNode.java
index 0f414c08..23a8fafc 100644
--- a/src/main/java/com/seaofnodes/simple/node/IfNode.java
+++ b/src/main/java/com/seaofnodes/simple/node/IfNode.java
@@ -46,7 +46,7 @@ public class IfNode extends MultiNode {
             return prior instanceof ProjNode proj
               // Repeated test, dominated on one side.  Test result is the same.
               ? (proj._idx==0 ? TypeTuple.IF_TRUE : TypeTuple.IF_FALSE)
-              : null;           // Repeated test not dominated on one side
+              : dom._type;      // Repeated test not dominated on one side
 
         return TypeTuple.IF_BOTH;
     }
diff --git a/src/main/java/com/seaofnodes/simple/node/LoopNode.java b/src/main/java/com/seaofnodes/simple/node/LoopNode.java
index 81f707ec..c192bc20 100644
--- a/src/main/java/com/seaofnodes/simple/node/LoopNode.java
+++ b/src/main/java/com/seaofnodes/simple/node/LoopNode.java
@@ -16,7 +16,8 @@ public class LoopNode extends RegionNode {
 
     @Override
     public Type compute() {
-        return inProgress() ? Type.CONTROL : super.compute();
+        if( inProgress() ) return Type.CONTROL;
+        return entry()._type;
     }
 
     @Override
diff --git a/src/main/java/com/seaofnodes/simple/node/Node.java b/src/main/java/com/seaofnodes/simple/node/Node.java
index 35b07781..2d6d5ad5 100644
--- a/src/main/java/com/seaofnodes/simple/node/Node.java
+++ b/src/main/java/com/seaofnodes/simple/node/Node.java
@@ -1,5 +1,6 @@
 package com.seaofnodes.simple.node;
 
+import com.seaofnodes.simple.IRPrinter;
 import com.seaofnodes.simple.Utils;
 import com.seaofnodes.simple.type.Type;
 
@@ -100,6 +101,7 @@ public abstract class Node {
     public final String print() {
         return _print0(new StringBuilder(), new BitSet()).toString();
     }
+
     // This is the common print: check for repeats, check for DEAD and print
     // "DEAD" else call the per-Node print1.
     final StringBuilder _print0(StringBuilder sb, BitSet visited) {
@@ -138,6 +140,8 @@ public abstract class Node {
         sb.append("\n");
     }
 
+    public String p(int depth) { return IRPrinter.prettyPrint(this,depth); }
+
     public boolean isMultiHead() { return false; }
     public boolean isMultiTail() { return false; }
 
@@ -153,8 +157,6 @@ public abstract class Node {
 
     public int nIns() { return _inputs.size(); }
 
-    public Node out(int i) { return _outputs.get(i); }
-
     public int nOuts() { return _outputs.size(); }
 
     public boolean isUnused() { return nOuts() == 0; }
@@ -443,7 +445,10 @@ public abstract class Node {
      * Used to allow repeating tests in the same JVM.  This just resets the
      * Node unique id generator, and is done as part of making a new Parser.
      */
-    public static void reset() { UNIQUE_ID = 1; _disablePeephole=false; }
+    public static void reset() {
+        UNIQUE_ID = 1;
+        _disablePeephole=false;
+    }
 
     /**
      * Debugging utility to find a Node by index
diff --git a/src/main/java/com/seaofnodes/simple/node/PhiNode.java b/src/main/java/com/seaofnodes/simple/node/PhiNode.java
index 2a4ff52a..c7b5b357 100644
--- a/src/main/java/com/seaofnodes/simple/node/PhiNode.java
+++ b/src/main/java/com/seaofnodes/simple/node/PhiNode.java
@@ -85,10 +85,13 @@ public class PhiNode extends Node {
      */
     private Node singleUniqueInput() {
         Node live = null;
-        for( int i=1; i<nIns(); i++ )
+        for( int i=1; i<nIns(); i++ ) {
+            if( region() instanceof LoopNode loop && loop.entry()._type == Type.XCONTROL )
+                return null;    // Dead entry loops just ignore and let the loop collapse
             if( region().in(i)._type != Type.XCONTROL && in(i) != this )
                 if( live == null || live == in(i) ) live = in(i);
                 else return null;
+        }
         return live;
     }
 
diff --git a/src/main/java/com/seaofnodes/simple/node/RegionNode.java b/src/main/java/com/seaofnodes/simple/node/RegionNode.java
index acaadaf1..453d028f 100644
--- a/src/main/java/com/seaofnodes/simple/node/RegionNode.java
+++ b/src/main/java/com/seaofnodes/simple/node/RegionNode.java
@@ -3,6 +3,7 @@ package com.seaofnodes.simple.node;
 import com.seaofnodes.simple.Utils;
 import com.seaofnodes.simple.type.Type;
 
+import java.util.Collection;
 import java.util.BitSet;
 
 public class RegionNode extends Node {
diff --git a/src/main/java/com/seaofnodes/simple/node/ScopeNode.java b/src/main/java/com/seaofnodes/simple/node/ScopeNode.java
index a4fcbe65..e5d03d93 100644
--- a/src/main/java/com/seaofnodes/simple/node/ScopeNode.java
+++ b/src/main/java/com/seaofnodes/simple/node/ScopeNode.java
@@ -39,6 +39,10 @@ public class ScopeNode extends Node {
         for( int j=0; j<nIns(); j++ ) {
             sb.append(names[j]).append(":");
             Node n = in(j);
+            while( n instanceof ScopeNode loop ) {
+                sb.append("Lazy_");
+                n = loop.in(j);
+            }
             n._print0(sb, visited).append(" ");
         }
         sb.setLength(sb.length()-1);
@@ -107,8 +111,18 @@ public class ScopeNode extends Node {
         var idx = syms.get(name);
         if( idx == null ) return update(name,n,nestingLevel-1); // Not found in this scope, recursively look in parent scope
         Node old = in(idx);
-        // If n is null we are looking up rather than updating, hence return existing value
-        return n==null ? old : setDef(idx,n);
+        if( old instanceof ScopeNode loop ) {
+            // Lazy Phi!
+            old = loop.in(idx) instanceof PhiNode phi && loop.ctrl()==phi.region()
+                // Loop already has a real Phi, use it
+                ? loop.in(idx)
+                // Set real Phi in the loop head
+                // The phi takes its one input (no backedge yet) from a recursive
+                // lookup, which might have insert a Phi in every loop nest.
+                : loop.setDef(idx,new PhiNode(name,loop.ctrl(),loop.update(name,null,nestingLevel),null).peephole());
+            setDef(idx,old);
+        }
+        return n==null ? old : setDef(idx,n); // Not lazy, so this is the answer
     }
 
     public Node ctrl() { return in(0); }
@@ -131,6 +145,11 @@ public class ScopeNode extends Node {
      * tables), nor deep (would dup the Scope, the HashMap tables, but then
      * also the program Nodes).
      * <p>
+     * If the {@code loop} flag is set, the edges are filled in as the original
+     * Scope, as a indication of Lazy Phis at loop heads.  The goal here is to
+     * not make Phis at loop heads for variables which are never touched in the
+     * loop body.
+     * <p>
      * The new Scope is a full-fledged Node with proper use<->def edges.
      */
     public ScopeNode dup() { return dup(false); }
@@ -144,17 +163,10 @@ public class ScopeNode extends Node {
             dup._scopes.push(new HashMap<>(syms));
 
         dup.addDef(ctrl());      // Control input is just copied
-        for( int i=1; i<nIns(); i++ ) {
-            if ( !loop ) { dup.addDef(in(i)); }
-            else {
-                String[] names = reverseNames(); // Get the variable names
-                // Create a phi node with second input as null - to be filled in
-                // by endLoop() below
-                dup.addDef(new PhiNode(names[i], ctrl(), in(i), null).peephole());
-                // Ensure our node has the same phi in case we created one
-                setDef(i, dup.in(i));
-            }
-        }
+        for( int i=1; i<nIns(); i++ )
+            // For lazy phis on loops we use a sentinel
+            // that will trigger phi creation on update
+            dup.addDef(loop ? this : in(i));
         return dup;
     }
 
@@ -172,7 +184,9 @@ public class ScopeNode extends Node {
         // Note that we skip i==0, which is bound to '$ctrl'
         for (int i = 1; i < nIns(); i++)
             if( in(i) != that.in(i) ) // No need for redundant Phis
-                setDef(i, new PhiNode(ns[i], r, in(i), that.in(i)).peephole());
+                // If we are in lazy phi mode we need to a lookup
+                // by name as it will trigger a phi creation
+                setDef(i, new PhiNode(ns[i], r, this.lookup(ns[i]), that.lookup(ns[i])).peephole());
         that.kill();            // Kill merged scope
         return r.unkeep().peephole();
     }
@@ -184,14 +198,26 @@ public class ScopeNode extends Node {
         assert ctrl instanceof LoopNode loop && loop.inProgress();
         ctrl.setDef(2,back.ctrl());
         for( int i=1; i<nIns(); i++ ) {
-            PhiNode phi = (PhiNode)in(i);
-            assert phi.region()==ctrl && phi.in(2)==null;
-            phi.setDef(2,back.in(i));
-            // Do an eager useless-phi removal
-            Node in = phi.peephole();
-            if( in != phi )
-                phi.subsume(in);
+            if( back.in(i) != this ) {
+                PhiNode phi = (PhiNode)in(i);
+                assert phi.region()==ctrl && phi.in(2)==null;
+                phi.setDef(2,back.in(i));
+            }
+            if( exit.in(i) == this ) // Replace a lazy-phi on the exit path also
+                exit.setDef(i,in(i));
         }
         back.kill();            // Loop backedge is dead
+        // Now one-time do a useless-phi removal
+        for( int i=1; i<nIns(); i++ ) {
+            if( in(i) instanceof PhiNode phi ) {
+                // Do an eager useless-phi removal
+                Node in = phi.peephole();
+                if( in != phi ) {
+                    phi.subsume(in);
+                    setDef(i,in); // Set the update back into Scope
+                }
+            }
+        }
+
     }
 }