8338661: StackMapTable is invalid if frames appear in dead code

Reviewed-by: liach

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

@@ -53,6 +53,14 @@ public class AttributeHolder {
         Util.writeAttributes(buf, attributes);
     }
 
+    @SuppressWarnings("unchecked")
+    <A> A get(AttributeMapper<A> am) {
+        for (Attribute<?> a : attributes)
+            if (a.attributeMapper() == am)
+                return (A)a;
+        return null;
+    }
+
     boolean isPresent(AttributeMapper<?> am) {
         for (Attribute<?> a : attributes)
             if (a.attributeMapper() == am)

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

@@ -25,20 +25,23 @@
  */
 package jdk.internal.classfile.impl;
 
+import java.lang.classfile.Attributes;
 import java.lang.classfile.TypeKind;
+import java.lang.classfile.attribute.StackMapFrameInfo;
+import java.lang.classfile.attribute.StackMapTableAttribute;
 import java.lang.classfile.constantpool.ConstantDynamicEntry;
 import java.lang.classfile.constantpool.DynamicConstantPoolEntry;
 import java.lang.classfile.constantpool.MemberRefEntry;
+import java.lang.constant.ClassDesc;
 import java.lang.constant.MethodTypeDesc;
 import java.nio.ByteBuffer;
 import java.util.ArrayDeque;
 import java.util.BitSet;
 import java.util.List;
 import java.util.Queue;
+import java.util.stream.Collectors;
 
 import static java.lang.classfile.ClassFile.*;
-import java.lang.constant.ClassDesc;
-import java.util.stream.Collectors;
 
 public final class StackCounter {
 
@@ -47,6 +50,7 @@ public final class StackCounter {
     static StackCounter of(DirectCodeBuilder dcb, BufWriterImpl buf) {
         return new StackCounter(
                 dcb,
+                dcb.attributes.get(Attributes.stackMapTable()),
                 buf.thisClass().asSymbol(),
                 dcb.methodInfo.methodName().stringValue(),
                 dcb.methodInfo.methodTypeSymbol(),
@@ -97,6 +101,7 @@ public final class StackCounter {
     }
 
     public StackCounter(LabelContext labelContext,
+                     StackMapTableAttribute smta,
                      ClassDesc thisClass,
                      String methodName,
                      MethodTypeDesc methodDesc,
@@ -111,8 +116,20 @@ public final class StackCounter {
         this.bytecode = bytecode;
         this.cp = cp;
         targets = new ArrayDeque<>();
-        maxStack = stack = rets = 0;
+        stack = rets = 0;
+        maxStack = handlers.isEmpty() ? 0 : 1;
         for (var h : handlers) targets.add(new Target(labelContext.labelToBci(h.handler), 1));
+        if (smta != null) {
+            for (var smfi : smta.entries()) {
+                int frameStack = smfi.stack().size();
+                for (var vti : smfi.stack()) {
+                    if (vti == StackMapFrameInfo.SimpleVerificationTypeInfo.ITEM_LONG
+                     || vti == StackMapFrameInfo.SimpleVerificationTypeInfo.ITEM_DOUBLE) frameStack++;
+                }
+                if (maxStack < frameStack) maxStack = frameStack;
+                targets.add(new Target(labelContext.labelToBci(smfi.target()), frameStack));
+            }
+        }
         maxLocals = isStatic ? 0 : 1;
         maxLocals += Util.parameterSlots(methodDesc);
         bcs = new RawBytecodeHelper(bytecode);

test/jdk/jdk/classfile/StackMapsTest.java

@@ -24,18 +24,25 @@
 /*
  * @test
  * @summary Testing Classfile stack maps generator.
- * @bug 8305990 8320222 8320618 8335475 8338623
+ * @bug 8305990 8320222 8320618 8335475 8338623 8338661
  * @build testdata.*
  * @run junit StackMapsTest
  */
 
 import java.lang.classfile.*;
 import java.lang.classfile.attribute.CodeAttribute;
+import java.lang.classfile.attribute.StackMapFrameInfo;
+import java.lang.classfile.attribute.StackMapTableAttribute;
 import java.lang.classfile.components.ClassPrinter;
+import java.lang.constant.ClassDesc;
+import java.lang.constant.ConstantDescs;
+import java.lang.constant.MethodTypeDesc;
+import java.lang.reflect.AccessFlag;
 import java.net.URI;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
+import java.util.List;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
@@ -45,11 +52,6 @@ import static helpers.TestUtil.assertEmpty;
 import static java.lang.classfile.ClassFile.ACC_STATIC;
 import static java.lang.constant.ConstantDescs.*;
 
-import java.lang.constant.ClassDesc;
-import java.lang.constant.ConstantDescs;
-import java.lang.constant.MethodTypeDesc;
-import java.lang.reflect.AccessFlag;
-
 /**
  * StackMapsTest
  */
@@ -368,4 +370,30 @@ class StackMapsTest {
             assertEquals(1, code.maxStack());
         }
     }
+
+    @Test
+    void testDeadCodeCountersWithCustomSMTA() {
+        ClassDesc bar = ClassDesc.of("Bar");
+        byte[] bytes = ClassFile.of(ClassFile.StackMapsOption.DROP_STACK_MAPS).build(bar, clb -> clb
+                .withMethodBody(
+                        "foo", MethodTypeDesc.of(ConstantDescs.CD_long), ACC_STATIC, cob -> {
+                            cob.lconst_0().lreturn();
+                            Label f2 = cob.newBoundLabel();
+                            cob.lstore(0);
+                            Label f3 = cob.newBoundLabel();
+                            cob.lload(0).lreturn().with(
+                                    StackMapTableAttribute.of(List.of(
+                                    StackMapFrameInfo.of(f2,
+                                            List.of(),
+                                            List.of(StackMapFrameInfo.SimpleVerificationTypeInfo.ITEM_LONG)),
+                                    StackMapFrameInfo.of(f3,
+                                            List.of(StackMapFrameInfo.SimpleVerificationTypeInfo.ITEM_LONG),
+                                            List.of()))));
+                        }
+                ));
+        assertEmpty(ClassFile.of().verify(bytes));
+        var code = (CodeAttribute) ClassFile.of().parse(bytes).methods().getFirst().code().orElseThrow();
+        assertEquals(2, code.maxLocals());
+        assertEquals(2, code.maxStack());
+    }
 }

test/micro/org/openjdk/bench/jdk/classfile/CodeAttributeTools.java

@@ -122,6 +122,7 @@ public class CodeAttributeTools {
     public void benchmarkStackCounter(Blackhole bh) {
         for (var d : data) bh.consume(new StackCounter(
                 d.labelContext(),
+                null,
                 d.thisClass(),
                 d.methodName(),
                 d.methodDesc(),