Strip JST-Injected Enum Extensions

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 {}

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,22 +92,79 @@ 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("<clinit>"));
+        Optional<MethodNode> $valuesOpt = tryFindMethod(classNode, mth -> mth.name.equals("$values"));
+        boolean $valuesPresent = $valuesOpt.isPresent();
+
+        boolean requiresRewrite = false;
+        List<FieldNode> 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 <clinit>
+
+            var dropValuesGenerator = new ListGeneratorAdapter(new InsnList());
+            removeFieldsValuesArray(classType, dropValuesGenerator, vanillaEntryCount);
+
+            // Remove construction and field store of the relevant fields in <clinit>
+            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<EnumPrototype> 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("<clinit>"));
-        Optional<MethodNode> $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<String> ctors = classNode.methods.stream()
                 .filter(mth -> mth.name.equals("<init>"))
                 .filter(RuntimeEnumExtender::isAllowedConstructor)
                 .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<FieldNode> enumEntriesToRemove) {
+        Set<String> 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<FieldNode> enumEntries) {
         // values = Arrays.copyOf(values, values.length + listSize);
         generator.dup();

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
+    <T extends Enum<T>> 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<T> 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
     <T extends Enum<T>> void testExtendEnum() throws Exception {
         installation.setupProductionClient();