Fix issues with PE: move context level for non-local returns to bytecode node

src/trufflesom/src/trufflesom/compiler/bc/BytecodeGenerator.java

@@ -146,7 +146,7 @@ public static void emitRETURNSELF(final BytecodeMethodGenContext mgenc) {
   }
 
   public static void emitRETURNNONLOCAL(final BytecodeMethodGenContext mgenc) {
-    emit2(mgenc, RETURN_NON_LOCAL, mgenc.getMaxContextLevel(), 0);
+    emit1(mgenc, RETURN_NON_LOCAL, 0);
   }
 
   public static void emitRETURNFIELD(final BytecodeMethodGenContext mgenc, final byte idx) {

src/trufflesom/src/trufflesom/compiler/bc/BytecodeMethodGenContext.java

@@ -329,7 +329,7 @@ private BytecodeLoopNode constructBytecodeBody(final long coord) {
     BackJump[] loops = inlinedLoops.toArray(new BackJump[0]);
 
     return new BytecodeLoopNode(bytecodes, locals.size(), literalsArr, maxStackDepth,
-        frameOnStackMarkerIndex, loops);
+        frameOnStackMarkerIndex, loops, getMaxContextLevel());
   }
 
   public byte[] getBytecodeArray() {

src/trufflesom/src/trufflesom/compiler/bc/Disassembler.java

@@ -320,7 +320,9 @@ public static void dumpMethod(final List<Byte> bytecodes, final String indent,
         }
 
         case RETURN_NON_LOCAL: {
-          Universe.errorPrintln("context: " + bytecodes.get(b + 1));
+          if (m != null) {
+            Universe.errorPrintln("context: " + m.getContextLevel());
+          }
           break;
         }
 

src/trufflesom/src/trufflesom/interpreter/bc/Bytecodes.java

@@ -297,7 +297,7 @@ public static final boolean isOneOf(final byte bytecode, final byte[] arr) {
         2, // SEND
         2, // SUPER_SEND
         1, // RETURN_LOCAL
-        2, // RETURN_NON_LOCAL
+        1, // RETURN_NON_LOCAL
         1, // RETURN_SELF
 
         1, // RETURN_FIELD_0

src/trufflesom/src/trufflesom/interpreter/nodes/bc/BytecodeLoopNode.java

@@ -155,14 +155,16 @@ public class BytecodeLoopNode extends NoPreEvalExprNode implements ScopeReferenc
 
   @Children private final Node[] quickenedField;
 
+  @CompilationFinal private int contextLevel;
+
   private final int numLocals;
   private final int maxStackDepth;
 
   private final int frameOnStackMarkerIndex;
 
   public BytecodeLoopNode(final byte[] bytecodes, final int numLocals,
       final Object[] literals, final int maxStackDepth,
-      final int frameOnStackMarkerIndex, final BackJump[] inlinedLoops) {
+      final int frameOnStackMarkerIndex, final BackJump[] inlinedLoops, int contextLevel) {
     this.bytecodesField = bytecodes;
     this.numLocals = numLocals;
     this.literalsAndConstantsField = literals;
@@ -172,13 +174,19 @@ public BytecodeLoopNode(final byte[] bytecodes, final int numLocals,
     this.frameOnStackMarkerIndex = frameOnStackMarkerIndex;
 
     this.quickenedField = new Node[bytecodes.length];
+    this.contextLevel = contextLevel;
+  }
+
+  public int getContextLevel() {
+    return contextLevel;
   }
 
   @Override
   public Node deepCopy() {
     return new BytecodeLoopNode(
         bytecodesField.clone(), numLocals, literalsAndConstantsField,
-        maxStackDepth, frameOnStackMarkerIndex, inlinedLoopsField).initialize(sourceCoord);
+        maxStackDepth, frameOnStackMarkerIndex, inlinedLoopsField, contextLevel).initialize(
+            sourceCoord);
   }
 
   public String getNameOfLocal(final int idx) {
@@ -763,7 +771,7 @@ public Object executeGeneric(final VirtualFrame frame) {
 
           Object result = stack[stackPointer];
           // stackPointer -= 1;
-          doReturnNonLocal(frame, bytecodeIndex, result);
+          doReturnNonLocal(frame, result);
           return Nil.nilObject;
         }
 
@@ -1195,11 +1203,8 @@ private SClass getHolder() {
   }
 
   @InliningCutoff
-  private void doReturnNonLocal(final VirtualFrame frame, final int bytecodeIndex,
-      final Object result) {
-    byte contextIdx = bytecodesField[bytecodeIndex + 1];
-
-    MaterializedFrame ctx = determineContext(frame, contextIdx);
+  private void doReturnNonLocal(final VirtualFrame frame, final Object result) {
+    MaterializedFrame ctx = determineContext(frame, contextLevel);
     FrameOnStackMarker marker =
         (FrameOnStackMarker) ctx.getObject(frameOnStackMarkerIndex);
 
@@ -1265,6 +1270,9 @@ public void replaceAfterScopeChange(final ScopeAdaptationVisitor inliner) {
       }
     } else {
       boolean requiresChangesToContextLevels = inliner.outerScopeChanged();
+      if (requiresChangesToContextLevels) {
+        contextLevel -= 1;
+      }
       adapt(inliner, requiresChangesToContextLevels);
     }
   }
@@ -1532,12 +1540,10 @@ private void inlineInto(final BytecodeMethodGenContext mgenc,
         }
 
         case RETURN_NON_LOCAL: {
-          byte contextIdx = bytecodes[i + 1];
-          byte newCtx = (byte) (contextIdx - 1);
-          if (newCtx == 0) {
-            emitRETURNLOCAL(mgenc);
-          } else {
+          if (mgenc.isBlockMethod()) {
             emitRETURNNONLOCAL(mgenc);
+          } else {
+            emitRETURNLOCAL(mgenc);
           }
           break;
         }
@@ -1734,22 +1740,8 @@ private void adapt(final ScopeAdaptationVisitor inliner,
         case POP_FIELD_1:
         case SEND:
         case SUPER_SEND:
-        case RETURN_LOCAL: {
-          break;
-        }
-
-        case RETURN_NON_LOCAL: {
-          byte contextIdx = bytecodes[i + 1];
-          if (requiresChangesToContextLevels && contextIdx >= inliner.contextLevel) {
-            // we don't simplify to return local, because they had different bytecode length
-            // and, well, I don't think this should happen
-            assert contextIdx - 1 > 0 : "I wouldn't expect a RETURN_LOCAL equivalent here, "
-                + " because we are in a block, or it is already a return local";
-            bytecodes[i + 1] = (byte) (contextIdx - 1);
-          }
-          break;
-        }
-
+        case RETURN_LOCAL:
+        case RETURN_NON_LOCAL:
         case RETURN_SELF:
         case RETURN_FIELD_0:
         case RETURN_FIELD_1:

src/trufflesom/src/trufflesom/primitives/basics/BlockPrims.java

@@ -49,7 +49,6 @@ public SAbstractObject doSBlock(@SuppressWarnings("unused") final SBlock receive
     }
   }
 
-  @ReportPolymorphism
   @GenerateNodeFactory
   @GenerateWrapper
   @Primitive(className = "Block", primitive = "value")
@@ -100,7 +99,6 @@ public WrapperNode createWrapper(final ProbeNode probe) {
     }
   }
 
-  @ReportPolymorphism
   @GenerateWrapper
   @GenerateNodeFactory
   @Primitive(className = "Block2", primitive = "value:", selector = "value:", inParser = false,
@@ -146,7 +144,6 @@ public WrapperNode createWrapper(final ProbeNode probe) {
     }
   }
 
-  @ReportPolymorphism
   @GenerateNodeFactory
   @Primitive(className = "Block3", primitive = "value:with:", selector = "value:with:",
       inParser = false, receiverType = SBlock.class)

tests/trufflesom/tests/BytecodeBlockTests.java

@@ -192,12 +192,12 @@ private void blockIfReturnNonLocal(final String sel, final byte jumpBytecode) {
         + " #end\n"
         + "]");
 
-    assertEquals(17, bytecodes.length);
+    assertEquals(16, bytecodes.length);
     check(bytecodes,
         t(5, Bytecodes.SEND),
-        new BC(jumpBytecode, 6),
+        new BC(jumpBytecode, 5),
         Bytecodes.PUSH_ARG1,
-        new BC(Bytecodes.RETURN_NON_LOCAL, 1),
+        Bytecodes.RETURN_NON_LOCAL,
         Bytecodes.POP);
   }