Fix datagen sometimes crashing because of concurrent keybind registration (#4771)

Author: screret

src/main/java/com/gregtechceu/gtceu/core/mixins/GTMixinPlugin.java

@@ -26,10 +26,11 @@ public String getRefMapperConfig() {
     private static final String MIXIN_PACKAGE = "com.gregtechceu.gtceu.core.mixins.";
     private static final Map<String, String> MOD_COMPAT_MIXINS = new HashMap<>();
 
-    private static final String DEV_PACKAGE = MIXIN_PACKAGE + "dev.";
+    private static final String DEV_PACKAGE = "dev.";
+    private static final String DATAGEN_PACKAGE = "datagen.";
 
     static {
-        MOD_COMPAT_MIXINS.put("roughlyenoughitems", MIXIN_PACKAGE + "rei");
+        MOD_COMPAT_MIXINS.put("roughlyenoughitems", "rei.");
         addModCompatMixin("emi");
         addModCompatMixin("jei");
         addModCompatMixin("top");
@@ -40,8 +41,24 @@ public String getRefMapperConfig() {
 
     @Override
     public boolean shouldApplyMixin(String targetClassName, String mixinClassName) {
+        if (!mixinClassName.startsWith(MIXIN_PACKAGE)) {
+            // skip checking mixins that aren't in our package
+            // this should never happen, but better safe than sorry
+            return true;
+        }
+        mixinClassName = mixinClassName.substring(MIXIN_PACKAGE.length());
+
         if (mixinClassName.startsWith(DEV_PACKAGE)) {
-            return !FMLLoader.isProduction();
+            if (FMLLoader.isProduction()) {
+                // don't load dev-only mixins in prod
+                return false;
+            }
+            mixinClassName = mixinClassName.substring(DEV_PACKAGE.length());
+            if (mixinClassName.startsWith(DATAGEN_PACKAGE)) {
+                // only load datagen mixins in datagen
+                return FMLLoader.getLaunchHandler().isData();
+            }
+            return true;
         }
         for (var compatMod : MOD_COMPAT_MIXINS.entrySet()) {
             if (mixinClassName.startsWith(compatMod.getValue())) {
@@ -66,7 +83,7 @@ public void preApply(String targetClassName, ClassNode targetClass, String mixin
     public void postApply(String targetClassName, ClassNode targetClass, String mixinClassName, IMixinInfo mixinInfo) {}
 
     private static void addModCompatMixin(String modId) {
-        MOD_COMPAT_MIXINS.put(modId, MIXIN_PACKAGE + modId);
+        MOD_COMPAT_MIXINS.put(modId, modId + ".");
     }
 
     private static boolean isModLoaded(String modId) {

src/main/java/com/gregtechceu/gtceu/core/mixins/dev/datagen/FixParallelKeyMappingRegistrationMixin.java

@@ -0,0 +1,37 @@
+package com.gregtechceu.gtceu.core.mixins.dev.datagen;
+
+import net.minecraft.client.KeyMapping;
+import net.minecraftforge.client.settings.KeyMappingLookup;
+
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
+import com.mojang.blaze3d.platform.InputConstants;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+
+import java.util.EnumMap;
+import java.util.List;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * This mixin is only loaded in datagen contexts
+ */
+@Mixin(value = KeyMappingLookup.class, remap = false)
+public class FixParallelKeyMappingRegistrationMixin {
+
+    /**
+     * @author screret
+     * @reason Fix keybind registration in datagen.<br>
+     *         <p>
+     *         The map isn't thread safe and some mods register their keybinds in places they probably shouldn't.<br>
+     *         (IDK why it only crashes in datagen, though)
+     */
+    @WrapOperation(method = "<clinit>",
+                   at = @At(value = "INVOKE",
+                            target = "Ljava/util/EnumMap;put(Ljava/lang/Enum;Ljava/lang/Object;)Ljava/lang/Object;"))
+    private static <K extends Enum<K>, V> V gtceu$makeKeyMappingMapsConcurrent(EnumMap<K, V> instance, K key, V value,
+                                                                               Operation<V> original) {
+        // replace `value`, which is a non-threadsafe HashMap, with a ConcurrentHashMap
+        return original.call(instance, key, new ConcurrentHashMap<InputConstants.Key, List<KeyMapping>>());
+    }
+}

src/main/resources/gtceu.mixins.json

@@ -63,6 +63,7 @@
         "TagManagerMixin",
         "TagValueAccessor",
         "client.ItemEntityMixin",
+        "dev.datagen.FixParallelKeyMappingRegistrationMixin",
         "dev.test.GameTestRegistryMixin",
         "emi.EmiRecipeFillerMixin",
         "emi.FillRecipePacketMixin",