8337225: Demote maxStack and maxLocals from CodeModel to CodeAttribute

Reviewed-by: asotona

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

@@ -43,17 +43,7 @@ import jdk.internal.javac.PreviewFeature;
 @PreviewFeature(feature = PreviewFeature.Feature.CLASSFILE_API)
 public sealed interface CodeModel
         extends CompoundElement<CodeElement>, AttributedElement, MethodElement
-        permits CodeAttribute, BufferedCodeBuilder.Model, CodeImpl {
+        permits CodeAttribute, BufferedCodeBuilder.Model {
-
-    /**
-     * {@return the maximum size of the local variable table}
-     */
-    int maxLocals();
-
-    /**
-     * {@return the maximum size of the operand stack}
-     */
-    int maxStack();
 
     /**
      * {@return the enclosing method, if known}

src/java.base/share/classes/java/lang/classfile/attribute/CodeAttribute.java

@@ -47,6 +47,16 @@ import jdk.internal.javac.PreviewFeature;
 public sealed interface CodeAttribute extends Attribute<CodeAttribute>, CodeModel
         permits BoundAttribute.BoundCodeAttribute {
 
+    /**
+     * {@return the maximum size of the local variable table}
+     */
+    int maxLocals();
+
+    /**
+     * {@return the maximum size of the operand stack}
+     */
+    int maxStack();
+
     /**
      * {@return The length of the code array in bytes}
      */

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

@@ -35,7 +35,7 @@ import java.lang.classfile.CompoundElement;
 
 public abstract sealed class AbstractUnboundModel<E extends ClassFileElement>
         extends AbstractElement
-        implements CompoundElement<E>, AttributedElement, Util.Writable
+        implements CompoundElement<E>, AttributedElement
         permits BufferedCodeBuilder.Model, BufferedFieldBuilder.Model, BufferedMethodBuilder.Model {
     private final List<E> elements;
     private List<Attribute<?>> attributes;
@@ -63,8 +63,11 @@ public abstract sealed class AbstractUnboundModel<E extends ClassFileElement>
     public List<Attribute<?>> attributes() {
         if (attributes == null)
             attributes = elements.stream()
-                                 .filter(e -> e instanceof Attribute)
+                                 .<Attribute<?>>mapMulti((e, sink) -> {
-                                 .<Attribute<?>>map(e -> (Attribute<?>) e)
+                                     if (e instanceof Attribute<?> attr) {
+                                         sink.accept(attr);
+                                     }
+                                 })
                                  .toList();
         return attributes;
     }

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

@@ -32,9 +32,6 @@ import java.lang.classfile.constantpool.ConstantPoolBuilder;
 import java.lang.classfile.Label;
 import java.lang.classfile.MethodModel;
 import java.lang.classfile.instruction.ExceptionCatch;
-import java.lang.classfile.instruction.IncrementInstruction;
-import java.lang.classfile.instruction.LoadInstruction;
-import java.lang.classfile.instruction.StoreInstruction;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -47,7 +44,6 @@ public final class BufferedCodeBuilder
     private final ClassFileImpl context;
     private final List<CodeElement> elements = new ArrayList<>();
     private final LabelImpl startLabel, endLabel;
-    private final CodeModel original;
     private final MethodInfo methodInfo;
     private boolean finished;
     private int maxLocals;
@@ -60,12 +56,8 @@ public final class BufferedCodeBuilder
         this.context = context;
         this.startLabel = new LabelImpl(this, -1);
         this.endLabel = new LabelImpl(this, -1);
-        this.original = original;
         this.methodInfo = methodInfo;
-        this.maxLocals = Util.maxLocals(methodInfo.methodFlags(), methodInfo.methodTypeSymbol());
+        this.maxLocals = TerminalCodeBuilder.setupTopLocal(methodInfo, original);
-        if (original != null)
-            this.maxLocals = Math.max(this.maxLocals, original.maxLocals());
-
         elements.add(startLabel);
     }
 
@@ -162,27 +154,16 @@ public final class BufferedCodeBuilder
         @Override
         public List<ExceptionCatch> exceptionHandlers() {
             return elements.stream()
-                           .filter(x -> x instanceof ExceptionCatch)
+                           .<ExceptionCatch>mapMulti((x, sink) -> {
-                           .map(x -> (ExceptionCatch) x)
+                               if (x instanceof ExceptionCatch ec) {
+                                   sink.accept(ec);
+                               }
+                           })
                            .toList();
         }
 
-        @Override
+        int curTopLocal() {
-        public int maxLocals() {
+            return BufferedCodeBuilder.this.curTopLocal();
-            for (CodeElement element : elements) {
-                if (element instanceof LoadInstruction i)
-                    maxLocals = Math.max(maxLocals, i.slot() + i.typeKind().slotSize());
-                else if (element instanceof StoreInstruction i)
-                    maxLocals = Math.max(maxLocals, i.slot() + i.typeKind().slotSize());
-                else if (element instanceof IncrementInstruction i)
-                    maxLocals = Math.max(maxLocals, i.slot() + 1);
-            }
-            return maxLocals;
-        }
-
-        @Override
-        public int maxStack() {
-            throw new UnsupportedOperationException("nyi");
         }
 
         @Override
@@ -200,11 +181,6 @@ public final class BufferedCodeBuilder
             });
         }
 
