Typed constant indices

codecs/src/main/scala/org/polyvariant/classfile/codecs/ClassFileCodecs.scala

@@ -26,6 +26,7 @@ import scodec.bits.ByteVector
 import scodec.bits._
 
 import java.nio.charset.StandardCharsets
+import scala.reflect.TypeTest
 
 object ClassFileCodecs {
 
@@ -35,24 +36,27 @@ object ClassFileCodecs {
   val u2: Codec[Int] = uint(16)
   val u4: Codec[Long] = ulong(32)
 
-  val constantPoolIndex: Codec[ConstantIndex] = u2.as[ConstantIndex]
+  def constantPoolIndex[T <: Constant]: Codec[ConstantIndex[T]] = u2.as[ConstantIndex[T]]
+
   val constantPoolIndexNarrow: Codec[ConstantIndexNarrow] = u1.as[ConstantIndexNarrow]
 
   private val fieldRefCommon =
-    ("class index" | constantPoolIndex) ::
-      ("name and type index" | constantPoolIndex)
+    ("class index" | constantPoolIndex[Constant.ClassInfo]) ::
+      ("name and type index" | constantPoolIndex[Constant])
 
   val methodRef: Codec[Constant.MethodRefInfo] = fieldRefCommon.as[Constant.MethodRefInfo]
   val fieldRef: Codec[Constant.FieldRefInfo] = fieldRefCommon.as[Constant.FieldRefInfo]
   val interfaceMethodRef: Codec[Constant.InterfaceMethodRefInfo] = fieldRefCommon
     .as[Constant.InterfaceMethodRefInfo]
 
   val nameAndType: Codec[Constant.NameAndTypeInfo] =
-    (("name index" | constantPoolIndex) :: ("descriptor index" | constantPoolIndex))
+    (("name index" | constantPoolIndex[Constant.Utf8Info]) ::
+      ("descriptor index" | constantPoolIndex[Constant.Utf8Info]))
       .as[Constant.NameAndTypeInfo]
 
-  val classConstant: Codec[Constant.ClassInfo] = ("name index" | constantPoolIndex)
-    .as[Constant.ClassInfo]
+  val classConstant: Codec[Constant.ClassInfo] =
+    ("name index" | constantPoolIndex[Constant.Utf8Info])
+      .as[Constant.ClassInfo]
 
   // def encodeString(s: String) = {
   // val charArray = s
@@ -76,8 +80,9 @@ object ClassFileCodecs {
     .xmap(bytes => new String(bytes.toArray), s => ByteVector(s.getBytes()))
     .as[Constant.Utf8Info]
 
-  val stringConstant: Codec[Constant.StringInfo] = ("string index" | constantPoolIndex)
-    .as[Constant.StringInfo]
+  val stringConstant: Codec[Constant.StringInfo] =
+    ("string index" | constantPoolIndex[Constant.Utf8Info])
+      .as[Constant.StringInfo]
 
   private val numeric = "bytes" | bytes(4)
   private val bigNumeric = ("high bytes" | bytes(4)) :: ("low bytes" | bytes(4))
@@ -87,30 +92,31 @@ object ClassFileCodecs {
   val longConstant: Codec[Constant.LongInfo] = bigNumeric.as[Constant.LongInfo]
   val doubleConstant: Codec[Constant.DoubleInfo] = bigNumeric.as[Constant.DoubleInfo]
 
-  val methodType: Codec[Constant.MethodTypeInfo] = ("descriptor index" | constantPoolIndex)
-    .as[Constant.MethodTypeInfo]
+  val methodType: Codec[Constant.MethodTypeInfo] =
+    ("descriptor index" | constantPoolIndex[Constant.NameAndTypeInfo])
+      .as[Constant.MethodTypeInfo]
 
   val methodHandle: Codec[Constant.MethodHandleInfo] =
     (("reference kind" | mappedEnum(
       u1,
       MethodReferenceKind.values.map(k => k -> k.ordinal.toShort).toMap,
     )) ::
-      ("reference index" | constantPoolIndex))
+      ("reference index" | constantPoolIndex[Constant]))
       .as[Constant.MethodHandleInfo]
 
   private val dynamicCommon =
     ("bootstrap method attr index" | u2) ::
-      ("name and type index" | constantPoolIndex)
+      ("name and type index" | constantPoolIndex[Constant.NameAndTypeInfo])
 
   val dynamic: Codec[Constant.DynamicInfo] = dynamicCommon.as[Constant.DynamicInfo]
 
   val invokeDynamic: Codec[Constant.InvokeDynamicInfo] = dynamicCommon
     .as[Constant.InvokeDynamicInfo]
 
-  val module: Codec[Constant.ModuleInfo] = ("name index" | constantPoolIndex)
+  val module: Codec[Constant.ModuleInfo] = ("name index" | constantPoolIndex[Constant.Utf8Info])
     .as[Constant.ModuleInfo]
 
-  val pkg: Codec[Constant.PackageInfo] = ("name index" | constantPoolIndex)
+  val pkg: Codec[Constant.PackageInfo] = ("name index" | constantPoolIndex[Constant.Utf8Info])
     .as[Constant.PackageInfo]
 
   val constantEntry: Codec[Constant] =
@@ -193,7 +199,7 @@ object ClassFileCodecs {
 
   val attribute: Codec[AttributeInfo] =
     "attribute" | (
-      ("name index" | constantPoolIndex) ::
+      ("name index" | constantPoolIndex[Constant.Utf8Info]) ::
         variableSizeBytesLong(
           "attribute length" | u4,
           "info" | bytes,
@@ -207,16 +213,16 @@ object ClassFileCodecs {
     "field info" |
       (
         "access flags" | fieldAccessFlags ::
-          ("name index" | constantPoolIndex) ::
-          ("descriptor index" | constantPoolIndex) ::
+          ("name index" | constantPoolIndex[Constant.Utf8Info]) ::
+          ("descriptor index" | constantPoolIndex[Constant.Utf8Info]) ::
           attributes
       ).as[FieldInfo]
 
   val methodInfo =
     "method info" | (
       ("access flags" | methodAccessFlags) ::
-        ("name index" | constantPoolIndex) ::
-        ("descriptor index" | constantPoolIndex) ::
+        ("name index" | constantPoolIndex[Constant.Utf8Info]) ::
+        ("descriptor index" | constantPoolIndex[Constant.Utf8Info]) ::
         attributes
     ).as[MethodInfo]
 
@@ -235,11 +241,11 @@ object ClassFileCodecs {
         ("major version" | u2) ::
         constantPool ::
         classAccessFlags ::
-        ("this class" | constantPoolIndex) ::
-        ("super class" | constantPoolIndex) ::
+        ("this class" | constantPoolIndex[Constant.ClassInfo]) ::
+        ("super class" | constantPoolIndex[Constant.ClassInfo]) ::
         listOfN(
           "interface count" | u2,
-          "interface index" | constantPoolIndex,
+          "interface index" | constantPoolIndex[Constant.ClassInfo],
         ) ::
         listOfN(
           "field count" | u2,

codecs/src/main/scala/org/polyvariant/classfile/codecs/InstructionCodecs.scala

@@ -45,6 +45,7 @@ object InstructionCodecs {
     val offset = ("offset" | u2).as[Offset]
     val offsetWide = ("offset" | u4).as[OffsetWide]
 
+    val anyConstant = constantPoolIndex[Constant]
     discriminated[Instruction]
       .by(u1)
       .typecase(0x32, localIndex.as[dload])
@@ -55,7 +56,7 @@ object InstructionCodecs {
       .singleton(0x2b, aload_1)
       .singleton(0x2c, aload_2)
       .singleton(0x2d, aload_3)
-      .typecase(0xbd, constantPoolIndex.as[anewarray])
+      .typecase(0xbd, anyConstant.as[anewarray])
       .singleton(0xb0, areturn)
       .singleton(0xbe, arraylength)
       .typecase(0x3a, localIndex.as[astore])
@@ -69,7 +70,7 @@ object InstructionCodecs {
       .typecase(0x10, byte.as[bipush])
       .singleton(0x34, caload)
       .singleton(0x55, castore)
-      .typecase(0xc0, constantPoolIndex.as[checkcast])
+      .typecase(0xc0, anyConstant.as[checkcast])
       .singleton(0x90, d2f)
       .singleton(0x8e, d2i)
       .singleton(0x8f, d2l)
@@ -129,8 +130,8 @@ object InstructionCodecs {
       .singleton(0x45, fstore_2)
       .singleton(0x46, fstore_3)
       .singleton(0x66, fsub)
-      .typecase(0xb4, constantPoolIndex.as[getfield])
-      .typecase(0xb2, constantPoolIndex.as[getstatic])
+      .typecase(0xb4, anyConstant.as[getfield])
+      .typecase(0xb2, anyConstant.as[getstatic])
       .typecase(0xa7, offset.as[goto])
       .typecase(0xc8, offsetWide.as[goto_w])
       .singleton(0x91, i2b)
@@ -175,15 +176,15 @@ object InstructionCodecs {
       .singleton(0x1d, iload_3)
       .singleton(0xfe, imul)
       .singleton(0x74, ineg)
-      .typecase(0xc1, constantPoolIndex.as[instanceof])
-      .typecase(0xba, (constantPoolIndex :: constant(hex"0000")).dropUnits.as[invokedynamic])
+      .typecase(0xc1, anyConstant.as[instanceof])
+      .typecase(0xba, (anyConstant :: constant(hex"0000")).dropUnits.as[invokedynamic])
       .typecase(
         0xb9,
-        (constantPoolIndex :: ("count" | u1) :: constant(hex"0")).dropUnits.as[invokeinterface],
+        (anyConstant :: ("count" | u1) :: constant(hex"0")).dropUnits.as[invokeinterface],
       )
-      .typecase(0xb7, constantPoolIndex.as[invokespecial])
-      .typecase(0xb8, constantPoolIndex.as[invokestatic])
-      .typecase(0xb6, constantPoolIndex.as[invokevirtual])
+      .typecase(0xb7, anyConstant.as[invokespecial])
+      .typecase(0xb8, anyConstant.as[invokestatic])
+      .typecase(0xb6, anyConstant.as[invokevirtual])
       .singleton(0x80, ior)
       .singleton(0x70, irem)
       .singleton(0xac, ireturn)
@@ -210,8 +211,8 @@ object InstructionCodecs {
       .singleton(0x09, lconst_0)
       .singleton(0x0a, lconst_1)
       .typecase(0x12, constantPoolIndexNarrow.as[ldc])
-      .typecase(0x13, constantPoolIndex.as[ldc_w])
-      .typecase(0x14, constantPoolIndex.as[ldc2_w])
+      .typecase(0x13, anyConstant.as[ldc_w])
+      .typecase(0x14, anyConstant.as[ldc2_w])
       .singleton(0x6d, ldiv)
       .typecase(0x16, localIndex.as[lload])
       .singleton(0x1e, lload_0)
@@ -236,14 +237,14 @@ object InstructionCodecs {
       .singleton(0x83, lxor)
       .singleton(0xc2, monitorenter)
       .singleton(0xc3, monitorexit)
-      .typecase(0xc5, (constantPoolIndex :: u1).as[multianewarray])
-      .typecase(0xbb, constantPoolIndex.as[_new])
+      .typecase(0xc5, (anyConstant :: u1).as[multianewarray])
+      .typecase(0xbb, anyConstant.as[_new])
       .typecase(0xbc, arrayType.as[newarray])
       .singleton(0x00, nop)
       .singleton(0x57, pop)
       .singleton(0x58, pop2)
-      .typecase(0xb5, constantPoolIndex.as[putfield])
-      .typecase(0xb3, constantPoolIndex.as[putstatic])
+      .typecase(0xb5, anyConstant.as[putfield])
+      .typecase(0xb3, anyConstant.as[putstatic])
       .typecase(0xa9, localIndex.as[ret])
       .singleton(0xb1, _return)
       .singleton(0x35, saload)

core/src/main/scala/org/polyvariant/classfile/ClassFile.scala

@@ -26,9 +26,9 @@ case class ClassFile(
   majorVersion: Int,
   constants: ConstantPool,
   accessFlags: Set[ClassAccessFlag],
-  thisClass: ConstantIndex,
-  superClass: ConstantIndex,
-  interfaces: List[ConstantIndex],
+  thisClass: ConstantIndex[Constant.ClassInfo],
+  superClass: ConstantIndex[Constant.ClassInfo],
+  interfaces: List[ConstantIndex[Constant.ClassInfo]],
   fields: List[FieldInfo],
   methods: List[MethodInfo],
   attributes: List[AttributeInfo],
@@ -37,7 +37,7 @@ case class ClassFile(
 final case class ConstantPool private (private val constants: Array[Constant | Null])
   extends AnyVal {
   // ConstantIndex is 1-based, so we subtract 1
-  def apply(index: ConstantIndex): Constant = constants(index.value - 1)
+  def apply[A <: Constant](index: ConstantIndex[A]): Constant = constants(index.value - 1)
 
   // Size declared in length byte
   def declaredSize: Int = constants.length
@@ -47,7 +47,7 @@ final case class ConstantPool private (private val constants: Array[Constant | N
 
   def constantList: List[Constant] = constants.toList.collect { case c: Constant => c }
 
-  def indexOf(constant: Constant): Option[ConstantIndex] =
+  def indexOf[C <: Constant](constant: C): Option[ConstantIndex[C]] =
     constants.indexOf(constant) match {
       case -1 => None
       case i  => Some(ConstantIndex(i + 1))
@@ -74,7 +74,7 @@ object ConstantPool {
 
 }
 
-case class ConstantIndex(value: Int) {
+case class ConstantIndex[C <: Constant](value: Int) {
 
   def toNarrowEither: Either[this.type, ConstantIndexNarrow] =
     if (value.isValidShort)
@@ -87,7 +87,7 @@ case class ConstantIndex(value: Int) {
 case class ConstantIndexNarrow(value: Short)
 
 object ConstantIndex {
-  val init: ConstantIndex = ConstantIndex(1)
+  val init: ConstantIndex[Nothing] = ConstantIndex(1)
 }
 
 enum Constant {
@@ -98,23 +98,50 @@ enum Constant {
       case _                           => 1
     }
 
-  case ClassInfo(nameIndex: ConstantIndex)
-  case FieldRefInfo(classIndex: ConstantIndex, nameAndTypeIndex: ConstantIndex)
-  case MethodRefInfo(classIndex: ConstantIndex, nameAndTypeIndex: ConstantIndex)
-  case InterfaceMethodRefInfo(classIndex: ConstantIndex, nameAndTypeIndex: ConstantIndex)
-  case StringInfo(stringIndex: ConstantIndex)
+  case ClassInfo(nameIndex: ConstantIndex[Constant.Utf8Info])
+
+  case FieldRefInfo(
+    classIndex: ConstantIndex[Constant.ClassInfo],
+    nameAndTypeIndex: ConstantIndex[Constant],
+  )
+
+  case MethodRefInfo(
+    classIndex: ConstantIndex[Constant.ClassInfo],
+    nameAndTypeIndex: ConstantIndex[Constant],
+  )
+
+  case InterfaceMethodRefInfo(
+    classIndex: ConstantIndex[Constant.ClassInfo],
+    nameAndTypeIndex: ConstantIndex[Constant],
+  )
+
+  case StringInfo(stringIndex: ConstantIndex[Constant.Utf8Info])
   case IntegerInfo(bytes: ByteVector)
   case FloatInfo(bytes: ByteVector)
   case LongInfo(highBytes: ByteVector, lowBytes: ByteVector)
   case DoubleInfo(highBytes: ByteVector, lowBytes: ByteVector)
-  case NameAndTypeInfo(nameIndex: ConstantIndex, descriptorIndex: ConstantIndex)
+
+  case NameAndTypeInfo(
+    nameIndex: ConstantIndex[Constant.Utf8Info],
+    descriptorIndex: ConstantIndex[Constant.Utf8Info],
+  )
+
   case Utf8Info(bytes: String)
-  case MethodHandleInfo(referenceType: MethodReferenceKind, referenceIndex: ConstantIndex)
-  case MethodTypeInfo(descriptorIndex: ConstantIndex)
-  case DynamicInfo(bootstrapMethodAttrIndex: Int, nameAndTypeIndex: ConstantIndex)
-  case InvokeDynamicInfo(bootstrapMethodAttrIndex: Int, nameAndTypeIndex: ConstantIndex)
-  case ModuleInfo(nameIndex: ConstantIndex)
-  case PackageInfo(nameIndex: ConstantIndex)
+  case MethodHandleInfo(referenceType: MethodReferenceKind, referenceIndex: ConstantIndex[Constant])
+  case MethodTypeInfo(descriptorIndex: ConstantIndex[Constant.NameAndTypeInfo])
+
+  case DynamicInfo(
+    bootstrapMethodAttrIndex: Int,
+    nameAndTypeIndex: ConstantIndex[Constant.NameAndTypeInfo],
+  )
+
+  case InvokeDynamicInfo(
+    bootstrapMethodAttrIndex: Int,
+    nameAndTypeIndex: ConstantIndex[Constant.NameAndTypeInfo],
+  )
+
+  case ModuleInfo(nameIndex: ConstantIndex[Constant.Utf8Info])
+  case PackageInfo(nameIndex: ConstantIndex[Constant.Utf8Info])
 }
 
 object Constant {
@@ -142,19 +169,19 @@ enum MethodReferenceKind {
 
 case class FieldInfo(
   accessFlags: Set[FieldAccessFlag],
-  nameIndex: ConstantIndex,
-  descriptorIndex: ConstantIndex,
+  nameIndex: ConstantIndex[Constant.Utf8Info],
+  descriptorIndex: ConstantIndex[Constant.Utf8Info],
   attributes: List[AttributeInfo],
 )
 
-case class AttributeInfo(
-  nameIndex: ConstantIndex,
-  info: ByteVector,
-)
-
 case class MethodInfo(
   accessFlags: Set[MethodAccessFlag],
-  nameIndex: ConstantIndex,
-  descriptorIndex: ConstantIndex,
+  nameIndex: ConstantIndex[Constant.Utf8Info],
+  descriptorIndex: ConstantIndex[Constant.Utf8Info],
   attributes: List[AttributeInfo],
 )
+
+case class AttributeInfo(
+  nameIndex: ConstantIndex[Constant.Utf8Info],
+  info: ByteVector,
+)