8341275: Explicit counter support for Code attribute building

src/java.base/share/classes/java/lang/classfile/CodeBuilder.java

@@ -3416,4 +3416,32 @@ default CodeBuilder tableswitch(Label defaultTarget, List<SwitchCase> cases) {
         }
         return tableswitch(low, high, defaultTarget, cases);
     }
+
+    // advanced user hints
+
+    /**
+     * Provides this builder with explicit {@code max_stack} and {@code
+     * max_locals} values when {@link StackMapsOption#DROP_STACK_MAPS} is set.
+     * No automatic counting for the two max values will be done; the resulting
+     * {@code Code} attribute will instead use the provided values.
+     * <p>
+     * This call does not affect the two max values if any of the following is
+     * true:
+     * <ul>
+     * <li>{@code DROP_STACK_MAPS} is not set;
+     * <li>this builder does not build a complete {@code Code} attribute;
+     * <li>this builder's output is redirected to a {@link CodeTransform};
+     * <li>the built {@code Code} attribute is subsequently inflated in a
+     * transformation;
+     * </ul>
+     * <p>
+     * This call always validates the provided max values.
+     *
+     * @param maxStack the max stack slots used by this method body
+     * @param maxLocals the max local variable slots used by this method body
+     * @return this builder
+     * @throws IllegalArgumentException if any of the values is not a {@code u2}
+     * @since 25
+     */
+    CodeBuilder withExplicitStackAndLocals(int maxStack, int maxLocals);
 }

src/java.base/share/classes/java/lang/runtime/SwitchBootstraps.java

@@ -511,8 +511,8 @@ private static Consumer<CodeBuilder> generateTypeSwitchSkeleton(Class<?> selecto
             if (labelConstants.length == 0) {
                 cb.loadConstant(0)
                   .ireturn()
-                  .with(StackMapTableAttribute.of(stackMapFrames));
-                DirectCodeBuilder.withMaxs(cb, 2, locals.size()); // checkIndex uses 2
+                  .with(StackMapTableAttribute.of(stackMapFrames))
+                  .withExplicitStackAndLocals(2, locals.size()); // checkIndex uses 2
                 return;
             }
             cb.iload(RESTART_IDX);
@@ -712,8 +712,8 @@ private static Consumer<CodeBuilder> generateTypeSwitchSkeleton(Class<?> selecto
             cb.labelBinding(dflt)
               .loadConstant(labelConstants.length)
               .ireturn()
-              .with(StackMapTableAttribute.of(stackMapFrames));
-            DirectCodeBuilder.withMaxs(cb, 3, locals.size()); // enum labels use 3 stack, others use 2
+              .with(StackMapTableAttribute.of(stackMapFrames))
+              .withExplicitStackAndLocals(3, locals.size()); // enum labels use 3 stack, others use 2
         };
     }
 

src/java.base/share/classes/jdk/internal/classfile/impl/BufferedCodeBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -48,6 +48,8 @@ public final class BufferedCodeBuilder
     private final MethodInfo methodInfo;
     private boolean finished;
     private int maxLocals;
