Fix bug loop tree depth on inf loops

3 files changed, 51 insertions, 4 deletions

diff --git a/chapter18/src/main/java/com/seaofnodes/simple/node/CFGNode.java b/chapter18/src/main/java/com/seaofnodes/simple/node/CFGNode.java
index ece46e80..49398838 100644
--- a/chapter18/src/main/java/com/seaofnodes/simple/node/CFGNode.java
+++ b/chapter18/src/main/java/com/seaofnodes/simple/node/CFGNode.java
@@ -99,7 +99,7 @@ public abstract class CFGNode extends Node {
     // ------------------------------------------------------------------------
     // Tag all CFG Nodes with their containing LoopNode; LoopNodes themselves
     // also refer to *their* containing LoopNode, as well as have their depth.
-    // Start is a LoopNode which contains all at depth 1.
+    // Start contains the function bodies at depth 0; real loops add depth.
     public void buildLoopTree(StopNode stop) {
         _ltree = stop._ltree = Parser.XCTRL._ltree = new LoopTree((StartNode)this);
         _bltWalk(2,null,stop, new BitSet());
@@ -108,6 +108,14 @@ public abstract class CFGNode extends Node {
         // Pre-walked?
         if( _pre!=0 ) return pre;
         _pre = pre++;
+        // Seed the function's outer tree before walking its body.  Synthetic
+        // exits can be added before the post-order walk reaches the entry.
+        if( this instanceof FunNode ) _ltree = stop._ltree;
+        if( this instanceof ReturnNode ) {
+            _ltree = fun._ltree;
+            post.set(_nid);
+            return pre; // Stop/Start links describe the outside world, not a loop.
+        }
         // Pre-walk
         for( Node use : _outputs )
             if( use instanceof CFGNode usecfg && !skip(usecfg) )
diff --git a/chapter18/src/main/java/com/seaofnodes/simple/node/LoopNode.java b/chapter18/src/main/java/com/seaofnodes/simple/node/LoopNode.java
index 097ddea7..f26df076 100644
--- a/chapter18/src/main/java/com/seaofnodes/simple/node/LoopNode.java
+++ b/chapter18/src/main/java/com/seaofnodes/simple/node/LoopNode.java
@@ -30,7 +30,7 @@ public class LoopNode extends RegionNode {
 
     // If this is an unreachable loop, it may not have an exit.  If it does not
     // (i.e., infinite loop), force an exit to make it reachable.
-    public StopNode forceExit( FunNode fun, StopNode stop ) {
+    public FunNode forceExit( FunNode fun, StopNode stop ) {
         // Walk the backedge, then immediate dominator tree util we hit this
         // Loop again.  If we ever hit a CProj from an If (as opposed to
         // directly on the If) we found our exit.
@@ -39,7 +39,7 @@ public class LoopNode extends RegionNode {
             if( x instanceof CProjNode exit && exit.in(0) instanceof IfNode iff ) {
                 CFGNode other = iff.cproj(1-exit._idx);
                 if( other!=null && other.loopDepth() < loopDepth() )
-                    return stop; // Found an exit, not an infinite loop
+                    return fun; // Found an exit, not an infinite loop
             }
             x = x.idom();
         }
@@ -56,6 +56,8 @@ public class LoopNode extends RegionNode {
         // XCtrl or a bunch of prior NeverNode exits.
         Node top = new ConstantNode(Type.TOP).peephole();
         ReturnNode ret = fun.ret();
+        iff._ltree = f._ltree = _ltree;
+        ret._ltree = t._ltree = fun._ltree;
         Node ctrl = ret.ctrl(), mem = ret.mem(), expr = ret.expr();
         if( ctrl._type != Type.XCONTROL ) {
             // Perfect aligned exit?
@@ -76,10 +78,11 @@ public class LoopNode extends RegionNode {
             mem  = top;
             expr = top;
         }
+        ((CFGNode)ctrl)._ltree = fun._ltree;
         ret.setDef(0,ctrl);
         ret.setDef(1,mem );
         ret.setDef(2,expr);
 
-        return stop;
+        return fun;
     }
 }
diff --git a/chapter18/src/test/java/com/seaofnodes/simple/Chapter18Test.java b/chapter18/src/test/java/com/seaofnodes/simple/Chapter18Test.java
index d13514b5..c08bc5b2 100644
--- a/chapter18/src/test/java/com/seaofnodes/simple/Chapter18Test.java
+++ b/chapter18/src/test/java/com/seaofnodes/simple/Chapter18Test.java
@@ -9,6 +9,42 @@ import static org.junit.Assert.fail;
 import org.junit.Ignore;
 
 public class Chapter18Test {
+    @Test public void testNeverExitLoopDepth() {
+        String[] sources = {
+            "if(arg) while(1) {} return 7;",
+            "if(arg==1) while(1) {} if(arg==2) while(1) {} return 7;",
+            "while(1) {}",
+            "int x=0; while(arg) { if(arg==2) while(1) { x++; } x++; arg--; } return x;"
+        };
+        for( String src : sources ) {
+            CodeGen code = new CodeGen(src).parse().opto().typeCheck();
+            code._start.buildLoopTree(code._stop);
+            int[] nevers = {0};
+            code._stop.walk(n -> {
+                if( n instanceof FunNode fun )
+                    assertEquals("Function entry is outside its loops",0,fun.loopDepth());
+                if( n instanceof ReturnNode ret ) {
+                    assertEquals(ret.fun().loopDepth(),ret.loopDepth());
+                    assertEquals(ret.fun().loopDepth(),ret.cfg0().loopDepth());
+                }
+                if( n instanceof LoopNode loop && !(loop instanceof StartNode) &&
+                    loop.back() instanceof CProjNode back && back.in(0) instanceof NeverNode never ) {
+                    nevers[0]++;
+                    assertSame(loop,never.loop());
+                    assertSame(loop,back.loop());
+                    CProjNode exit = never.cproj(1-back._idx);
+                    CFGNode fun = loop;
+                    while( !(fun instanceof FunNode) ) fun=fun.idom();
+                    assertSame(fun.loop(),exit.loop());
+                    assertEquals(fun.loopDepth(),exit.loopDepth());
+                    assertTrue(loop.loopDepth()>fun.loopDepth());
+                }
+                return null;
+            });
+            assertTrue("Exercise a synthetic exit",nevers[0]>0);
+        }
+    }
+
     @Test public void testFunctionLocalConstantChains() {
         String src = "val f = { int x -> x ? f(x-1)*305420988+x/305420988 : 1; }; "+
                      "val g = { int x -> x ? g(x-1)/305420988+x*305420988 : 1; }; "+