Propagate loop count throughout the lexical context

This should make sure that the method containing a loop gets compiled properly, and not just a block that contains a loop.

Signed-off-by: Stefan Marr <git@stefan-marr.de>

Author: smarr

src/som/compiler/MethodGenerationContext.java

@@ -190,13 +190,22 @@ public SMethod assemble(final Universe universe, ExpressionNode methodBody) {
         new som.interpreter.Method(getSourceSectionForMethod(sourceSection),
             frameDescriptor, methodBody, universe, getLexicalContext());
 
+    setOuterMethodInLexicalScopes(truffleMethod);
+
     SMethod meth = (SMethod) universe.newMethod(signature, truffleMethod, false,
         embeddedBlockMethods.toArray(new SMethod[0]));
 
     // return the method - the holder field is to be set later on!
     return meth;
   }
 
+  private void setOuterMethodInLexicalScopes(final som.interpreter.Method method) {
+    for (SMethod m : embeddedBlockMethods) {
+      som.interpreter.Method blockMethod = (som.interpreter.Method) m.getInvokable();
+      blockMethod.setOuterContextMethod(method);
+    }
+  }
+
   private SourceSection getSourceSectionForMethod(final SourceSection ssBody) {
     SourceSection ssMethod = ssBody.getSource().createSection(
         holderGenc.getName().getString() + ">>" + signature.toString(),

src/som/interpreter/Invokable.java

@@ -43,4 +43,8 @@ public final boolean isSplittable() {
   public final RootCallTarget createCallTarget() {
     return Truffle.getRuntime().createCallTarget(this);
   }
+
+  public abstract void propagateLoopCountThroughoutLexicalScope(final long count);
+
+  public abstract boolean isBlock();
 }

src/som/interpreter/LexicalContext.java

@@ -1,13 +1,17 @@
 package som.interpreter;
 
+import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.FrameDescriptor;
 
 
 public final class LexicalContext {
   private final FrameDescriptor frameDescriptor;
   private final LexicalContext  lexicalContext;
+  @CompilationFinal private Method outerMethod;
 
-  public LexicalContext(final FrameDescriptor frameDescriptor, final LexicalContext outerContext) {
+  public LexicalContext(final FrameDescriptor frameDescriptor,
+      final LexicalContext outerContext) {
     this.frameDescriptor = frameDescriptor;
     this.lexicalContext  = outerContext;
   }
@@ -19,4 +23,14 @@ public FrameDescriptor getFrameDescriptor() {
   public LexicalContext getOuterContext() {
     return lexicalContext;
   }
+
+  public Method getOuterMethod() {
+    return outerMethod;
+  }
+
+  public void setOuterMethod(final Method method) {
+    CompilerAsserts.neverPartOfCompilation();
+    assert outerMethod == null; // should not have been set before
+    outerMethod = method;
+  }
 }

src/som/interpreter/Method.java

@@ -59,8 +59,29 @@ public Invokable cloneWithNewLexicalContext(final LexicalContext outerContext) {
         outerContext);
     ExpressionNode  inlinedBody = Inliner.doInline(getUninitializedBody(),
         inlinedContext);
-    return new Method(getSourceSection(), inlinedFrameDescriptor, inlinedBody,
+    Method clone = new Method(getSourceSection(), inlinedFrameDescriptor, inlinedBody,
         universe, outerContext);
+    inlinedContext.setOuterMethod(clone);
+    return clone;
+  }
+
+  @Override
+  public boolean isBlock() {
+    return outerContext != null;
+  }
+
+  public void setOuterContextMethod(final Method method) {
+    outerContext.setOuterMethod(method);
+  }
+
+  @Override
+  public void propagateLoopCountThroughoutLexicalScope(final long count) {
+    assert count >= 0;
+
+    if (outerContext != null) {
+      outerContext.getOuterMethod().propagateLoopCountThroughoutLexicalScope(count);
+    }
+    reportLoopCount((count > Integer.MAX_VALUE) ? Integer.MAX_VALUE : (int) count);
   }
 
   @Override

src/som/interpreter/Primitive.java

@@ -28,8 +28,19 @@ public RootNode split() {
     return cloneWithNewLexicalContext(null);
   }
 
+  @Override
+  public boolean isBlock() {
+    return false;
+  }
+
   @Override
   public String toString() {
     return "Primitive " + expressionOrSequence.getClass().getSimpleName() + "@" + Integer.toHexString(hashCode());
   }
+
+  @Override
+  public void propagateLoopCountThroughoutLexicalScope(final long count) {
+    throw new UnsupportedOperationException(
+        "This should not happen, primitives don't have lexically nested loops.");
+  }
 }

src/som/interpreter/nodes/specialized/IntToByDoMessageNode.java

@@ -1,5 +1,6 @@
 package som.interpreter.nodes.specialized;
 
+import som.interpreter.Invokable;
 import som.interpreter.nodes.ExpressionNode;
 import som.interpreter.nodes.PreevaluatedExpression;
 import som.interpreter.nodes.nary.QuaternaryExpressionNode;
@@ -8,13 +9,11 @@
 
 import com.oracle.truffle.api.CompilerAsserts;
 import com.oracle.truffle.api.CompilerDirectives;
-import com.oracle.truffle.api.LoopCountReceiver;
 import com.oracle.truffle.api.Truffle;
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.DirectCallNode;
 import com.oracle.truffle.api.nodes.Node;
-import com.oracle.truffle.api.nodes.RootNode;
 
 
 public abstract class IntToByDoMessageNode extends QuaternaryExpressionNode
@@ -74,7 +73,7 @@ public final long doIntToByDo(final VirtualFrame frame, final long receiver, fin
       }
     } finally {
       if (CompilerDirectives.inInterpreter()) {
-        reportLoopCount((int) limit - receiver);
+        reportLoopCount((long) limit - receiver);
       }
     }
     return receiver;
@@ -83,14 +82,11 @@ public final long doIntToByDo(final VirtualFrame frame, final long receiver, fin
   protected final void reportLoopCount(final long count) {
     CompilerAsserts.neverPartOfCompilation();
     Node current = getParent();
-    while (current != null && !(current instanceof RootNode)) {
+    while (current != null && !(current instanceof Invokable)) {
       current = current.getParent();
     }
     if (current != null) {
-      RootNode root = (RootNode) current;
-      if (root.getCallTarget() instanceof LoopCountReceiver) {
-        ((LoopCountReceiver) root.getCallTarget()).reportLoopCount((int) count);
-      }
+      ((Invokable) current).propagateLoopCountThroughoutLexicalScope(count);
     }
   }
 }

src/som/interpreter/nodes/specialized/IntToDoMessageNode.java

@@ -1,5 +1,6 @@
 package som.interpreter.nodes.specialized;
 
+import som.interpreter.Invokable;
 import som.interpreter.nodes.ExpressionNode;
 import som.interpreter.nodes.PreevaluatedExpression;
 import som.interpreter.nodes.nary.TernaryExpressionNode;
@@ -9,7 +10,6 @@
 import com.oracle.truffle.api.CompilerAsserts;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.SlowPath;
-import com.oracle.truffle.api.LoopCountReceiver;
 import com.oracle.truffle.api.Truffle;
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
@@ -88,10 +88,7 @@ private void reportLoopCount(final long count) {
       current = current.getParent();
     }
     if (current != null) {
-      RootNode root = (RootNode) current;
-      if (root.getCallTarget() instanceof LoopCountReceiver) {
-        ((LoopCountReceiver) root.getCallTarget()).reportLoopCount((int) count);
-      }
+      ((Invokable) current).propagateLoopCountThroughoutLexicalScope(count);
     }
   }
 }

src/som/interpreter/nodes/specialized/WhileWithStaticBlocksNode.java

@@ -1,5 +1,6 @@
 package som.interpreter.nodes.specialized;
 
+import som.interpreter.Invokable;
 import som.interpreter.nodes.literals.BlockNode;
 import som.interpreter.nodes.nary.BinaryExpressionNode;
 import som.vm.Universe;
@@ -9,7 +10,6 @@
 import com.oracle.truffle.api.CallTarget;
 import com.oracle.truffle.api.CompilerAsserts;
 import com.oracle.truffle.api.CompilerDirectives;
-import com.oracle.truffle.api.LoopCountReceiver;
 import com.oracle.truffle.api.Truffle;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.DirectCallNode;
@@ -73,7 +73,7 @@ public final void executeEvaluatedVoid(final VirtualFrame frame,
   protected final SObject doWhileConditionally(final VirtualFrame frame,
       final SBlock loopCondition,
       final SBlock loopBody) {
-    int iterationCount = 0;
+    long iterationCount = 0;
     boolean loopConditionResult = (boolean) conditionValueSend.call(
         frame, new Object[] {loopCondition});
 
@@ -98,17 +98,14 @@ protected final SObject doWhileConditionally(final VirtualFrame frame,
     return universe.nilObject;
   }
 
-  protected final void reportLoopCount(final int count) {
+  protected final void reportLoopCount(final long count) {
     CompilerAsserts.neverPartOfCompilation();
     Node current = getParent();
     while (current != null && !(current instanceof RootNode)) {
       current = current.getParent();
     }
     if (current != null) {
-      RootNode root = (RootNode) current;
-      if (root.getCallTarget() instanceof LoopCountReceiver) {
-        ((LoopCountReceiver) root.getCallTarget()).reportLoopCount(count);
-      }
+      ((Invokable) current).propagateLoopCountThroughoutLexicalScope(count);
     }
   }