Attempt to improve parity/reliability of dynamic_entity_renderers

Author: embeddedt

common/src/main/java/org/embeddedt/modernfix/entity/EntityRendererMap.java

@@ -3,6 +3,8 @@
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.Iterators;
+import com.google.common.collect.Maps;
 import net.minecraft.client.renderer.entity.EntityRenderer;
 import net.minecraft.client.renderer.entity.EntityRendererProvider;
 import net.minecraft.client.renderer.entity.EntityRenderers;
@@ -12,14 +14,20 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.AbstractCollection;
+import java.util.AbstractSet;
 import java.util.Collection;
+import java.util.Iterator;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 
+@SuppressWarnings("OptionalAssignedToNull")
 public class EntityRendererMap implements Map<EntityType<?>, EntityRenderer<?>> {
     private final Map<EntityType<?>, EntityRendererProvider<?>> rendererProviders;
-    private final LoadingCache<EntityType<?>, EntityRenderer<?>> rendererMap;
+    private final LoadingCache<EntityType<?>, Optional<EntityRenderer<?>>> rendererMap;
     private final EntityRendererProvider.Context context;
 
     public EntityRendererMap(Map<EntityType<?>, EntityRendererProvider<?>> rendererProviders, EntityRendererProvider.Context context) {
@@ -28,22 +36,22 @@ public EntityRendererMap(Map<EntityType<?>, EntityRendererProvider<?>> rendererP
         this.rendererMap = CacheBuilder.newBuilder().build(new RenderConstructor());
     }
 
-    class RenderConstructor extends CacheLoader<EntityType<?>, EntityRenderer<?>> {
+    class RenderConstructor extends CacheLoader<EntityType<?>, Optional<EntityRenderer<?>>> {
         @Override
-        public EntityRenderer<?> load(EntityType<?> key) throws Exception {
+        public Optional<EntityRenderer<?>> load(EntityType<?> key) throws Exception {
             EntityRendererProvider<?> provider = rendererProviders.get(key);
+            if(provider == null)
+                return Optional.empty();
             synchronized(EntityRenderers.class) {
                 EntityRenderer<?> renderer;
                 try {
-                    if(provider == null)
-                        throw new RuntimeException("Provider not registered");
                     renderer = provider.create(context);
                     ModernFix.LOGGER.info("Loaded entity {}", BuiltInRegistries.ENTITY_TYPE.getKey(key));
                 } catch(RuntimeException e) {
                     ModernFix.LOGGER.error("Failed to create entity model for " + BuiltInRegistries.ENTITY_TYPE.getKey(key) + ":", e);
                     renderer = new ErroredEntityRenderer<>(context);
                 }
-                return renderer;
+                return Optional.ofNullable(renderer);
             }
         }
     }
@@ -71,10 +79,8 @@ public boolean containsValue(Object o) {
     @Override
     public EntityRenderer<?> get(Object o) {
         try {
-            EntityRenderer<?> renderer = rendererMap.get((EntityType<?>)o);
-            if(renderer == null)
-                throw new AssertionError("Returned entity renderer should never be null");
-            return renderer;
+            Optional<EntityRenderer<?>> renderer = rendererMap.get((EntityType<?>)o);
+            return renderer.orElse(null);
         } catch (IllegalStateException e) {
             return null; /* emulate value not being present if recursive load occurs */
         } catch (ExecutionException e) {
@@ -85,21 +91,21 @@ public EntityRenderer<?> get(Object o) {
     @Nullable
     @Override
     public EntityRenderer<?> put(EntityType<?> entityType, EntityRenderer<?> entityRenderer) {
-        EntityRenderer<?> old = rendererMap.getIfPresent(entityType);
-        rendererMap.put(entityType, entityRenderer);
-        return old;
+        Optional<EntityRenderer<?>> old = rendererMap.getIfPresent(entityType);
+        rendererMap.put(entityType, Optional.ofNullable(entityRenderer));
+        return old != null ? old.orElse(null) : null;
     }
 
     @Override
     public EntityRenderer<?> remove(Object o) {
-        EntityRenderer<?> r = rendererMap.getIfPresent(o);
+        Optional<EntityRenderer<?>> old = rendererMap.getIfPresent(o);
         rendererMap.invalidate(o);
-        return r;
+        return old != null ? old.orElse(null) : null;
     }
 
     @Override
     public void putAll(@NotNull Map<? extends EntityType<?>, ? extends EntityRenderer<?>> map) {
-        rendererMap.putAll(map);
+        rendererMap.putAll(Maps.transformValues(map, Optional::ofNullable));
     }
 
     @Override
@@ -116,12 +122,64 @@ public Set<EntityType<?>> keySet() {
     @NotNull
     @Override
     public Collection<EntityRenderer<?>> values() {
-        return rendererMap.asMap().values();
+        return new AbstractCollection<>() {
+            @Override
+            public Iterator<EntityRenderer<?>> iterator() {
+                return Iterators.transform(Iterators.unmodifiableIterator(rendererProviders.keySet().iterator()), EntityRendererMap.this::get);
+            }
+
+            @Override
+            public int size() {
+                return rendererProviders.size();
+            }
+        };
+    }
+
+    private class Entry implements Map.Entry<EntityType<?>, EntityRenderer<?>> {
+        private final EntityType<?> key;
+
+        private Entry(EntityType<?> key) {
+            this.key = key;
+        }
+
+        @Override
+        public EntityType<?> getKey() {
+            return key;
+        }
+
+        @Override
+        public EntityRenderer<?> getValue() {
+            return get(key);
+        }
+
+        @Override
+        public EntityRenderer<?> setValue(EntityRenderer<?> value) {
+            return put(key, value);
+        }
     }
 
     @NotNull
     @Override
     public Set<Map.Entry<EntityType<?>, EntityRenderer<?>>> entrySet() {
-        return rendererMap.asMap().entrySet();
+        return new AbstractSet<>() {
+            @Override
+            public Iterator<Map.Entry<EntityType<?>, EntityRenderer<?>>> iterator() {
+                return Iterators.transform(Iterators.unmodifiableIterator(rendererProviders.keySet().iterator()), Entry::new);
+            }
+
+            @Override
+            public boolean contains(Object o) {
+                if (o instanceof Map.Entry<?,?> e) {
+                    return rendererProviders.containsKey(e.getKey()) && Objects.equals(get(e.getKey()), e.getValue());
+                } else {
+                    return false;
+                }
+            }
+
+            @Override
+            public int size() {
+                return rendererProviders.size();
+            }
+        };
     }
 }