Optimize NBT memory usage

src/main/java/org/embeddedt/modernfix/core/config/ModernFixEarlyConfig.java

@@ -83,6 +83,7 @@ public class ModernFixEarlyConfig {
         this.addMixinRule("perf.deduplicate_location", false);
         this.addMixinRule("perf.cache_blockstate_cache_arrays", true);
         this.addMixinRule("perf.cache_model_materials", true);
+        this.addMixinRule("perf.nbt_memory_usage", true);
         this.addMixinRule("perf.patchouli_deduplicate_books", modPresent("patchouli"));
         this.addMixinRule("perf.datapack_reload_exceptions", true);
         this.addMixinRule("perf.async_locator", true);

src/main/java/org/embeddedt/modernfix/mixin/perf/nbt_memory_usage/CompoundTagMixin.java

@@ -0,0 +1,45 @@
+package org.embeddedt.modernfix.mixin.perf.nbt_memory_usage;
+
+import com.google.common.collect.Maps;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.nbt.Tag;
+import org.embeddedt.modernfix.util.CanonizingStringMap;
+import org.spongepowered.asm.mixin.*;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.ModifyArg;
+import org.spongepowered.asm.mixin.injection.Redirect;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import java.util.Map;
+
+@Mixin(CompoundTag.class)
+public class CompoundTagMixin {
+    @Shadow @Final @Mutable
+    private Map<String, Tag> tags;
+
+    /**
+     * Ensure that the backing map is always a CanonizingStringMap.
+     */
+    @Redirect(method = "<init>(Ljava/util/Map;)V", at = @At(value = "FIELD", target = "Lnet/minecraft/nbt/CompoundTag;tags:Ljava/util/Map;", ordinal = 0))
+    private void replaceTagMap(CompoundTag tag, Map<String, Tag> incomingMap) {
+        if(incomingMap instanceof CanonizingStringMap)
+            this.tags = incomingMap;
+        else {
+            this.tags = new CanonizingStringMap<>();
+            this.tags.putAll(incomingMap);
+        }
+    }
+
+    /**
+     * @author embeddedt
+     * @reason use more efficient method when copying canonizing string map
+     */
+    @Inject(method = "copy()Lnet/minecraft/nbt/Tag;", at = @At("HEAD"), cancellable = true)
+    public void copyEfficient(CallbackInfoReturnable<Tag> cir) {
+        if(this.tags instanceof CanonizingStringMap) {
+            cir.setReturnValue(new CompoundTag(CanonizingStringMap.deepCopy((CanonizingStringMap<Tag>)this.tags, Tag::copy)));
+        }
+    }
+}

src/main/java/org/embeddedt/modernfix/util/CanonizingStringMap.java

@@ -0,0 +1,147 @@
+package org.embeddedt.modernfix.util;
+
+import com.google.common.collect.Interner;
+import com.google.common.collect.Interners;
+import it.unimi.dsi.fastutil.objects.*;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.*;
+import java.util.function.Function;
+
+/**
+ * Replacement backing map for CompoundTags. Uses an array map for tags with 4 or less entries,
+ * and a hash map for larger tags.
+ */
+public class CanonizingStringMap<T> implements Map<String, T> {
+    private Object2ObjectMap<String, T> backingMap;
+
+    private static final int GROWTH_THRESHOLD = 4;
+    private static final Interner<String> KEY_INTERNER = Interners.newStrongInterner();
+
+    public CanonizingStringMap() {
+        this(new Object2ObjectArrayMap<>());
+    }
+
+    protected CanonizingStringMap(Object2ObjectMap<String, T> newMap) {
+        this.backingMap = newMap;
+    }
+
+    @Override
+    public int size() {
+        return backingMap.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return backingMap.isEmpty();
+    }
+
+    @Override
+    public boolean containsKey(Object o) {
+        return backingMap.containsKey(o);
+    }
+
+    @Override
+    public boolean containsValue(Object o) {
+        return backingMap.containsValue(o);
+    }
+
+    @Override
+    public T get(Object o) {
+        return backingMap.get(o);
+    }
+
+    @Nullable
+    @Override
+    public T put(String s, T t) {
+        if(backingMap.size() >= GROWTH_THRESHOLD && !(backingMap instanceof Object2ObjectOpenHashMap) && !backingMap.containsKey(s)) {
+            // map will grow to GROWTH_THRESHOLD + 1 entries, change to hashmap
+            backingMap = new Object2ObjectOpenHashMap<>(backingMap);
+        }
+        s = KEY_INTERNER.intern(s);
+        return backingMap.put(s, t);
+    }
+
+    @Override
+    public T remove(Object o) {
+        T value = backingMap.remove(o);
+        // need to shrink to be consistent with new maps
+        if(backingMap.size() <= GROWTH_THRESHOLD && backingMap instanceof Object2ObjectOpenHashMap) {
+            backingMap = new Object2ObjectArrayMap<>(backingMap);
+        }
+        return value;
+    }
+
+    @Override
+    public void putAll(@NotNull Map<? extends String, ? extends T> map) {
+        if(map.size() == 0)
+            return;
+        map.forEach((String key, T val) -> {
+            key = KEY_INTERNER.intern(key);
+            backingMap.put(key, val);
+        });
+        // if it's too big to be an array, grow it
+        if(backingMap.size() > GROWTH_THRESHOLD && !(backingMap instanceof Object2ObjectOpenHashMap)) {
+            backingMap = new Object2ObjectOpenHashMap<>(backingMap);
+        }
+    }
+
+    @Override
+    public void clear() {
+        if(!(this.backingMap instanceof Object2ObjectArrayMap))
+            this.backingMap = new Object2ObjectArrayMap<>();
+        else
+            this.backingMap.clear();
+    }
+
+    @NotNull
+    @Override
+    public Set<String> keySet() {
+        return Collections.unmodifiableSet(this.backingMap.keySet());
+    }
+
+    @NotNull
+    @Override
+    public Collection<T> values() {
+        return Collections.unmodifiableCollection(this.backingMap.values());
+    }
+
+    @NotNull
+    @Override
+    public Set<Entry<String, T>> entrySet() {
+        return Collections.unmodifiableSet(this.backingMap.entrySet());
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        CanonizingStringMap<?> that = (CanonizingStringMap<?>)o;
+        if(that.backingMap.size() != backingMap.size())
+            return false;
+        return backingMap.object2ObjectEntrySet().containsAll(that.backingMap.object2ObjectEntrySet());
+    }
+
+    /**
+     * We deliberately use a hashcode that will be consistent regardless of underlying map type.
+     */
+    @Override
+    public int hashCode() {
+        final ObjectIterator<Object2ObjectMap.Entry<String, T>> i = Object2ObjectMaps.fastIterator(backingMap);
+        int h = 0, n = backingMap.size();
+        while (n-- != 0)
+            h += i.next().hashCode();
+        return h;
+    }
+
+    public static <T> CanonizingStringMap<T> deepCopy(CanonizingStringMap<T> inputMap, Function<T, T> deepCopier) {
+        Object2ObjectMap<String, T> copiedBackingMap;
+        if(inputMap.backingMap instanceof Object2ObjectArrayMap)
+            copiedBackingMap = ((Object2ObjectArrayMap<String, T>)inputMap.backingMap).clone();
+        else
+            copiedBackingMap = ((Object2ObjectOpenHashMap<String, T>)inputMap.backingMap).clone();
+        copiedBackingMap.replaceAll((k, v) -> deepCopier.apply(v));
+        return new CanonizingStringMap<>(copiedBackingMap);
+    }
+}

src/main/resources/META-INF/accesstransformer.cfg

@@ -23,4 +23,5 @@ public net.minecraft.block.AbstractBlock$Properties field_235818_t_ # hasPostPro
 public net.minecraft.block.AbstractBlock$Properties field_235819_u_ # emissiveRendering
 public net.minecraft.block.AbstractBlock$Properties field_235806_h_ # requiresCorrectToolForDrops
 public net.minecraft.block.AbstractBlock$Properties field_200959_g # destroyTime
-public net.minecraft.world.server.ServerChunkProvider$ChunkExecutor
+public net.minecraft.world.server.ServerChunkProvider$ChunkExecutor
+public net.minecraft.nbt.CompoundNBT <init>(Ljava/util/Map;)V # <init>

src/main/resources/modernfix.mixins.json

@@ -59,6 +59,7 @@
     "perf.kubejs.RecipeJSMixin",
     "perf.kubejs.IDFilterMixin",
     "perf.kubejs.CustomIngredientMixin",
+    "perf.nbt_memory_usage.CompoundTagMixin",
     "perf.fast_registry_validation.ForgeRegistryMixin",
     "perf.cache_strongholds.ChunkGeneratorMixin",
     "perf.cache_upgraded_structures.StructureManagerMixin",