-        @Override
-        public void writeTo(BufWriterImpl buf) {
-            DirectCodeBuilder.build(methodInfo, cb -> elements.forEach(cb), constantPool, context, null).writeTo(buf);
-        }
-
         @Override
         public String toString() {
             return String.format("CodeModel[id=%s]", Integer.toHexString(System.identityHashCode(this)));

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

@@ -103,7 +103,7 @@ public final class BufferedFieldBuilder
 
         @Override
         public void writeTo(DirectClassBuilder builder) {
-            builder.withField(name, desc, new Consumer<FieldBuilder>() {
+            builder.withField(name, desc, new Consumer<>() {
                 @Override
                 public void accept(FieldBuilder fieldBuilder) {
                     elements.forEach(fieldBuilder);
@@ -111,13 +111,6 @@ public final class BufferedFieldBuilder
             });
         }
 
-        @Override
-        public void writeTo(BufWriterImpl buf) {
-            DirectFieldBuilder fb = new DirectFieldBuilder(constantPool, context, name, desc, null);
-            elements.forEach(fb);
-            fb.writeTo(buf);
-        }
-
         @Override
         public String toString() {
             return String.format("FieldModel[fieldName=%s, fieldType=%s, flags=%d]", name.stringValue(), desc.stringValue(), flags.flagsMask());

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

@@ -196,7 +196,11 @@ public final class BufferedMethodBuilder
 
         @Override
         public Optional<CodeModel> code() {
-            throw new UnsupportedOperationException("nyi");
+            return elements.stream().<CodeModel>mapMulti((e, sink) -> {
+                if (e instanceof CodeModel cm) {
+                    sink.accept(cm);
+                }
+            }).findFirst();
         }
 
         @Override
@@ -209,13 +213,6 @@ public final class BufferedMethodBuilder
             });
         }
 
-        @Override
-        public void writeTo(BufWriterImpl buf) {
-            DirectMethodBuilder mb = new DirectMethodBuilder(constantPool, context, name, desc, methodFlags(), null);
-            elements.forEach(mb);
-            mb.writeTo(buf);
-        }
-
         @Override
         public String toString() {
             return String.format("MethodModel[methodName=%s, methodType=%s, flags=%d]",

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

@@ -43,7 +43,7 @@ import static java.lang.classfile.ClassFile.*;
 
 public final class CodeImpl
         extends BoundAttribute.BoundCodeAttribute
-        implements CodeModel, LabelContext {
+        implements LabelContext {
 
     static final Instruction[] SINGLETON_INSTRUCTIONS = new Instruction[256];
 

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

@@ -127,12 +127,10 @@ public final class DirectCodeBuilder
         this.transformFwdJumps = transformFwdJumps;
         this.transformBackJumps = context.shortJumpsOption() == ClassFile.ShortJumpsOption.FIX_SHORT_JUMPS;
         bytecodesBufWriter = (original instanceof CodeImpl cai) ? new BufWriterImpl(constantPool, context, cai.codeLength())
-                                                               : new BufWriterImpl(constantPool, context);
+                : new BufWriterImpl(constantPool, context);
         this.startLabel = new LabelImpl(this, 0);
         this.endLabel = new LabelImpl(this, -1);
-        this.topLocal = Util.maxLocals(methodInfo.methodFlags(), methodInfo.methodTypeSymbol());
+        this.topLocal = TerminalCodeBuilder.setupTopLocal(methodInfo, original);
-        if (original != null)
-            this.topLocal = Math.max(this.topLocal, original.maxLocals());
     }
 
     @Override
@@ -312,8 +310,9 @@ public final class DirectCodeBuilder
 
             private void writeCounters(boolean codeMatch, BufWriterImpl buf) {
                 if (codeMatch) {
-                    buf.writeU2(original.maxStack());
+                    var originalAttribute = (CodeImpl) original;
-                    buf.writeU2(original.maxLocals());
+                    buf.writeU2(originalAttribute.maxStack());
+                    buf.writeU2(originalAttribute.maxLocals());
                 } else {
                     StackCounter cntr = StackCounter.of(DirectCodeBuilder.this, buf);
                     buf.writeU2(cntr.maxStack());

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

@@ -25,8 +25,24 @@
 package jdk.internal.classfile.impl;
 
 import java.lang.classfile.CodeBuilder;
+import java.lang.classfile.CodeModel;
+import java.lang.classfile.attribute.CodeAttribute;
 
 public sealed interface TerminalCodeBuilder extends CodeBuilder, LabelContext
         permits DirectCodeBuilder, BufferedCodeBuilder {
     int curTopLocal();
+
+    static int setupTopLocal(MethodInfo methodInfo, CodeModel original) {
+        int paramSlots = Util.maxLocals(methodInfo.methodFlags(), methodInfo.methodTypeSymbol());
+        if (original == null) {
+            return paramSlots;
+        }
+        if (original instanceof CodeAttribute attr) {
+            return Math.max(paramSlots, attr.maxLocals());
+        }
+        if (original instanceof BufferedCodeBuilder.Model buffered) {
+            return Math.max(paramSlots, buffered.curTopLocal());
+        }
+        throw new InternalError("Unknown code model " + original);
+    }
 }

test/jdk/jdk/classfile/BuilderBlockTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2024, 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
@@ -23,36 +23,48 @@
 
 /*
  * @test
+ * @bug 8337225
  * @summary Testing ClassFile builder blocks.
  * @run junit BuilderBlockTest
  */
-import java.lang.constant.ClassDesc;
+import java.io.IOException;
-
-import static java.lang.constant.ConstantDescs.*;
-
-import java.lang.constant.MethodTypeDesc;
-import java.lang.reflect.Method;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-
-import helpers.ByteArrayClassLoader;
 import java.lang.classfile.AccessFlags;
-import java.lang.reflect.AccessFlag;
+import java.lang.classfile.Attributes;
 import java.lang.classfile.ClassFile;
+import java.lang.classfile.ClassTransform;
+import java.lang.classfile.CodeBuilder;
+import java.lang.classfile.CodeElement;
+import java.lang.classfile.CodeTransform;
 import java.lang.classfile.Label;
+import java.lang.classfile.MethodModel;
+import java.lang.classfile.MethodTransform;
 import java.lang.classfile.Opcode;
 import java.lang.classfile.TypeKind;
+import java.lang.classfile.constantpool.StringEntry;
+import java.lang.classfile.instruction.ConstantInstruction;
+import java.lang.constant.ClassDesc;
+import java.lang.constant.MethodTypeDesc;
+import java.lang.reflect.AccessFlag;
+import java.lang.reflect.Method;
+import java.net.URI;
+import java.nio.file.Path;
+
+import helpers.ByteArrayClassLoader;
 import jdk.internal.classfile.impl.LabelImpl;
-import static org.junit.jupiter.api.Assertions.*;
 import org.junit.jupiter.api.Test;
 
+import static java.lang.constant.ConstantDescs.CD_int;
+import static java.lang.constant.ConstantDescs.CD_void;
+
+import static org.junit.jupiter.api.Assertions.*;
+
 /**
  * BuilderBlockTest
  */
 class BuilderBlockTest {
 
-    static final String testClassName = "AdaptCodeTest$TestClass";
+    static final String testClassName = "BuilderBlockTest$TestClass";
-    static final Path testClassPath = Paths.get("target/test-classes/" + testClassName + ".class");
+    static final Path testClassPath = Path.of(URI.create(BuilderBlockTest.class.getResource(testClassName + ".class").toString()));
 
     @Test
     void testStartEnd() throws Exception {
@@ -305,4 +317,81 @@ class BuilderBlockTest {
                           }));
         });
     }
+
+    private static final CodeTransform ALLOCATE_LOCAL_EXAMINER = CodeTransform.ofStateful(() -> new CodeTransform() {
+        boolean foundItem = false;
+
+        @Override
+        public void atStart(CodeBuilder builder) {
+            foundItem = false;
+        }
+
+        @Override
+        public void accept(CodeBuilder cob, CodeElement coe) {
+            cob.with(coe);
+            if (coe instanceof ConstantInstruction.LoadConstantInstruction ldc
+                    && ldc.constantEntry() instanceof StringEntry se
+                    && se.utf8().equalsString("Output")) {
+                assertFalse(foundItem);
+                foundItem = true;
+                var i = cob.allocateLocal(TypeKind.IntType);
+                assertEquals(7, i, "Allocated new int slot");
+            }
+        }
+
+        @Override
+        public void atEnd(CodeBuilder builder) {
+            assertTrue(foundItem);
+        }
+    });
+
+    // Test updating local variable slot management from
+    // source code models in transformingCode;
+    // CodeBuilder.transform(CodeModel, CodeTransform) is
+    // not managed for now
+    @Test
+    void testAllocateLocalTransformingCodeAttribute() throws IOException {
+        var cf = ClassFile.of();
+        var code = cf.parse(testClassPath)
+                .methods()
+                .stream()
+                .filter(f -> f.methodName().equalsString("work"))
+                .findFirst()
+                .orElseThrow()
+                .findAttribute(Attributes.code())
+                .orElseThrow();
+        ClassFile.of().build(ClassDesc.of("Foo"), cb -> cb
+                .withMethod("foo", MethodTypeDesc.ofDescriptor("(IJI)V"), 0, mb -> mb
+                        .transformCode(code, ALLOCATE_LOCAL_EXAMINER)));
+    }
+
+    @Test
+    void testAllocateLocalTransformingBufferedCode() throws IOException {
+        var cf = ClassFile.of();
+        var testClass = cf.parse(testClassPath);
+        ClassTransform bufferingTransform = (clb, cle) -> {
+            if (cle instanceof MethodModel mm && mm.methodName().equalsString("work")) {
+                clb.withMethodBody(mm.methodName(), mm.methodType(), mm.flags().flagsMask(), cob -> {
+                    int d = cob.allocateLocal(TypeKind.IntType);
+                    int e = cob.allocateLocal(TypeKind.IntType);
+
+                    assertEquals(5, d);
+                    assertEquals(6, e);
+
+                    mm.code().ifPresent(code -> code.forEach(cob));
+                });
+            }
+        };
+        cf.transformClass(testClass, bufferingTransform.andThen(ClassTransform.transformingMethods(MethodTransform.transformingCode(ALLOCATE_LOCAL_EXAMINER))));
+    }
+
+    public static class TestClass {
+        public void work(int a, long b, int c) {
+            int d = Math.addExact(a, 25);
+            int e = Math.multiplyExact(d, c);
+            System.out.println("Output");
+            System.out.println(e + b);
+            throw new IllegalArgumentException("foo");
+        }
+    }
 }

test/jdk/jdk/classfile/CorpusTest.java

@@ -39,6 +39,7 @@ import org.junit.jupiter.api.parallel.Execution;
 import org.junit.jupiter.api.parallel.ExecutionMode;
 
 import java.io.ByteArrayInputStream;
+import java.lang.classfile.attribute.CodeAttribute;
 import java.util.*;
 
 import static helpers.ClassRecord.assertEqualsDeep;
@@ -222,9 +223,11 @@ class CorpusTest {
             var m1 = itStack.next();
             var m2 = itNoStack.next();
             var text1 = m1.methodName().stringValue() + m1.methodType().stringValue() + ": "
-                      + m1.code().map(c -> c.maxLocals() + " / " + c.maxStack()).orElse("-");
+                      + m1.code().map(CodeAttribute.class::cast)
+                                 .map(c -> c.maxLocals() + " / " + c.maxStack()).orElse("-");
             var text2 = m2.methodName().stringValue() + m2.methodType().stringValue() + ": "
-                      + m2.code().map(c -> c.maxLocals() + " / " + c.maxStack()).orElse("-");
+                      + m2.code().map(CodeAttribute.class::cast)
+                                 .map(c -> c.maxLocals() + " / " + c.maxStack()).orElse("-");
             assertEquals(text1, text2);
         }
         assertFalse(itNoStack.hasNext());

test/jdk/jdk/classfile/DiscontinuedInstructionsTest.java

@@ -26,6 +26,7 @@
  * @summary Testing ClassFile handling JSR and RET instructions.
  * @run junit DiscontinuedInstructionsTest
  */
+import java.lang.classfile.attribute.CodeAttribute;
 import java.lang.constant.ClassDesc;
 import java.lang.constant.MethodTypeDesc;
 import java.util.ArrayList;
@@ -63,7 +64,7 @@ class DiscontinuedInstructionsTest {
                         .pop()
                         .with(DiscontinuedInstruction.RetInstruction.of(355))));
 
-        var c = cc.parse(bytes).methods().get(0).code().get();
+        var c = (CodeAttribute) cc.parse(bytes).methods().get(0).code().get();
         assertEquals(356, c.maxLocals());
         assertEquals(6, c.maxStack());
 

test/jdk/jdk/classfile/StackMapsTest.java

@@ -30,6 +30,7 @@
  */
 
 import java.lang.classfile.*;
+import java.lang.classfile.attribute.CodeAttribute;
 import java.lang.classfile.components.ClassPrinter;
 import java.net.URI;
 import java.nio.file.FileSystem;
@@ -329,7 +330,7 @@ class StackMapsTest {
         var cm = ClassFile.of().parse(bytes);
         for (var method : cm.methods()) {
             var name = method.methodName();
-            var code = method.code().orElseThrow();
+            var code = (CodeAttribute) method.code().orElseThrow();
             if (name.equalsString("a")) {
                 assertEquals(0, code.maxLocals()); // static method
                 assertEquals(0, code.maxStack());