Bytecode DSL: fix childBci calculation for yield operations

Author: DSouzaM

truffle/src/com.oracle.truffle.api.bytecode.test/src/com/oracle/truffle/api/bytecode/test/CustomYieldTest.java

@@ -652,6 +652,40 @@ public void testImplicitTagInstrumentation() {
         });
     }
 
+    @Test
+    public void testYieldQuickeningRegressionTest() {
+        /*
+         * Regression test for a quickening bug. The yield's childBci calculation did not account
+         * for tag.resume instructions, and under the right circumstances, the parent operation
+         * would "quicken" an instruction operand, leading to unexpected results.
+         */
+        runInstrumentationTest((context, instrumenter) -> {
+            BytecodeRootNodes<ComplexCustomYieldTestRootNode> nodes = ComplexCustomYieldTestRootNodeGen.create(BytecodeDSLTestLanguage.REF.get(null), BytecodeConfig.DEFAULT, b -> {
+                b.beginRoot();
+                b.emitLoadConstant(42);
+                b.beginReturn();
+                b.beginAddConstantsYield(1);
+                b.emitLoadArgument(0);
+                b.endAddConstantsYield(1);
+                b.endReturn();
+                b.endRoot();
+            });
+            ComplexCustomYieldTestRootNode root = nodes.getNode(0);
+            root.getBytecodeNode().setUncachedThreshold(0);
+            List<Object> yieldValues = new ArrayList<>();
+            AtomicInteger resumeCount = new AtomicInteger();
+            instrumenter.attachExecutionEventFactory(SourceSectionFilter.newBuilder().tagIs(StatementTag.class).build(), createFactory(yieldValues, resumeCount));
+
+            ContinuationResult cont = (ContinuationResult) root.getCallTarget().call(0);
+            assertEquals(2, cont.getResult());
+            assertEquals(123, cont.continueWith(123));
+
+            cont = (ContinuationResult) root.getCallTarget().call(0);
+            assertEquals(2, cont.getResult());
+            assertEquals(123, cont.continueWith(123));
+        });
+    }
+
     private static void runInstrumentationTest(BiConsumer<Context, Instrumenter> test) {
         Context context = Context.create(BytecodeDSLTestLanguage.ID);
         try {

truffle/src/com.oracle.truffle.dsl.processor/src/com/oracle/truffle/dsl/processor/bytecode/generator/BytecodeRootNodeElement.java

@@ -4797,7 +4797,7 @@ private CodeExecutableElement createEnd(OperationModel operation) {
                     buildEmitOperationInstruction(b, operation, constantOperandIndices);
 
                     if (model.enableTagInstrumentation) {
-                        b.statement("doEmitTagResume()");
+                        b.declaration(type(int.class), "tagResumeBci", "doEmitTagResume()");
                     }
                     break;
                 default:
@@ -4847,6 +4847,10 @@ private CodeExecutableElement createEnd(OperationModel operation) {
                 b.end();
 
                 emitCallAfterChild(b, operation, "true", "nextBci");
+            } else if (model.enableTagInstrumentation && (operation.kind == OperationKind.YIELD || operation.kind == OperationKind.CUSTOM_YIELD)) {
+                // The "childBci" can change depending on whether tag.resume was emitted.
+                // We don't BE yields/tag.resume but the BE machinery needs a valid bci.
+                emitCallAfterChild(b, operation, "true", "tagResumeBci != -1 ? tagResumeBci : state.bci - " + operation.instruction.getInstructionLength());
             } else {
                 String nextBci;
                 if (operation.instruction != null) {
@@ -7109,24 +7113,27 @@ private CodeExecutableElement createDoEmitTagResume() {
                 throw new AssertionError("cannot produce method");
             }
 
-            CodeExecutableElement ex = new CodeExecutableElement(Set.of(PRIVATE), type(void.class), "doEmitTagResume");
+            CodeExecutableElement ex = new CodeExecutableElement(Set.of(PRIVATE), type(int.class), "doEmitTagResume");
             CodeTreeBuilder b = ex.createBuilder();
+            b.declaration(type(int.class), "tagResumeBci", "-1");
 
             b.startIf().string("tags == 0").end().startBlock();
-            b.returnDefault();
+            b.startReturn().string("tagResumeBci").end();
             b.end();
 
             buildOperationStackWalkFromBottom(b, "state.rootOperationSp", () -> {
                 b.startSwitch().string("operation.operation").end().startBlock();
                 OperationModel op = model.findOperation(OperationKind.TAG);
                 b.startCase().tree(createOperationConstant(op)).end();
                 b.startBlock();
+                b.startAssign("tagResumeBci").string("state.bci").end();
                 buildEmitInstruction(b, model.tagResumeInstruction, operationStack.read(op, operationFields.nodeId));
                 b.statement("break");
                 b.end(); // case tag
 
                 b.end(); // switch
             });
+            b.startReturn().string("tagResumeBci").end();
 
             return ex;
         }