8332614: Type-checked ConstantPool.entryByIndex and ClassReader.readEntryOrNull

Reviewed-by: asotona

Author: liach

src/java.base/share/classes/java/lang/classfile/ClassReader.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
@@ -76,17 +76,14 @@ public sealed interface ClassReader extends ConstantPool
 
     // Constant pool
 
-    /**
-     * {@return the UTF8 constant pool entry at the given index of the constant
-     * pool}  The given index must correspond to a valid constant pool index
-     * whose slot holds a UTF8 constant.
-     * @param index the index into the constant pool
-     */
-    Utf8Entry utf8EntryByIndex(int index);
-
     /**
      * {@return the constant pool entry whose index is given at the specified
      * offset within the classfile}
+     *
+     * @apiNote
+     * If only a particular type of entry is expected, use {@link #readEntry(
+     * int, Class) readEntry(int, Class)}.
+     *
      * @param offset the offset of the index within the classfile
      * @throws ConstantPoolException if the index is out of range of the
      *         constant pool size, or zero
@@ -108,12 +105,31 @@ public sealed interface ClassReader extends ConstantPool
      * {@return the constant pool entry whose index is given at the specified
      * offset within the classfile, or null if the index at the specified
      * offset is zero}
+     *
+     * @apiNote
+     * If only a particular type of entry is expected, use {@link #readEntryOrNull(
+     * int, Class) readEntryOrNull(int, Class)}.
+     *
      * @param offset the offset of the index within the classfile
      * @throws ConstantPoolException if the index is out of range of the
      *         constant pool size
      */
     PoolEntry readEntryOrNull(int offset);
 
+    /**
+     * {@return the constant pool entry of a given type whose index is given
+     * at the specified offset within the classfile, or null if the index at
+     * the specified offset is zero}
+     *
+     * @param <T> the entry type
+     * @param offset the offset of the index within the classfile
+     * @param cls the entry type
+     * @throws ConstantPoolException if the index is out of range of the
+     *         constant pool size, or zero, or the entry is not of the given type
+     * @since 24
+     */
+    <T extends PoolEntry> T readEntryOrNull(int offset, Class<T> cls);
+
     /**
      * {@return the UTF8 entry whose index is given at the specified
      * offset within the classfile}

src/java.base/share/classes/java/lang/classfile/constantpool/ConstantPool.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
@@ -46,6 +46,10 @@ public sealed interface ConstantPool extends Iterable<PoolEntry>
     /**
      * {@return the entry at the specified index}
      *
+     * @apiNote
+     * If only a particular type of entry is expected, use {@link #entryByIndex(
+     * int, Class) entryByIndex(int, Class)}.
+     *
      * @param index the index within the pool of the desired entry
      * @throws ConstantPoolException if the index is out of range of the
      *         constant pool, or is considered unusable
@@ -57,6 +61,18 @@ public sealed interface ConstantPool extends Iterable<PoolEntry>
      */
     int size();
 
+    /**
+     * {@return the entry of a given type at the specified index}
+     *
+     * @param <T> the entry type
+     * @param index the index within the pool of the desired entry
+     * @param cls the entry type
+     * @throws ConstantPoolException if the index is out of range of the
+     *         constant pool, or the entry is not of the given type
+     * @since 24
+     */
+    <T extends PoolEntry> T entryByIndex(int index, Class<T> cls);
+
     /**
      * {@return an iterator over pool entries}
      */

