From 2cc683f39d12e5c6ed8f2ad5f00d4f1a182f5c2e Mon Sep 17 00:00:00 2001
From: embeddedt <42941056+embeddedt@users.noreply.github.com>
Date: Sat, 27 Dec 2025 19:33:30 -0500
Subject: [PATCH] Make maps provided to ModifyBakingResult mutable

---
 .../dynresources/DynamicModelSystem.java      |   2 +-
 .../dynresources/DynamicRegistryMap.java      | 161 ++++++++++++++++++
 2 files changed, 162 insertions(+), 1 deletion(-)
 create mode 100644 src/main/java/org/embeddedt/modernfix/dynresources/DynamicRegistryMap.java

diff --git a/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java b/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java
index 0a1c8180..3bb117e3 100644
--- a/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java
+++ b/src/main/java/org/embeddedt/modernfix/dynresources/DynamicModelSystem.java
@@ -172,7 +172,7 @@ public class DynamicModelSystem {
                 }
             }
         });
-        return Maps.asMap(input.keySet(), k -> {
+        return new DynamicRegistryMap<>(input.keySet(),k -> {
             if (k != null) {
                 Object value = bakedCache.getUnchecked(k);
                 if (value == NULL_BAKED) {
diff --git a/src/main/java/org/embeddedt/modernfix/dynresources/DynamicRegistryMap.java b/src/main/java/org/embeddedt/modernfix/dynresources/DynamicRegistryMap.java
new file mode 100644
index 00000000..65c8e6ef
--- /dev/null
+++ b/src/main/java/org/embeddedt/modernfix/dynresources/DynamicRegistryMap.java
@@ -0,0 +1,161 @@
+package org.embeddedt.modernfix.dynresources;
+
+import com.google.common.collect.Collections2;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.AbstractSet;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+/**
+ * A map that behaves like Guava's Maps.asMap but allows for additional entries to be written that override the backing
+ * map's entries.
+ */
+public final class DynamicRegistryMap<K, V> implements Map<K, V> {
+    private static final Object NULL_OVERRIDE = new Object();
+    private final Set<K> originalKeys;
+    private final Function<K, V> fallbackGetter;
+    private final ConcurrentHashMap<K, Object> overrides;
+    private final EntrySet entrySet;
+
+    public DynamicRegistryMap(Set<K> originalKeys, Function<K, V> fallbackGetter) {
+        this.originalKeys = originalKeys;
+        this.fallbackGetter = fallbackGetter;
+        this.overrides = new ConcurrentHashMap<>();
+        this.entrySet = new EntrySet();
+    }
+
+    @Override
+    public int size() {
+        return originalKeys.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return originalKeys.isEmpty();
+    }
+
+    @Override
+    public boolean containsKey(Object o) {
+        if (o == null) {
+            return false;
+        }
+        var override = overrides.get(o);
+        if (override == NULL_OVERRIDE) {
+            return false;
+        }
+        return override != null || originalKeys.contains(o);
+    }
+
+    @Override
+    public boolean containsValue(Object o) {
+        if (o == null || o == NULL_OVERRIDE) {
+            return false;
+        }
+        return overrides.containsValue(o);
+    }
+
+    @Override
+    public V get(Object o) {
+        Object value = overrides.get(o);
+        if (value == NULL_OVERRIDE) {
+            return null;
+        } else if (value != null) {
+            return (V) value;
+        } else {
+            return fallbackGetter.apply((K)o);
+        }
+    }
+
+    @Override
+    public @Nullable V put(K k, V v) {
+        if (v == null) {
+            return remove(k);
+        }
+        overrides.put(k, v);
+        return null;
+    }
+
+    @Override
+    public V remove(Object o) {
+        overrides.put((K)o, NULL_OVERRIDE);
+        return null;
+    }
+
+    @Override
+    public void putAll(@NotNull Map<? extends K, ? extends V> map) {
+        map.forEach(this::put);
+    }
+
+    @Override
+    public void clear() {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public @NotNull Set<K> keySet() {
+        return Collections.unmodifiableSet(originalKeys);
+    }
+
+    @Override
+    public @NotNull Collection<V> values() {
+        return Collections2.transform(originalKeys, this::get);
+    }
+
+    @Override
+    public @NotNull Set<Entry<K, V>> entrySet() {
+        return this.entrySet;
+    }
+
+    private class ModelEntry implements Map.Entry<K, V> {
+        private final K key;
+
+        private ModelEntry(K key) {
+            this.key = key;
+        }
+
+        @Override
+        public K getKey() {
+            return key;
+        }
+
+        @Override
+        public V getValue() {
+            return get(key);
+        }
+
+        @Override
+        public V setValue(V value) {
+            return put(key, value);
+        }
+    }
+
+    private class EntrySet extends AbstractSet<Map.Entry<K, V>> {
+        @Override
+        public Iterator<Entry<K, V>> iterator() {
+            var iterator = originalKeys.iterator();
+            return new Iterator<>() {
+                @Override
+                public boolean hasNext() {
+                    return iterator.hasNext();
+                }
+
+                @Override
+                public Entry<K, V> next() {
+                    return new ModelEntry(iterator.next());
+                }
+            };
+        }
+
+        @Override
+        public int size() {
+            return DynamicRegistryMap.this.size();
+        }
+    }
+}