GR-73325: Preserve template state on handler return

Author: mur47x111

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/api/directives/BytecodeInterpreterDirectives.java

@@ -60,15 +60,6 @@ private BytecodeInterpreterDirectives() {
          */
         boolean threading() default true;
 
-        /**
-         * Indicates whether this handler can be dispatched while a template variable is in a
-         * non-zero state. If {@code false}, non-zero template handler tables dispatch this handler's
-         * opcodes to a generated fallback stub so control returns to the interpreter switch before
-         * the handler executes. The default is {@code true}, meaning the handler is compatible with
-         * all template states.
-         */
-        boolean templateCompatible() default true;
-
         /**
          * Indicates whether execution of this handler should include a safepoint check.
          */
@@ -132,21 +123,6 @@ enum ExpansionKind {
                  */
                 boolean nonNull() default true;
 
-                /**
-                 * Marks this field as scratch state when template mode is enabled. Scratch fields
-                 * are carried between threaded bytecode handler stubs, but they are not initialized
-                 * from the original Java object on entry, are not written back to the original Java
-                 * object on exit, and are not preserved through pending exception state. This is
-                 * intended for register-resident interpreter state that is valid only while threaded
-                 * execution remains inside the generated handler stubs. When template mode is not
-                 * enabled, this metadata is ignored and the field remains an ordinary expanded
-                 * argument.
-                 * <p>
-                 * This property is only supported for fields of {@link ExpansionKind#VIRTUAL}
-                 * arguments. It cannot be combined with {@link #templateVariable()}.
-                 */
-                boolean scratch() default false;
-
                 /**
                  * Marks the expanded field as the template variable used to specialize template
                  * variants and select the template variant of the next threaded bytecode handler.
@@ -238,5 +214,4 @@ enum ExpansionKind {
     @Target({ElementType.METHOD})
     public @interface BytecodeInterpreterFetchOpcode {
     }
-
 }

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/hotspot/stubs/HotSpotTruffleBytecodeHandlerStub.java

@@ -64,7 +64,7 @@ protected StructuredGraph getGraph(DebugContext debug, CompilationIdentifier com
         try {
             HotSpotGraphKit kit = new HotSpotGraphKit(debug, callsite.getEnclosingMethod(), providers, providers.getGraphBuilderPlugins(), compilationId, callsite.getStubName(), false, true);
             return BytecodeHandlerStubHelper.createStub(kit, callsite.getEnclosingMethod(), callsite.getBci(), false, null, null, callsite.getHandlerConfig(), callsite.getTargetMethod(),
-                            0, null);
+                            0, null, null);
         } catch (Exception e) {
             throw GraalError.shouldNotReachHere(e); // ExcludeFromJacocoGeneratedReport
         }

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/phases/util/BytecodeHandlerCallSite.java

@@ -30,7 +30,6 @@
 import java.util.function.Function;
 
 import jdk.graal.compiler.debug.GraalError;
-import jdk.graal.compiler.nodes.ConstantNode;
 import jdk.graal.compiler.nodes.FixedGuardNode;
 import jdk.graal.compiler.nodes.FixedNode;
 import jdk.graal.compiler.nodes.LogicNode;
@@ -118,24 +117,18 @@ public String getStubName() {
 
     /**
      * Constructs the stub ABI argument list at the caller. Expanded arguments are lowered to field
-     * loads from their Java owner objects, except scratch fields which are initialized with a
-     * default value because their Java object field value is intentionally ignored. Non-expanded
-     * arguments are forwarded unchanged.
+     * loads from their Java owner objects; non-expanded arguments are forwarded unchanged.
      */
     public ValueNode[] createCallerArguments(ValueNode[] oldArguments, FixedNode insertBefore, Function<ResolvedJavaField, ResolvedJavaField> fieldMap) {
         List<ArgumentInfo> stubAbiArgumentInfos = handlerConfig.getStubAbiArgumentInfos();
         List<ValueNode> newArguments = new ArrayList<>();
         StructuredGraph graph = insertBefore.graph();
         for (ArgumentInfo argumentInfo : stubAbiArgumentInfos) {
             if (argumentInfo.isExpanded()) {
-                if (argumentInfo.scratch()) {
-                    newArguments.add(ConstantNode.defaultForKind(argumentInfo.type().getJavaKind(), graph));
-                } else {
-                    ValueNode owner = oldArguments[argumentInfo.originalIndex()];
-                    LoadFieldNode load = LoadFieldNode.create(graph.getAssumptions(), owner, fieldMap.apply(argumentInfo.field()));
-                    graph.addBeforeFixed(insertBefore, graph.add(load));
-                    newArguments.add(load);
-                }
+                ValueNode owner = oldArguments[argumentInfo.originalIndex()];
+                LoadFieldNode load = LoadFieldNode.create(graph.getAssumptions(), owner, fieldMap.apply(argumentInfo.field()));
+                graph.addBeforeFixed(insertBefore, graph.add(load));
+                newArguments.add(load);
             } else {
                 newArguments.add(oldArguments[argumentInfo.originalIndex()]);
             }
@@ -159,7 +152,7 @@ public void updateCallerReturns(FixedNode newInvoke, ValueNode[] oldArguments, F
         List<ArgumentInfo> stubAbiArgumentInfos = handlerConfig.getStubAbiArgumentInfos();
 
         for (ArgumentInfo argumentInfo : stubAbiArgumentInfos) {
-            if (argumentInfo.isExpanded() && !argumentInfo.isImmutable() && !argumentInfo.scratch()) {
+            if (argumentInfo.isExpanded() && !argumentInfo.isImmutable()) {
                 ReadArgumentNode fetchReturn = graph.unique(new ReadArgumentNode(newInvoke,
                                 argumentInfo.type().getJavaKind(), argumentInfo.index()));
                 ValueNode owner = oldArguments[argumentInfo.originalIndex()];

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/phases/util/BytecodeHandlerConfig.java

@@ -42,7 +42,7 @@
  * Aggregates metadata derived from a {@code BytecodeInterpreterHandlerConfig} annotation which is
  * resolved against a concrete bytecode handler signature. The resulting model describes the stub
  * ABI: argument expansion, non-null guarantees, {@code returnValue}/copy-from-return slots, mutable
- * expanded fields, scratch fields, and the maximum opcode for which a handler exists.
+ * expanded fields and the maximum opcode for which a handler exists.
  */
 public final class BytecodeHandlerConfig {
 
@@ -184,11 +184,11 @@ private static int appendReceiver(List<ArgumentInfo> arguments, AnnotationValue
 
         switch (expansionKind) {
             case NONE -> {
-                arguments.add(new ArgumentInfo(declaringClass, nextIndex++, originalIndex, false, false, false, null, null, false, true, nonNull, false, 0));
+                arguments.add(new ArgumentInfo(declaringClass, nextIndex++, originalIndex, false, false, false, null, null, false, true, nonNull, 0));
             }
             case VIRTUAL -> throw GraalError.shouldNotReachHere("Receiver cannot be VIRTUAL");
             case MATERIALIZED -> {
-                arguments.add(new ArgumentInfo(declaringClass, nextIndex++, originalIndex, false, true, false, null, null, false, true, nonNull, false, 0));
+                arguments.add(new ArgumentInfo(declaringClass, nextIndex++, originalIndex, false, true, false, null, null, false, true, nonNull, 0));
                 nextIndex = appendMaterializedFields(arguments, receiverConfig, declaringClass, declaringClass, originalIndex, nextIndex, templateModeEnabled);
             }
             default -> throw GraalError.shouldNotReachHere("Unknown expansion kind " + expansionKind);
@@ -205,30 +205,28 @@ private static int appendParameter(List<ArgumentInfo> arguments, AnnotationValue
 
         switch (expansionKind) {
             case NONE -> {
-                arguments.add(new ArgumentInfo(parameterType, nextIndex++, originalIndex, copyFromReturn, false, false, null, null, false, !copyFromReturn, nonNull, false, 0));
+                arguments.add(new ArgumentInfo(parameterType, nextIndex++, originalIndex, copyFromReturn, false, false, null, null, false, !copyFromReturn, nonNull, 0));
             }
             case VIRTUAL -> {
                 List<AnnotationValue> fields = parameterConfig.getList("fields", AnnotationValue.class);
                 for (ResolvedJavaField javaField : parameterType.getInstanceFields(true)) {
                     ResolvedJavaType fieldType = javaField.getType().resolve(declaringClass);
                     boolean fieldNonNull = false;
                     AnnotationValue fieldConfig = findFieldConfig(fields, javaField.getName());
-                    boolean scratch = templateModeEnabled && fieldConfig != null && fieldConfig.getBoolean("scratch");
                     int templateVariants = templateModeEnabled ? getTemplateVariants(fieldConfig, javaField) : 0;
-                    GraalError.guarantee(!scratch || templateVariants == 0, "Scratch field %s cannot be a template variable", javaField.format("%H.%n"));
                     if (!fieldType.isPrimitive()) {
-                        fieldNonNull = !scratch && fieldConfig != null && fieldConfig.getBoolean("nonNull");
+                        fieldNonNull = fieldConfig != null && fieldConfig.getBoolean("nonNull");
                         GraalError.guarantee(templateVariants == 0, "Field %s is marked as a template variable", javaField.format("%H.%n"));
                     } else if (templateVariants > 0) {
                         GraalError.guarantee(fieldType.getJavaKind() == JavaKind.Int, "Template variable field %s must be int", javaField.format("%H.%n"));
                     }
                     int abiIndex = templateVariants > 0 ? -1 : nextIndex++;
                     arguments.add(new ArgumentInfo(fieldType, abiIndex, originalIndex, false, false, true, parameterType, javaField, true, javaField.isFinal(),
-                                    fieldNonNull, scratch, templateVariants));
+                                    fieldNonNull, templateVariants));
                 }
             }
             case MATERIALIZED -> {
-                arguments.add(new ArgumentInfo(parameterType, nextIndex++, originalIndex, copyFromReturn, true, false, null, null, false, true, nonNull, false, 0));
+                arguments.add(new ArgumentInfo(parameterType, nextIndex++, originalIndex, copyFromReturn, true, false, null, null, false, true, nonNull, 0));
                 nextIndex = appendMaterializedFields(arguments, parameterConfig, parameterType, declaringClass, originalIndex, nextIndex, templateModeEnabled);
             }
             default -> throw GraalError.shouldNotReachHere("Unknown expansion kind " + expansionKind);
@@ -243,16 +241,13 @@ private static int appendMaterializedFields(List<ArgumentInfo> arguments, Annota
         for (ResolvedJavaField javaField : expandedType.getInstanceFields(true)) {
             AnnotationValue fieldConfig = findFieldConfig(fields, javaField.getName());
             if (fieldConfig != null) {
-                if (templateModeEnabled) {
-                    GraalError.guarantee(!fieldConfig.getBoolean("scratch"), "Scratch field %s must belong to a VIRTUAL argument", javaField.format("%H.%n"));
-                }
                 ResolvedJavaType fieldType = javaField.getType().resolve(declaringClass);
                 boolean fieldNonNull = !fieldType.isPrimitive() && fieldConfig.getBoolean("nonNull");
                 if (templateModeEnabled) {
                     GraalError.guarantee(getTemplateVariants(fieldConfig, javaField) == 0, "Field %s is marked as a template variable", javaField.format("%H.%n"));
                 }
                 arguments.add(new ArgumentInfo(fieldType, nextIndex++, originalIndex, false, false, true, declaringClass, javaField, false, javaField.isFinal(),
-                                fieldNonNull, false, 0));
+                                fieldNonNull, 0));
             }
         }
         return nextIndex;
@@ -336,6 +331,22 @@ public List<ResolvedJavaType> getStubAbiArgumentTypes() {
         return stubAbiArgumentTypes;
     }
 
+    /**
+     * Returns the number of thread-local slots required to preserve ordinary ABI arguments and
+     * template variables when a handler chain returns to Java.
+     */
+    public int getPendingStateSlotCount() {
+        return stubAbiArgumentInfos.size() + templateVariableArguments.size();
+    }
+
+    /**
+     * Returns the thread-local slot assigned to a template variable. Template slots follow the
+     * dense stub ABI slots and do not form part of the stub calling convention.
+     */
+    public int getTemplateVariablePendingStateSlot(int templateVariableIndex) {
+        return stubAbiArgumentInfos.size() + templateVariableIndex;
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (this == obj) {
@@ -397,10 +408,6 @@ private enum ExpansionKind {
      * produce multiple {@link ArgumentInfo}s that share the same {@link #originalIndex()} but map to
      * different stub ABI {@link #index()} values. In template mode, template variables have index
      * {@code -1} and are excluded from the stub ABI.
-     * <p>
-     * When template mode is enabled, scratch fields are still ABI arguments so threaded handlers
-     * can pass them to each other, but they are initialized with default values at the Java caller
-     * boundary and are not written back to the original owner object.
      */
     public record ArgumentInfo(ResolvedJavaType type,
                     int index,
@@ -413,7 +420,6 @@ public record ArgumentInfo(ResolvedJavaType type,
                     boolean isOwnerVirtual,
                     boolean isImmutable,
                     boolean nonNull,
-                    boolean scratch,
                     int templateVariants) {
         public boolean isTemplateVariable() {
             return templateVariants > 0;
@@ -424,7 +430,7 @@ public boolean isTemplateVariable() {
          * edge: the {@code copyFromReturn} value and mutable fields of virtual-expanded arguments.
          */
         public boolean needsPendingExceptionState() {
-            return copyFromReturn || (isExpanded && isOwnerVirtual && !isImmutable && !scratch);
+            return copyFromReturn || (isExpanded && isOwnerVirtual && !isImmutable);
         }
     }
 }