src/java.base/share/classes/jdk/internal/classfile/impl/AbstractBoundLocalVariable.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
@@ -46,7 +46,7 @@ protected int nameIndex() {
 
     public Utf8Entry name() {
         if (nameEntry == null)
-            nameEntry = (Utf8Entry) code.constantPool().entryByIndex(nameIndex());
+            nameEntry = code.constantPool().entryByIndex(nameIndex(), Utf8Entry.class);
         return nameEntry;
     }
 
@@ -56,7 +56,7 @@ protected int secondaryIndex() {
 
     protected Utf8Entry secondaryEntry() {
         if (secondaryEntry == null)
-            secondaryEntry = (Utf8Entry) code.constantPool().entryByIndex(secondaryIndex());
+            secondaryEntry = code.constantPool().entryByIndex(secondaryIndex(), Utf8Entry.class);
         return secondaryEntry;
     }
 

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

@@ -687,10 +687,10 @@ public BoundLoadConstantInstruction(Opcode op, CodeImpl code, int pos) {
 
         @Override
         public LoadableConstantEntry constantEntry() {
-            return (LoadableConstantEntry)
-                    code.classReader.entryByIndex(op == Opcode.LDC
+            return code.classReader.entryByIndex(op == Opcode.LDC
                                                   ? code.classReader.readU1(pos + 1)
-                                                  : code.classReader.readU2(pos + 1));
+                                                  : code.classReader.readU2(pos + 1),
+                            LoadableConstantEntry.class);
         }
 
         @Override

src/java.base/share/classes/jdk/internal/classfile/impl/AnnotationReader.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
@@ -127,7 +127,7 @@ private static int skipElementValue(ClassReader classReader, int p) {
     }
 
     private static Annotation readAnnotation(ClassReader classReader, int p) {
-        Utf8Entry annotationClass = classReader.utf8EntryByIndex(classReader.readU2(p));
+        Utf8Entry annotationClass = classReader.entryByIndex(classReader.readU2(p), Utf8Entry.class);
         p += 2;
         List<AnnotationElement> elems = readAnnotationElementValuePairs(classReader, p);
         return new AnnotationImpl(annotationClass, elems);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2023, 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
@@ -632,7 +632,7 @@ private void structure() {
                 for (int i = 0; p < end; p += 6, i++) {
                     elements[i] = ModuleRequireInfo.of(classReader.readModuleEntry(p),
                             classReader.readU2(p + 2),
-                            (Utf8Entry) classReader.readEntryOrNull(p + 4));
+                            classReader.readEntryOrNull(p + 4, Utf8Entry.class));
                 }
                 requires = List.of(elements);
             }
@@ -771,15 +771,9 @@ public List<InnerClassInfo> classes() {
                 int p = payloadStart + 2;
                 InnerClassInfo[] elements = new InnerClassInfo[cnt];
                 for (int i = 0; i < cnt; i++) {
-                    ClassEntry innerClass = classReader.readClassEntry(p); // TODO FIXME
-                    int outerClassIndex = classReader.readU2(p + 2);
-                    ClassEntry outerClass = outerClassIndex == 0
-                            ? null
-                            : (ClassEntry) classReader.entryByIndex(outerClassIndex);
-                    int innerNameIndex = classReader.readU2(p + 4);
-                    Utf8Entry innerName = innerNameIndex == 0
-                            ? null
-                            : (Utf8Entry) classReader.entryByIndex(innerNameIndex);
+                    ClassEntry innerClass = classReader.readClassEntry(p);
+                    var outerClass = classReader.readEntryOrNull(p + 2, ClassEntry.class);
+                    var innerName = classReader.readEntryOrNull(p + 4, Utf8Entry.class);
                     int flags = classReader.readU2(p + 6);
                     p += 8;
                     elements[i] = InnerClassInfo.of(innerClass, Optional.ofNullable(outerClass), Optional.ofNullable(innerName), flags);
@@ -803,7 +797,7 @@ public ClassEntry enclosingClass() {
 
         @Override
         public Optional<NameAndTypeEntry> enclosingMethod() {
-            return Optional.ofNullable((NameAndTypeEntry) classReader.readEntryOrNull(payloadStart + 2));
+            return Optional.ofNullable(classReader.readEntryOrNull(payloadStart + 2, NameAndTypeEntry.class));
         }
     }
 

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

@@ -28,6 +28,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Function;
 
@@ -158,8 +159,7 @@ public ClassEntry thisClassEntry() {
     @Override
     public Optional<ClassEntry> superclassEntry() {
         if (superclass == null) {
-            int scIndex = readU2(thisClassPos + 2);
-            superclass = Optional.ofNullable(scIndex == 0 ? null : (ClassEntry) entryByIndex(scIndex));
+            superclass = Optional.ofNullable(readEntryOrNull(thisClassPos + 2, ClassEntry.class));
         }
         return superclass;
     }
@@ -338,10 +338,42 @@ void writeConstantPoolEntries(BufWriter buf) {
     // Constantpool
     @Override
     public PoolEntry entryByIndex(int index) {
-        return entryByIndex(index, 0, 0xff);
+        return entryByIndex(index, PoolEntry.class);
+    }
+
+    private static boolean checkTag(int tag, Class<?> cls) {
+        var type = switch (tag) {
+            // JVMS Table 4.4-B. Constant pool tags
+            case TAG_UTF8 -> AbstractPoolEntry.Utf8EntryImpl.class;
+            case TAG_INTEGER -> AbstractPoolEntry.IntegerEntryImpl.class;
+            case TAG_FLOAT -> AbstractPoolEntry.FloatEntryImpl.class;
+            case TAG_LONG -> AbstractPoolEntry.LongEntryImpl.class;
+            case TAG_DOUBLE -> AbstractPoolEntry.DoubleEntryImpl.class;
+            case TAG_CLASS -> AbstractPoolEntry.ClassEntryImpl.class;
+            case TAG_STRING -> AbstractPoolEntry.StringEntryImpl.class;
+            case TAG_FIELDREF -> AbstractPoolEntry.FieldRefEntryImpl.class;
+            case TAG_METHODREF -> AbstractPoolEntry.MethodRefEntryImpl.class;
+            case TAG_INTERFACEMETHODREF -> AbstractPoolEntry.InterfaceMethodRefEntryImpl.class;
+            case TAG_NAMEANDTYPE -> AbstractPoolEntry.NameAndTypeEntryImpl.class;
+            case TAG_METHODHANDLE -> AbstractPoolEntry.MethodHandleEntryImpl.class;
+            case TAG_METHODTYPE -> AbstractPoolEntry.MethodTypeEntryImpl.class;
+            case TAG_CONSTANTDYNAMIC -> AbstractPoolEntry.ConstantDynamicEntryImpl.class;
+            case TAG_INVOKEDYNAMIC -> AbstractPoolEntry.InvokeDynamicEntryImpl.class;
+            case TAG_MODULE -> AbstractPoolEntry.ModuleEntryImpl.class;
+            case TAG_PACKAGE -> AbstractPoolEntry.PackageEntryImpl.class;
+            default -> null;
+        };
+        return type != null && cls.isAssignableFrom(type);
+    }
+
+    static <T extends PoolEntry> T checkType(PoolEntry e, int index, Class<T> cls) {
+        if (cls.isInstance(e)) return cls.cast(e);
+        throw new ConstantPoolException("Not a " + cls.getSimpleName() + " at index: " + index);
     }
 
-    private PoolEntry entryByIndex(int index, int lowerBoundTag, int upperBoundTag) {
+    @Override
+    public <T extends PoolEntry> T entryByIndex(int index, Class<T> cls) {
+        Objects.requireNonNull(cls);
         if (index <= 0 || index >= constantPoolCount) {
             throw new ConstantPoolException("Bad CP index: " + index);
         }
@@ -352,9 +384,9 @@ private PoolEntry entryByIndex(int index, int lowerBoundTag, int upperBoundTag)
                 throw new ConstantPoolException("Unusable CP index: " + index);
             }
             int tag = readU1(offset);
-            if (tag < lowerBoundTag || tag > upperBoundTag) {
+            if (!checkTag(tag, cls)) {
                 throw new ConstantPoolException(
-                        "Bad tag (" + tag + ") at index (" + index + ") position (" + offset + ")");
+                        "Bad tag (" + tag + ") at index (" + index + ") position (" + offset + "), expected " + cls.getSimpleName());
             }
             final int q = offset + 1;
             info = switch (tag) {
@@ -374,7 +406,7 @@ private PoolEntry entryByIndex(int index, int lowerBoundTag, int upperBoundTag)
                 case TAG_NAMEANDTYPE -> new AbstractPoolEntry.NameAndTypeEntryImpl(this, index, (AbstractPoolEntry.Utf8EntryImpl) readUtf8Entry(q),
                                                                                    (AbstractPoolEntry.Utf8EntryImpl) readUtf8Entry(q + 2));
                 case TAG_METHODHANDLE -> new AbstractPoolEntry.MethodHandleEntryImpl(this, index, readU1(q),
-                                                                                     readEntry(q + 1, AbstractPoolEntry.AbstractMemberRefEntry.class, TAG_FIELDREF, TAG_INTERFACEMETHODREF));
+                                                                                     readEntry(q + 1, AbstractPoolEntry.AbstractMemberRefEntry.class));
                 case TAG_METHODTYPE -> new AbstractPoolEntry.MethodTypeEntryImpl(this, index, (AbstractPoolEntry.Utf8EntryImpl) readUtf8Entry(q));
                 case TAG_CONSTANTDYNAMIC -> new AbstractPoolEntry.ConstantDynamicEntryImpl(this, index, readU2(q), (AbstractPoolEntry.NameAndTypeEntryImpl) readNameAndTypeEntry(q + 2));
                 case TAG_INVOKEDYNAMIC -> new AbstractPoolEntry.InvokeDynamicEntryImpl(this, index, readU2(q), (AbstractPoolEntry.NameAndTypeEntryImpl) readNameAndTypeEntry(q + 2));
@@ -385,15 +417,7 @@ private PoolEntry entryByIndex(int index, int lowerBoundTag, int upperBoundTag)
             };
             cp[index] = info;
         }
-        return info;
-    }
-
-    @Override
-    public AbstractPoolEntry.Utf8EntryImpl utf8EntryByIndex(int index) {
-        if (entryByIndex(index, TAG_UTF8, TAG_UTF8) instanceof AbstractPoolEntry.Utf8EntryImpl utf8) {
-            return utf8;
-        }
-        throw new ConstantPoolException("Not a UTF8 - index: " + index);
+        return checkType(info, index, cls);
     }
 
     public int skipAttributeHolder(int offset) {
@@ -418,17 +442,8 @@ public PoolEntry readEntry(int pos) {
 
     @Override
     public <T extends PoolEntry> T readEntry(int pos, Class<T> cls) {
-        return readEntry(pos, cls, 0, 0xff);
-    }
-
-    private <T extends PoolEntry> T readEntry(int pos, Class<T> cls, int expectedTag) {
-        return readEntry(pos, cls, expectedTag, expectedTag);
-    }
-
-    private <T extends PoolEntry> T readEntry(int pos, Class<T> cls, int lowerBoundTag, int upperBoundTag) {
-        var e = entryByIndex(readU2(pos), lowerBoundTag, upperBoundTag);
-        if (cls.isInstance(e)) return cls.cast(e);
-        throw new ConstantPoolException("Not a " + cls.getSimpleName() + " at index: " + readU2(pos));
+        Objects.requireNonNull(cls);
+        return entryByIndex(readU2(pos), cls);
     }
 
     @Override
@@ -440,44 +455,49 @@ public PoolEntry readEntryOrNull(int pos) {
         return entryByIndex(index);
     }
 
+    @Override
+    public <T extends PoolEntry> T readEntryOrNull(int offset, Class<T> cls) {
+        Objects.requireNonNull(cls);
+        int index = readU2(offset);
+        if (index == 0) {
+            return null;
+        }
+        return entryByIndex(index, cls);
+    }
+
     @Override
     public Utf8Entry readUtf8Entry(int pos) {
-        int index = readU2(pos);
-        return utf8EntryByIndex(index);
+        return readEntry(pos, Utf8Entry.class);
     }
 
     @Override
     public Utf8Entry readUtf8EntryOrNull(int pos) {
-        int index = readU2(pos);
-        if (index == 0) {
-            return null;
-        }
-        return utf8EntryByIndex(index);
+        return readEntryOrNull(pos, Utf8Entry.class);
     }
 
     @Override
     public ModuleEntry readModuleEntry(int pos) {
-        return readEntry(pos, ModuleEntry.class, TAG_MODULE);
+        return readEntry(pos, ModuleEntry.class);
     }
 
     @Override
     public PackageEntry readPackageEntry(int pos) {
-        return readEntry(pos, PackageEntry.class, TAG_PACKAGE);
+        return readEntry(pos, PackageEntry.class);
     }
 
     @Override
     public ClassEntry readClassEntry(int pos) {
-        return readEntry(pos, ClassEntry.class, TAG_CLASS);
+        return readEntry(pos, ClassEntry.class);
     }
 
     @Override
     public NameAndTypeEntry readNameAndTypeEntry(int pos) {
-        return readEntry(pos, NameAndTypeEntry.class, TAG_NAMEANDTYPE);
+        return readEntry(pos, NameAndTypeEntry.class);
     }
 
     @Override
     public MethodHandleEntry readMethodHandleEntry(int pos) {
-        return readEntry(pos, MethodHandleEntry.class, TAG_METHODHANDLE);
+        return readEntry(pos, MethodHandleEntry.class);
     }
 
     @Override

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

@@ -199,7 +199,7 @@ public List<ExceptionCatch> exceptionHandlers() {
                 public void accept(int s, int e, int h, int c) {
                     ClassEntry catchTypeEntry = c == 0
                                                              ? null
-                                                             : (ClassEntry) constantPool().entryByIndex(c);
+                                                             : constantPool().entryByIndex(c, ClassEntry.class);
                     exceptionTable.add(new AbstractPseudoInstruction.ExceptionCatchImpl(getLabel(h), getLabel(s), getLabel(e), catchTypeEntry));
                 }
             });
@@ -337,7 +337,7 @@ private void generateCatchTargets(Consumer<CodeElement> consumer) {
             public void accept(int s, int e, int h, int c) {
                 ClassEntry catchType = c == 0
                                                     ? null
-                                                    : (ClassEntry) classReader.entryByIndex(c);
+                                                    : classReader.entryByIndex(c, ClassEntry.class);
                 consumer.accept(new AbstractPseudoInstruction.ExceptionCatchImpl(getLabel(h), getLabel(s), getLabel(e), catchType));
             }
         });

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2023, 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
@@ -28,7 +28,6 @@
 import java.lang.constant.MethodTypeDesc;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Objects;
 
 import java.lang.classfile.Attribute;
 import java.lang.classfile.Attributes;
@@ -38,6 +37,7 @@
 import java.lang.classfile.BufWriter;
 import java.lang.classfile.attribute.BootstrapMethodsAttribute;
 import java.lang.classfile.constantpool.*;
+import java.util.Objects;
 
 import static java.lang.classfile.ClassFile.TAG_CLASS;
 import static java.lang.classfile.ClassFile.TAG_CONSTANTDYNAMIC;
@@ -114,6 +114,12 @@ public PoolEntry entryByIndex(int index) {
         return pe;
     }
 
+    @Override
+    public <T extends PoolEntry> T entryByIndex(int index, Class<T> cls) {
+        Objects.requireNonNull(cls);
+        return ClassReaderImpl.checkType(entryByIndex(index), index, cls);
+    }
+
     @Override
     public BootstrapMethodEntryImpl bootstrapMethodEntry(int index) {
         if (index < 0 || index >= bootstrapMethodCount()) {

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

@@ -296,7 +296,7 @@ public StackCounter(LabelContext labelContext,
                         next();
                     }
                     case GETSTATIC, PUTSTATIC, GETFIELD, PUTFIELD -> {
-                        var tk = TypeKind.fromDescriptor(((MemberRefEntry)cp.entryByIndex(bcs.getIndexU2())).nameAndType().type().stringValue());
+                        var tk = TypeKind.fromDescriptor(cp.entryByIndex(bcs.getIndexU2(), MemberRefEntry.class).nameAndType().type());
                         switch (bcs.rawCode) {
                             case GETSTATIC ->
                                 addStackSlot(tk.slotSize());
@@ -368,7 +368,7 @@ private void processLdc(int index) {
             case TAG_DOUBLE, TAG_LONG ->
                 addStackSlot(+2);
             case TAG_CONSTANTDYNAMIC ->
-                addStackSlot(((ConstantDynamicEntry)cp.entryByIndex(index)).typeKind().slotSize());
+                addStackSlot(cp.entryByIndex(index, ConstantDynamicEntry.class).typeKind().slotSize());
             default ->
                 throw error("CP entry #%d %s is not loadable constant".formatted(index, cp.entryByIndex(index).tag()));
         }

src/java.base/share/classes/jdk/internal/classfile/impl/StackMapDecoder.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
@@ -231,7 +231,7 @@ private VerificationTypeInfo readVerificationTypeInfo() {
             case VT_LONG -> SimpleVerificationTypeInfo.ITEM_LONG;
             case VT_NULL -> SimpleVerificationTypeInfo.ITEM_NULL;
             case VT_UNINITIALIZED_THIS -> SimpleVerificationTypeInfo.ITEM_UNINITIALIZED_THIS;
-            case VT_OBJECT -> new ObjectVerificationTypeInfoImpl((ClassEntry)classReader.entryByIndex(u2()));
+            case VT_OBJECT -> new ObjectVerificationTypeInfoImpl(classReader.entryByIndex(u2(), ClassEntry.class));
             case VT_UNINITIALIZED -> new UninitializedVerificationTypeInfoImpl(ctx.getLabel(u2()));
             default -> throw new IllegalArgumentException("Invalid verification type tag: " + tag);
         };

src/java.base/share/classes/jdk/internal/classfile/impl/StackMapGenerator.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
@@ -25,6 +25,8 @@
  */
 package jdk.internal.classfile.impl;
 
+import java.lang.classfile.constantpool.InvokeDynamicEntry;
+import java.lang.classfile.constantpool.NameAndTypeEntry;
 import java.lang.constant.ClassDesc;
 import static java.lang.constant.ConstantDescs.*;
 import java.lang.constant.MethodTypeDesc;
@@ -397,7 +399,7 @@ public void writeBody(BufWriter b) {
     }
 
     private static Type cpIndexToType(int index, ConstantPoolBuilder cp) {
-        return Type.referenceType(((ClassEntry)cp.entryByIndex(index)).asSymbol());
+        return Type.referenceType(cp.entryByIndex(index, ClassEntry.class).asSymbol());
     }
 
     private void processMethod() {
@@ -700,7 +702,7 @@ private void processLdc(int index) {
             case TAG_METHODTYPE ->
                 currentFrame.pushStack(Type.METHOD_TYPE);
             case TAG_CONSTANTDYNAMIC ->
-                currentFrame.pushStack(((ConstantDynamicEntry)cp.entryByIndex(index)).asSymbol().constantType());
+                currentFrame.pushStack(cp.entryByIndex(index, ConstantDynamicEntry.class).asSymbol().constantType());
             default ->
                 throw generatorError("CP entry #%d %s is not loadable constant".formatted(index, cp.entryByIndex(index).tag()));
         }
@@ -747,7 +749,7 @@ private void processSwitch(RawBytecodeHelper bcs) {
     }
 
     private void processFieldInstructions(RawBytecodeHelper bcs) {
-        var desc = Util.fieldTypeSymbol(((MemberRefEntry)cp.entryByIndex(bcs.getIndexU2())).nameAndType());
+        var desc = Util.fieldTypeSymbol(cp.entryByIndex(bcs.getIndexU2(), MemberRefEntry.class).nameAndType());
         switch (bcs.rawCode) {
             case GETSTATIC ->
                 currentFrame.pushStack(desc);
@@ -771,8 +773,9 @@ private void processFieldInstructions(RawBytecodeHelper bcs) {
     private boolean processInvokeInstructions(RawBytecodeHelper bcs, boolean inTryBlock, boolean thisUninit) {
         int index = bcs.getIndexU2();
         int opcode = bcs.rawCode;
-        var cpe = cp.entryByIndex(index);
-        var nameAndType = opcode == INVOKEDYNAMIC ? ((DynamicConstantPoolEntry)cpe).nameAndType() : ((MemberRefEntry)cpe).nameAndType();
+        var nameAndType = opcode == INVOKEDYNAMIC
+                ? cp.entryByIndex(index, InvokeDynamicEntry.class).nameAndType()
+                : cp.entryByIndex(index, MemberRefEntry.class).nameAndType();
         String invokeMethodName = nameAndType.name().stringValue();
         var mDesc = Util.methodTypeSymbol(nameAndType);
         int bci = bcs.bci;

src/java.base/share/classes/jdk/internal/classfile/impl/TemporaryConstantPool.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
@@ -170,6 +170,11 @@ public int size() {
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public <T extends PoolEntry> T entryByIndex(int index, Class<T> cls) {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public BootstrapMethodEntry bootstrapMethodEntry(int index) {
         throw new UnsupportedOperationException();

src/java.base/share/classes/jdk/internal/classfile/impl/verifier/VerificationWrapper.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
@@ -166,11 +166,11 @@ int entryCount() {
         }
 
         String classNameAt(int index) {
-            return ((ClassEntry)cp.entryByIndex(index)).asInternalName();
+            return cp.entryByIndex(index, ClassEntry.class).asInternalName();
         }
 
         String dynamicConstantSignatureAt(int index) {
-            return ((DynamicConstantPoolEntry)cp.entryByIndex(index)).type().stringValue();
+            return cp.entryByIndex(index, DynamicConstantPoolEntry.class).type().stringValue();
         }
 
         int tagAt(int index) {
@@ -192,7 +192,7 @@ String refSignatureAt(int index) {
         }
 
         int refClassIndexAt(int index) {
-            return ((MemberRefEntry)cp.entryByIndex(index)).owner().index();
+            return cp.entryByIndex(index, MemberRefEntry.class).owner().index();
         }
     }
 }

test/jdk/jdk/classfile/AttributesTest.java

@@ -23,14 +23,23 @@
 
 /*
  * @test
- * @bug 8331291
- * @summary Testing Attributes API.
+ * @bug 8331291 8332614
+ * @summary Testing Attributes API and ClassReader.
  * @run junit AttributesTest
  */
 import java.lang.classfile.AttributeMapper;
+import java.lang.classfile.AttributedElement;
 import java.lang.classfile.Attributes;
+import java.lang.classfile.BufWriter;
+import java.lang.classfile.ClassFile;
+import java.lang.classfile.ClassReader;
+import java.lang.classfile.CustomAttribute;
+import java.lang.classfile.constantpool.ConstantPoolException;
+import java.lang.classfile.constantpool.InvokeDynamicEntry;
 import java.lang.classfile.constantpool.Utf8Entry;
+import java.lang.constant.ClassDesc;
 import java.lang.reflect.Field;
+
 import org.junit.jupiter.api.Test;
 
 import static org.junit.jupiter.api.Assertions.*;
@@ -52,4 +61,75 @@ void testAttributesMapping() throws Exception {
             }
         }
     }
+
+    private static final String TEST_ATTRIBUTE_NAME = "org.openjdk.classfile.test";
+    private static final AttributeMapper<TestAttribute> TEST_MAPPER = new AttributeMapper<>() {
+        @Override
+        public String name() {
+            return TEST_ATTRIBUTE_NAME;
+        }
+
+        @Override
+        public TestAttribute readAttribute(AttributedElement enclosing, ClassReader cf, int pos) {
+            int cpPos = pos - 6; // Attribute Name Utf8
+            // Test valid pos/index - NPE
+            assertThrows(NullPointerException.class, () -> cf.readEntry(cpPos, null));
+            assertThrows(NullPointerException.class, () -> cf.readEntryOrNull(cpPos, null));
+            assertThrows(NullPointerException.class, () -> cf.entryByIndex(1, null));
+
+            // Test valid pos/index - incorrect type
+            assertThrows(ConstantPoolException.class, () -> cf.readEntry(cpPos, InvokeDynamicEntry.class));
+            assertThrows(ConstantPoolException.class, () -> cf.readEntryOrNull(cpPos, InvokeDynamicEntry.class));
+            assertThrows(ConstantPoolException.class, () -> cf.entryByIndex(1, InvokeDynamicEntry.class));
+
+            // Passing tests
+            var utf8 = cf.readEntry(cpPos, Utf8Entry.class);
+            assertSame(utf8, cf.readEntryOrNull(cpPos, Utf8Entry.class));
+
+            // Test invalid pos/index - NPE thrown before CPE
+            assertThrows(NullPointerException.class, () -> cf.readEntry(-1, null));
+            assertThrows(NullPointerException.class, () -> cf.readEntryOrNull(-1, null));
+            assertThrows(NullPointerException.class, () -> cf.entryByIndex(-1, null));
+
+            return new TestAttribute(true);
+        }
+
+        @Override
+        public void writeAttribute(BufWriter buf, TestAttribute attr) {
+            buf.writeIndex(buf.constantPool().utf8Entry(name()));
+            buf.writeInt(0);
+        }
+
+        @Override
+        public AttributeStability stability() {
+            return AttributeStability.STATELESS;
+        }
+    };
+
+    private static final class TestAttribute extends CustomAttribute<TestAttribute> {
+        final boolean fromMapper;
+
+        TestAttribute(boolean fromMapper) {
+            super(TEST_MAPPER);
+            this.fromMapper = fromMapper;
+        }
+    }
+
+    @Test
+    void testClassReader() throws Exception {
+        var cf = ClassFile.of(ClassFile.AttributeMapperOption.of(utf8 -> {
+            if (utf8.equalsString(TEST_ATTRIBUTE_NAME)) {
+                return TEST_MAPPER;
+            }
+            return null;
+        }));
+
+        var cd = ClassDesc.of("Testing");
+        var bytes = cf.build(cd, clb -> clb
+                .with(new TestAttribute(false)));
+        assertTrue(cf.parse(bytes)
+                .findAttribute(TEST_MAPPER)
+                .orElseThrow()
+                .fromMapper);
+    }
 }