2 files changed, 15 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;
     }
 }