Make CraftTweaker loot modifiers always fire last

Signed-off-by: TheSilkMiner <thesilkminer@outlook.com>
CraftTweaker · Sep 11, 2021 · 0249573 · 0249573
1 parent ab180a9
commit 0249573
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 5 deletions.
diff --git a/src/main/java/com/blamejared/crafttweaker/impl/loot/modifiers/CTLootModifierManager.java b/src/main/java/com/blamejared/crafttweaker/impl/loot/modifiers/CTLootModifierManager.java
@@ -11,7 +11,6 @@
 import com.blamejared.crafttweaker.impl.loot.conditions.CTLootConditionBuilder;
 import com.blamejared.crafttweaker.impl.util.NameUtils;
 import com.blamejared.crafttweaker_annotations.annotations.Document;
-import com.google.common.collect.ImmutableMap;
 import net.minecraft.util.ResourceLocation;
 import net.minecraftforge.common.ForgeInternalHandler;
 import net.minecraftforge.common.loot.IGlobalLootModifier;
@@ -20,7 +19,6 @@
 
 import java.lang.invoke.MethodHandle;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -195,9 +193,10 @@ private Map<ResourceLocation, IGlobalLootModifier> getLmmMap() {
         try {
             final LootModifierManager lmm = MethodHandleHelper.invoke(() -> (LootModifierManager) LMM_GETTER.invokeExact());
             Map<ResourceLocation, IGlobalLootModifier> map = MethodHandleHelper.invoke(() -> (Map<ResourceLocation, IGlobalLootModifier>) LMM_MAP_GETTER.invokeExact(lmm));
-            if (map instanceof ImmutableMap) {
-                map = new HashMap<>(map);
-                final Map<ResourceLocation, IGlobalLootModifier> finalMap = map;
+            // Someone else may make the map mutable, but I explicitly want CT stuff to go last
+            if (!(map instanceof CraftTweakerPrivilegedLootModifierMap)) {
+                final Map<ResourceLocation, IGlobalLootModifier> finalMap = new CraftTweakerPrivilegedLootModifierMap(map);
+                map = finalMap;
                 MethodHandleHelper.invokeVoid(() -> this.setLmmMap(lmm, finalMap)); // Let's "mutabilize" the map
             }
             return map;

diff --git a/...om/blamejared/crafttweaker/impl/loot/modifiers/CraftTweakerPrivilegedLootModifierMap.java b/...om/blamejared/crafttweaker/impl/loot/modifiers/CraftTweakerPrivilegedLootModifierMap.java
@@ -0,0 +1,83 @@
+package com.blamejared.crafttweaker.impl.loot.modifiers;
+
+import com.blamejared.crafttweaker.CraftTweaker;
+import net.minecraft.util.ResourceLocation;
+import net.minecraftforge.common.loot.IGlobalLootModifier;
+
+import java.util.AbstractMap;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.Set;
+
+final class CraftTweakerPrivilegedLootModifierMap extends AbstractMap<ResourceLocation, IGlobalLootModifier> {
+    private final Map<ResourceLocation, IGlobalLootModifier> defaultMap;
+    private final Map<ResourceLocation, IGlobalLootModifier> ctMap;
+
+    CraftTweakerPrivilegedLootModifierMap(final Map<ResourceLocation, IGlobalLootModifier> clone) {
+        if (clone instanceof CraftTweakerPrivilegedLootModifierMap) {
+            final CraftTweakerPrivilegedLootModifierMap other = (CraftTweakerPrivilegedLootModifierMap) clone;
+            this.defaultMap = new HashMap<>(other.defaultMap);
+            this.ctMap = new LinkedHashMap<>(other.ctMap);
+        } else {
+            this.defaultMap = new HashMap<>();
+            this.ctMap = new LinkedHashMap<>();
+            this.putAll(clone);
+        }
+    }
+
+    private static boolean isCtModifier(final ResourceLocation key) {
+        return CraftTweaker.MODID.equals(key.getNamespace());
+    }
+
+    @Override
+    public int size() {
+        return this.defaultMap.size() + this.ctMap.size();
+    }
+
+    @Override
+    public boolean containsValue(final Object value) {
+        return this.defaultMap.containsValue(value) || this.ctMap.containsValue(value);
+    }
+
+    @Override
+    public boolean containsKey(final Object key) {
+        return this.defaultMap.containsKey(key) || this.ctMap.containsKey(key);
+    }
+
+    @Override
+    public IGlobalLootModifier get(final Object key) {
+        final IGlobalLootModifier v = this.defaultMap.get(key);
+        return v == null? this.ctMap.get(key) : v;
+    }
+
+    @Override
+    public IGlobalLootModifier put(final ResourceLocation key, final IGlobalLootModifier value) {
+        if(isCtModifier(key)) {
+            return this.ctMap.put(key, value);
+        } else {
+            return this.defaultMap.put(key, value);
+        }
+    }
+
+    @Override
+    public IGlobalLootModifier remove(final Object key) {
+        final IGlobalLootModifier v = this.defaultMap.remove(key);
+        return v == null? this.ctMap.remove(key) : v;
+    }
+
+    @Override
+    public void clear() {
+        this.defaultMap.clear();
+        this.ctMap.clear();
+    }
+
+    @Override
+    public Set<Entry<ResourceLocation, IGlobalLootModifier>> entrySet() {
+        final Set<Entry<ResourceLocation, IGlobalLootModifier>> set = new LinkedHashSet<>(this.size());
+        set.addAll(this.defaultMap.entrySet());
+        set.addAll(this.ctMap.entrySet());
+        return set;
+    }
+}