diff --git a/loader/src/main/java/net/neoforged/fml/common/asm/enumextension/ExtensionEnumEntry.java b/loader/src/main/java/net/neoforged/fml/common/asm/enumextension/ExtensionEnumEntry.java new file mode 100644 index 000000000..710a412c8 --- /dev/null +++ b/loader/src/main/java/net/neoforged/fml/common/asm/enumextension/ExtensionEnumEntry.java @@ -0,0 +1,21 @@ +/* + * Copyright (c) NeoForged and contributors + * SPDX-License-Identifier: LGPL-2.1-only + */ + +package net.neoforged.fml.common.asm.enumextension; + +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; +import org.jetbrains.annotations.ApiStatus; + +/** + * This annotation is added by JST to mark enum entries added by mod enum extensions. This allows for consistent ordering + * of enum entries whether they are added with JST in source, or at runtime. + */ +@ApiStatus.Internal +@Retention(RetentionPolicy.CLASS) +@Target(ElementType.FIELD) +public @interface ExtensionEnumEntry {} diff --git a/loader/src/main/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtender.java b/loader/src/main/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtender.java index e413dbcf7..26492212b 100644 --- a/loader/src/main/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtender.java +++ b/loader/src/main/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtender.java @@ -49,6 +49,7 @@ public class RuntimeEnumExtender implements ClassProcessor { private static final Type INDEXED_ANNOTATION = Type.getType(IndexedEnum.class); private static final Type NAMED_ANNOTATION = Type.getType(NamedEnum.class); private static final Type RESERVED_ANNOTATION = Type.getType(ReservedConstructor.class); + private static final Type JST_ADDED_ENTRY = Type.getType(ExtensionEnumEntry.class); private static final Type ENUM_PROXY = Type.getType(EnumProxy.class); private static final Type NET_CHECK = Type.getType(NetworkedEnum.NetworkCheck.class); private static final Type EXT_INFO = Type.getType(ExtensionInfo.class); @@ -91,14 +92,72 @@ public ComputeFlags processClass(TransformationContext context) { throw new IllegalStateException("Tried to extend non-enum class or non-extensible enum: " + classType); } + MethodNode clinit = findMethod(classNode, mth -> mth.name.equals("")); + Optional $valuesOpt = tryFindMethod(classNode, mth -> mth.name.equals("$values")); + boolean $valuesPresent = $valuesOpt.isPresent(); + + boolean requiresRewrite = false; + List enumEntriesToRemove = new ArrayList<>(); + for (var field : classNode.fields) { + if ((field.access & Opcodes.ACC_ENUM) != 0 && field.desc.equals(classType.getDescriptor())) { + if (field.invisibleAnnotations != null && field.invisibleAnnotations.stream().anyMatch(anno -> anno.desc.equals(JST_ADDED_ENTRY.getDescriptor()))) { + // We have to remove this field everywhere! + enumEntriesToRemove.add(field); + } + } + } + if (!enumEntriesToRemove.isEmpty()) { + requiresRewrite = true; + classNode.fields.removeAll(enumEntriesToRemove); + } + + int vanillaEntryCount = getVanillaEntryCount(classNode, classType); + + if (!enumEntriesToRemove.isEmpty()) { + // Remove enum entries from + + var dropValuesGenerator = new ListGeneratorAdapter(new InsnList()); + removeFieldsValuesArray(classType, dropValuesGenerator, vanillaEntryCount); + + // Remove construction and field store of the relevant fields in + for (FieldNode field : enumEntriesToRemove) { + var putStaticInsn = findFirstFieldAccess(clinit, Opcodes.PUTSTATIC, classType.getInternalName(), field.name, classType.getDescriptor()); + if (putStaticInsn == null) { + continue; + // This could happen if the MC jar was processed by InstallerTools, which adds the enum fields but no initializers for them + } + AbstractInsnNode insnToRemove = findFirstInstructionBefore(clinit, Opcodes.NEW, clinit.instructions.indexOf(putStaticInsn)); + // Drop all of these (inclusive) + while (insnToRemove != putStaticInsn && insnToRemove != null) { + var next = insnToRemove.getNext(); + clinit.instructions.remove(insnToRemove); + insnToRemove = next; + } + clinit.instructions.remove(putStaticInsn); + } + + // Remove enum entries from values array + if ($valuesPresent) { // javac + MethodNode $values = $valuesOpt.get(); + AbstractInsnNode $valuesAretInsn = findFirstInstructionBefore($values, Opcodes.ARETURN, $values.instructions.size() - 1); + $values.instructions.insertBefore($valuesAretInsn, dropValuesGenerator.insnList); + + // Replace any GETSTATIC of the relevant fields with ACONST_NULL + removeRemovedFieldLoads(classType, $values, enumEntriesToRemove); + } else { // ECJ + AbstractInsnNode putStaticInsn = findValuesArrayStore(classType, classNode, clinit, classType.getInternalName()); + clinit.instructions.insertBefore(putStaticInsn, dropValuesGenerator.insnList); + + // Replace any GETSTATIC of the relevant fields with ACONST_NULL + removeRemovedFieldLoads(classType, clinit, enumEntriesToRemove); + } + } + List protos = prototypes.getOrDefault(classType.getInternalName(), List.of()); if (protos.isEmpty()) { - return ComputeFlags.NO_REWRITE; + return requiresRewrite ? ComputeFlags.COMPUTE_FRAMES : ComputeFlags.NO_REWRITE; } - MethodNode clinit = findMethod(classNode, mth -> mth.name.equals("")); - Optional $valuesOpt = tryFindMethod(classNode, mth -> mth.name.equals("$values")); - boolean $valuesPresent = $valuesOpt.isPresent(); MethodNode getExtInfo = findMethod(classNode, mth -> mth.name.equals("getExtensionInfo") && mth.desc.equals(EXT_INFO_GETTER_DESC)); Set ctors = classNode.methods.stream() .filter(mth -> mth.name.equals("")) @@ -106,7 +165,6 @@ public ComputeFlags processClass(TransformationContext context) { .map(mth -> mth.desc) .collect(Collectors.toSet()); - int vanillaEntryCount = getVanillaEntryCount(classNode, classType); int idParamIdx = getParameterIndexFromAnnotation(classNode, INDEXED_ANNOTATION); int nameParamIdx = getParameterIndexFromAnnotation(classNode, NAMED_ANNOTATION); @@ -177,6 +235,30 @@ public static MethodInsnNode findFirstStaticMethodCall(MethodNode method, String return null; } + /** + * Finds the first field access in the given method matching the given opcode, owner, name and descriptor + * + * @param method the method to search in + * @param opcode the opcode of the field access to search for + * @param owner the field access's owner to search for + * @param name the field access's name + * @param descriptor the field access's descriptor + * @return the found field access node, null if none matched + */ + public static FieldInsnNode findFirstFieldAccess(MethodNode method, int opcode, String owner, String name, String descriptor) { + for (int i = 0; i < method.instructions.size(); i++) { + AbstractInsnNode node = method.instructions.get(i); + if (node instanceof FieldInsnNode fieldInsnNode && node.getOpcode() == opcode) { + if (fieldInsnNode.owner.equals(owner) && + fieldInsnNode.name.equals(name) && + fieldInsnNode.desc.equals(descriptor)) { + return fieldInsnNode; + } + } + } + return null; + } + /** * Finds the first instruction with matching opcode before the given index in reverse search * @@ -474,6 +556,24 @@ private static void returnValuesToExtender(Type classType, ListGeneratorAdapter } } + private void removeRemovedFieldLoads(Type classType, MethodNode targetMethod, List enumEntriesToRemove) { + Set removedFieldNames = enumEntriesToRemove.stream().map(field -> field.name).collect(Collectors.toSet()); + for (int i = 0; i < targetMethod.instructions.size(); i++) { + var insn = targetMethod.instructions.get(i); + if (insn instanceof FieldInsnNode fieldInsn && fieldInsn.getOpcode() == Opcodes.GETSTATIC && fieldInsn.owner.equals(classType.getInternalName()) && removedFieldNames.contains(fieldInsn.name)) { + // Replace with null load + targetMethod.instructions.set(insn, new InsnNode(Opcodes.ACONST_NULL)); + } + } + } + + private static void removeFieldsValuesArray(Type classType, ListGeneratorAdapter generator, int vanillaEntryCount) { + generator.dup(); + generator.push(vanillaEntryCount); + generator.invokeStatic(ARRAYS, new Method("copyOf", "([Ljava/lang/Object;I)[Ljava/lang/Object;")); + generator.checkCast(Type.getType("[" + classType.getDescriptor())); + } + private static void appendValuesArray(Type classType, ListGeneratorAdapter generator, List enumEntries) { // values = Arrays.copyOf(values, values.length + listSize); generator.dup(); diff --git a/loader/src/test/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtenderTest.java b/loader/src/test/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtenderTest.java index 75f8cdb82..feb638cfb 100644 --- a/loader/src/test/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtenderTest.java +++ b/loader/src/test/java/net/neoforged/fml/common/asm/enumextension/RuntimeEnumExtenderTest.java @@ -31,6 +31,49 @@ void testMissingPath() throws Exception { "ERROR: Enum extender file xyz, provided by mod testmod, does not exist"); } + @Test + > void testStripsJstEntries() throws Exception { + installation.setupProductionClient(); + + installation.buildModJar("enum_ext_test.jar") + .withModsToml(getModsTomlBuilderConsumer("extensions.json")) + .addTextFile("extensions.json", """ + { + "entries": [ + { + "enum": "testmod/TestEnum", + "name": "TESTMOD_NEW_CONSTANT", + "constructor": "()V", + "parameters": [] + } + ] + } + """) + .addClass("testmod.TestEnum", """ + @net.neoforged.fml.common.asm.enumextension.NamedEnum + public enum TestEnum implements net.neoforged.fml.common.asm.enumextension.IExtensibleEnum { + TEST_THING, + @net.neoforged.fml.common.asm.enumextension.ExtensionEnumEntry TESTMOD_NEW_CONSTANT; + public static net.neoforged.fml.common.asm.enumextension.ExtensionInfo getExtensionInfo() { + return net.neoforged.fml.common.asm.enumextension.ExtensionInfo.nonExtended(TestEnum.class); + } + } + """) + .build(); + launchAndLoad("neoforgeclient"); + + Class testEnum = getEnumClass("testmod.TestEnum"); + var injectedField = testEnum.getField("TESTMOD_NEW_CONSTANT"); + assertThat(injectedField.getDeclaredAnnotation(ExtensionEnumEntry.class)).isNull(); + var enumConstants = testEnum.getEnumConstants(); + assertThat(enumConstants.length).isEqualTo(2); + var enumEntry = enumConstants[1]; + var byName = Enum.valueOf(testEnum, "TESTMOD_NEW_CONSTANT"); + assertThat(enumEntry).isSameAs(byName); + assertThat(enumEntry.name()).isEqualTo("TESTMOD_NEW_CONSTANT"); + assertThat(enumEntry.ordinal()).isEqualTo(1); + } + @Test > void testExtendEnum() throws Exception { installation.setupProductionClient();