+    private int maxStackHint = -1;
+    private int maxLocalsHint = -1;
 
     public BufferedCodeBuilder(MethodInfo methodInfo,
                                SplitConstantPool constantPool,
@@ -126,6 +128,17 @@ public CodeBuilder with(CodeElement element) {
         return this;
     }
 
+    @Override
+    public CodeBuilder withExplicitStackAndLocals(int maxStack, int maxLocals) {
+        BytecodeHelpers.validateU2(maxStack);
+        BytecodeHelpers.validateU2(maxLocals);
+        if (context.dropStackMaps()) {
+            this.maxStackHint = maxStack;
+            this.maxLocalsHint = maxLocals;
+        }
+        return this;
+    }
+
     @Override
     public String toString() {
         return String.format("CodeModel[id=%d]", System.identityHashCode(this));
@@ -174,7 +187,17 @@ public Optional<MethodModel> parent() {
 
         @Override
         public void writeTo(DirectMethodBuilder builder) {
-            builder.withCode(Util.writingAll(this));
+            if (maxStackHint != -1 && maxLocalsHint != -1) {
+                builder.withCode(new Consumer<>() {
+                    @Override
+                    public void accept(CodeBuilder cob) {
+                        Model.this.forEach(cob);
+                        cob.withExplicitStackAndLocals(maxStackHint, maxLocalsHint);
+                    }
+                });
+            } else {
+                builder.withCode(Util.writingAll(this));
+            }
         }
 
         @Override

src/java.base/share/classes/jdk/internal/classfile/impl/BytecodeHelpers.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2025, Oracle and/or its affiliates. All rights reserved.
  * Copyright (c) 2024, Alibaba Group Holding Limited. All Rights Reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
@@ -54,8 +54,12 @@ public static IllegalArgumentException cannotConvertException(TypeKind from, Typ
         return new IllegalArgumentException(String.format("convert %s -> %s", from, to));
     }
 
+    public static IllegalArgumentException u2OutOfBounds(int u2) {
+        return new IllegalArgumentException("Invalid value for u2:".concat(Integer.toString(u2)));
+    }
+
     public static IllegalArgumentException slotOutOfBounds(int slot) {
-        return new IllegalArgumentException("Invalid slot index :".concat(Integer.toString(slot)));
+        return new IllegalArgumentException("Invalid slot index:".concat(Integer.toString(slot)));
     }
 
     public static IllegalArgumentException slotOutOfBounds(Opcode opcode, int slot) {
@@ -437,6 +441,11 @@ public static void validateSlot(int slot) {
             throw slotOutOfBounds(slot);
     }
 
+    public static void validateU2(int u2) {
+        if ((u2 & ~0xFFFF) != 0)
+            throw u2OutOfBounds(u2);
+    }
+
     public static boolean validateAndIsWideIinc(int slot, int val) {
         var ret = false;
         if ((slot & ~0xFF) != 0) {

src/java.base/share/classes/jdk/internal/classfile/impl/DirectCodeBuilder.java

@@ -174,10 +174,15 @@ public MethodInfo methodInfo() {
         return methodInfo;
     }
 
-    public static void withMaxs(CodeBuilder cob, int stacks, int locals) {
-        var dcb = (DirectCodeBuilder) cob;
-        dcb.maxStackHint = stacks;
-        dcb.maxLocalsHint = locals;
+    @Override
+    public CodeBuilder withExplicitStackAndLocals(int maxStack, int maxLocals) {
+        BytecodeHelpers.validateU2(maxStack);
+        BytecodeHelpers.validateU2(maxLocals);
+        if (context.dropStackMaps()) {
+            this.maxStackHint = maxStack;
+            this.maxLocalsHint = maxLocals;
+        }
+        return this;
     }
 
     private UnboundAttribute<CodeAttribute> content = null;

src/java.base/share/classes/jdk/internal/classfile/impl/NonterminalCodeBuilder.java

@@ -62,4 +62,11 @@ public ConstantPoolBuilder constantPool() {
     public Label newLabel() {
         return terminal.newLabel();
     }
+
+    @Override
+    public CodeBuilder withExplicitStackAndLocals(int maxStack, int maxLocals) {
+        BytecodeHelpers.validateU2(maxStack);
+        BytecodeHelpers.validateU2(maxLocals);
+        return this;
+    }
 }

test/jdk/jdk/classfile/BuilderExplicitMaxsTest.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+/*
+ * @test
+ * @bug 8341275
+ * @summary Testing CodeBuilder::withExplicitStackAndLocals
+ * @run junit BuilderExplicitMaxsTest
+ */
+
+import java.lang.classfile.*;
+import java.lang.constant.ClassDesc;
+import java.util.Arrays;
+import java.util.stream.Stream;
+
+import helpers.CodeBuilderType;
+import jdk.internal.classfile.impl.BufferedCodeBuilder;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import static java.lang.constant.ConstantDescs.MTD_void;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+/**
+ * Testing explicit max stacks and locals setter. Configurations:
+ * <ul>
+ * <li>Builder Type: Direct, Buffered, Block, Chained
+ * <li>Argument validity
+ * <li>Whether DROP_STACK_MAPS is set
+ * </ul>
+ */
+class BuilderExplicitMaxsTest {
+
+    static Stream<Arguments> arguments() {
+        return Arrays.stream(CodeBuilderType.values()).mapMulti((t, sink) -> {
+            sink.accept(Arguments.of(ClassFile.StackMapsOption.DROP_STACK_MAPS, t));
+            sink.accept(Arguments.of(ClassFile.StackMapsOption.STACK_MAPS_WHEN_REQUIRED, t));
+        });
+    }
+
+    @MethodSource("arguments")
+    @ParameterizedTest
+    void testValidArgs(ClassFile.StackMapsOption stackMapsOption, CodeBuilderType builderType) {
+        var cc = ClassFile.of(stackMapsOption);
+        var bytes = cc.build(ClassDesc.of("Foo"),
+                             builderType.asClassHandler("foo", MTD_void, 0, cob ->
+                                     cob.return_()
+                                        .withExplicitStackAndLocals(2, 3)));
+        var clz = ClassFile.of().parse(bytes);
+        var code = clz.methods().getFirst().findAttribute(Attributes.code()).orElseThrow();
+        if (builderType.terminal && stackMapsOption == ClassFile.StackMapsOption.DROP_STACK_MAPS) {
+            assertEquals(2, code.maxStack());
+            assertEquals(3, code.maxLocals());
+        } else {
+            assertEquals(0, code.maxStack());
+            assertEquals(1, code.maxLocals());
+        }
+    }
+
+    @MethodSource("arguments")
+    @ParameterizedTest
+    void testInvalidArgs(ClassFile.StackMapsOption stackMapsOption, CodeBuilderType builderType) {
+        var cc = ClassFile.of(stackMapsOption);
+        assertThrows(IllegalArgumentException.class, () ->
+                cc.build(ClassDesc.of("Foo"),
+                         builderType.asClassHandler("foo", MTD_void, 0, cob ->
+                                 cob.return_()
+                                    .withExplicitStackAndLocals(-1, 2))));
+        assertThrows(IllegalArgumentException.class, () ->
+                cc.build(ClassDesc.of("Foo"),
+                         builderType.asClassHandler("foo", MTD_void, 0, cob ->
+                                 cob.return_()
+                                    .withExplicitStackAndLocals(2, 100000))));
+    }
+}

test/jdk/jdk/classfile/BuilderParamTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -30,37 +30,36 @@
 import java.lang.constant.MethodTypeDesc;
 
 import java.lang.classfile.ClassFile;
-import org.junit.jupiter.api.Test;
 
-import static java.lang.constant.ConstantDescs.CD_void;
+import helpers.CodeBuilderType;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
 import static java.lang.classfile.ClassFile.ACC_STATIC;
 import static org.junit.jupiter.api.Assertions.*;
 
 /**
  * BuilderParamTest
  */
 class BuilderParamTest {
-    @Test
-    void testDirectBuilder() {
+    @EnumSource
+    @ParameterizedTest
+    void test(CodeBuilderType type) {
         var cc = ClassFile.of();
-        cc.build(ClassDesc.of("Foo"), cb -> {
-            cb.withMethod("foo", MethodTypeDesc.ofDescriptor("(IJI)V"), 0,
-                          mb -> mb.withCode(xb -> {
-                assertEquals(xb.receiverSlot(), 0);
-                assertEquals(xb.parameterSlot(0), 1);
-                assertEquals(xb.parameterSlot(1), 2);
-                assertEquals(xb.parameterSlot(2), 4);
-                xb.return_();
-            }));
-        });
-        cc.build(ClassDesc.of("Foo"), cb -> {
-            cb.withMethod("foo", MethodTypeDesc.ofDescriptor("(IJI)V"), ACC_STATIC,
-                          mb -> mb.withCode(xb -> {
-                              assertEquals(xb.parameterSlot(0), 0);
-                              assertEquals(xb.parameterSlot(1), 1);
-                              assertEquals(xb.parameterSlot(2), 3);
-                              xb.return_();
-                          }));
-        });
+        cc.build(ClassDesc.of("Foo"),
+                 type.asClassHandler("foo", MethodTypeDesc.ofDescriptor("(IJI)V"), 0, xb -> {
+                     assertEquals(0, xb.receiverSlot(), "this");
+                     assertEquals(1, xb.parameterSlot(0), "int");
+                     assertEquals(2, xb.parameterSlot(1), "long");
+                     assertEquals(4, xb.parameterSlot(2), "int");
+                     xb.return_();
+                 }));
+        cc.build(ClassDesc.of("Foo"),
+                 type.asClassHandler("foo", MethodTypeDesc.ofDescriptor("(IJI)V"), ACC_STATIC, xb -> {
+                     assertEquals(0, xb.parameterSlot(0), "int");
+                     assertEquals(1, xb.parameterSlot(1), "long");
+                     assertEquals(3, xb.parameterSlot(2), "int");
+                     xb.return_();
+                 }));
     }
 }

test/jdk/jdk/classfile/InstructionValidationTest.java

@@ -33,12 +33,12 @@
 import java.lang.classfile.constantpool.ConstantPoolBuilder;
 import java.lang.classfile.instruction.*;
 import java.lang.constant.ClassDesc;
-import java.lang.reflect.Parameter;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.ObjIntConsumer;
 import java.util.stream.Stream;
 
+import helpers.CodeBuilderType;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.function.Executable;
 
@@ -168,21 +168,10 @@ static void ensureFailFast(int value, Consumer<CodeBuilder> action) {
             action.accept(cob);
             fail("Bad slot access did not fail fast: "  + value);
         };
-        var cf = ClassFile.of();
-        CodeTransform noopCodeTransform = CodeBuilder::with;
-        // Direct builders
-        assertThrows(IllegalArgumentException.class, () -> cf.build(ClassDesc.of("Test"), clb -> clb
-                .withMethodBody("test", MTD_void, ACC_STATIC, checkedAction)));
-        // Chained builders
-        assertThrows(IllegalArgumentException.class, () -> cf.build(ClassDesc.of("Test"), clb -> clb
-                .withMethodBody("test", MTD_void, ACC_STATIC, cob -> cob
-                        .transforming(CodeTransform.ACCEPT_ALL, checkedAction))));
-        var classTemplate = cf.build(ClassDesc.of("Test"), clb -> clb
-                .withMethodBody("test", MTD_void, ACC_STATIC, CodeBuilder::return_));
-        // Indirect builders
-        assertThrows(IllegalArgumentException.class, () -> cf.transformClass(cf.parse(classTemplate),
-                ClassTransform.transformingMethodBodies(CodeTransform.endHandler(checkedAction)
-                        .andThen(noopCodeTransform))));
+        for (var builderType : CodeBuilderType.values()) {
+            assertThrows(IllegalArgumentException.class, () -> ClassFile.of().build(ClassDesc.of("Test"),
+                    builderType.asClassHandler("test", MTD_void, ACC_STATIC, checkedAction)), builderType.name());
+        }
     }
 
     static void check(boolean fails, Executable exec) {