Implemented cit application caching

Closes #31

12 files changed, 174 insertions, 26 deletions

diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java b/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java
index 982e1bc..790fb8c 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/ActiveCITs.java
@@ -10,9 +10,11 @@ import net.minecraft.util.Hand;
 import net.minecraft.util.Identifier;
 import net.minecraft.util.registry.Registry;
 import net.minecraft.world.World;
+import shcm.shsupercm.fabric.citresewn.config.CITResewnConfig;
 import shcm.shsupercm.fabric.citresewn.pack.cits.*;
 
 import java.util.*;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 public class ActiveCITs {
@@ -44,9 +46,8 @@ public class ActiveCITs {
     }
 
     public void dispose() {
-        for (CIT cit : cits) {
+        for (CIT cit : cits)
             cit.dispose();
-        }
         cits.clear();
         citItems.clear();
         citArmor.clear();
@@ -54,42 +55,71 @@ public class ActiveCITs {
         citEnchantments.clear();
     }
 
-    public BakedModel getItemModel(ItemStack stack, World world, LivingEntity entity, int seed) {
-        BakedModel bakedModel = null;
+    public CITItem getCITItem(ItemStack stack, World world, LivingEntity entity) {
         Hand hand = entity != null && stack == entity.getOffHandStack() ? Hand.OFF_HAND : Hand.MAIN_HAND;
 
         List<CITItem> citItems = this.citItems.get(stack.getItem());
         if (citItems != null)
             for (CITItem citItem : citItems)
-                if (citItem.test(stack, hand, world, entity)) {
-                    bakedModel = citItem.getItemModel(stack, hand, (ClientWorld) world, entity, seed);
-                    if (bakedModel != null)
-                        break;
-                }
-
-        if (bakedModel != null && bakedModel.getOverrides() != null)
-            bakedModel = bakedModel.getOverrides().apply(bakedModel, stack, (ClientWorld) world, entity, seed);
-
-        return bakedModel;
+                if (citItem.test(stack, hand, world, entity))
+                    return citItem;
+        return null;
     }
 
-    public Identifier getElytraTexture(ItemStack stack, World world, LivingEntity livingEntity) {
+    public CITElytra getCITElytra(ItemStack stack, World world, LivingEntity livingEntity) {
         for (CITElytra citElytra : citElytra)
             if (citElytra.test(stack, Hand.MAIN_HAND, world, livingEntity))
-                return citElytra.textureIdentifier;
-
+                return citElytra;
         return null;
     }
 
-    public Map<String, Identifier> getArmorTextures(ItemStack itemStack, World world, LivingEntity livingEntity) {
-        Item item = itemStack.getItem();
+    public CITArmor getCITArmor(ItemStack stack, World world, LivingEntity livingEntity) {
+        Item item = stack.getItem();
         if (item instanceof ArmorItem) {
             List<CITArmor> citArmor = this.citArmor.get(item);
             if (citArmor != null)
                 for (CITArmor armor : citArmor)
-                    if (armor.test(itemStack, null, world, livingEntity))
-                        return armor.textures;
+                    if (armor.test(stack, null, world, livingEntity))
+                        return armor;
         }
         return null;
     }
+
+    public BakedModel getItemModelCached(ItemStack stack, World world, LivingEntity entity, int seed) {
+        BakedModel bakedModel = null;
+
+        Supplier<CITItem> realtime = () -> getCITItem(stack, world, entity);
+
+        //noinspection ConstantConditions
+        CITItem citItem = CITResewnConfig.INSTANCE().cache_ms == 0 ? realtime.get() : ((CITItem.Cached) (Object) stack).citresewn_getCachedCITItem(realtime);
+
+        if (citItem != null)
+            bakedModel = citItem.getItemModel(stack, (ClientWorld) world, entity, seed);
+
+        return bakedModel;
+    }
+
+    public Identifier getElytraTextureCached(ItemStack stack, World world, LivingEntity livingEntity) {
+        Supplier<CITElytra> realtime = () -> getCITElytra(stack, world, livingEntity);
+
+        //noinspection ConstantConditions
+        CITElytra citElytra = CITResewnConfig.INSTANCE().cache_ms == 0 ? realtime.get() : ((CITElytra.Cached) (Object) stack).citresewn_getCachedCITElytra(realtime);
+
+        if (citElytra != null)
+            return citElytra.textureIdentifier;
+
+        return null;
+    }
+
+    public Map<String, Identifier> getArmorTexturesCached(ItemStack stack, World world, LivingEntity livingEntity) {
+        Supplier<CITArmor> realtime = () -> getCITArmor(stack, world, livingEntity);
+
+        //noinspection ConstantConditions
+        CITArmor citArmor = CITResewnConfig.INSTANCE().cache_ms == 0 ? realtime.get() : ((CITArmor.Cached) (Object) stack).citresewn_getCachedCITArmor(realtime);
+
+        if (citArmor != null)
+            return citArmor.textures;
+
+        return null;
+    }
 }
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfig.java b/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfig.java
index c1c5148..6cbd62e 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfig.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfig.java
@@ -13,6 +13,7 @@ public class CITResewnConfig {
     public boolean enabled = true;
     public boolean mute_errors = false;
     public boolean mute_warns = false;
+    public int cache_ms = 50;
     public boolean broken_paths = false;
 
     private static final File FILE = new File("config/citresewn.json");
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfigScreenFactory.java b/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfigScreenFactory.java
index 6fd51f7..34ab4f3 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfigScreenFactory.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/config/CITResewnConfigScreenFactory.java
@@ -7,6 +7,7 @@ import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.gui.screen.Screen;
 import net.minecraft.text.LiteralText;
 import net.minecraft.text.TranslatableText;
+import net.minecraft.util.Formatting;
 
 public class CITResewnConfigScreenFactory {
     public static Screen create(Screen parent) {
@@ -43,6 +44,25 @@ public class CITResewnConfigScreenFactory {
                 .setDefaultValue(defaultConfig.mute_warns)
                 .build());
 
+        category.addEntry(entryBuilder.startIntSlider(new TranslatableText("config.citresewn.cache_ms.title"), currentConfig.cache_ms / 50, 0, 5 * 20)
+                .setTooltip(new TranslatableText("config.citresewn.cache_ms.tooltip"))
+                .setSaveConsumer(newConfig -> currentConfig.cache_ms = newConfig * 50)
+                .setDefaultValue(currentConfig.cache_ms / 50)
+                .setTextGetter(ticks -> {
+                    if (ticks <= 1)
+                        return new TranslatableText("config.citresewn.cache_ms.ticks." + ticks).formatted(Formatting.AQUA);
+
+                    Formatting color = Formatting.DARK_RED;
+
+                    if (ticks <= 40) color = Formatting.RED;
+                    if (ticks <= 20) color = Formatting.GOLD;
+                    if (ticks <= 10) color = Formatting.DARK_GREEN;
+                    if (ticks <= 5) color = Formatting.GREEN;
+
+                    return new TranslatableText("config.citresewn.cache_ms.ticks.any", ticks).formatted(color);
+                })
+                .build());
+
         category.addEntry(entryBuilder.startBooleanToggle(new TranslatableText("config.citresewn.broken_paths.title"), currentConfig.broken_paths)
                 .setTooltip(new TranslatableText("config.citresewn.broken_paths.tooltip"))
                 .setSaveConsumer(newConfig -> currentConfig.broken_paths = newConfig)
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java
index f4f4fca..5028d62 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ArmorFeatureRendererMixin.java
@@ -35,7 +35,7 @@ public class ArmorFeatureRendererMixin<T extends LivingEntity, M extends BipedEn
         //compat Cosmetic Armor
         itemStack = OptionalCompat.getCosmeticArmor(itemStack, entity, armorSlot, false);
 
-        Map<String, Identifier> armorTextures = CITResewn.INSTANCE.activeCITs.getArmorTextures(itemStack, entity.world, entity);
+        Map<String, Identifier> armorTextures = CITResewn.INSTANCE.activeCITs.getArmorTexturesCached(itemStack, entity.world, entity);
         if (armorTextures != null) {
             armorTexturesCached = new WeakReference<>(armorTextures);
             return;
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ItemStackMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ItemStackMixin.java
new file mode 100644
index 0000000..6041c9b
--- /dev/null
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citarmor/ItemStackMixin.java
@@ -0,0 +1,25 @@
+package shcm.shsupercm.fabric.citresewn.mixin.citarmor;
+
+import net.minecraft.item.ItemStack;
+import org.spongepowered.asm.mixin.Mixin;
+import shcm.shsupercm.fabric.citresewn.config.CITResewnConfig;
+import shcm.shsupercm.fabric.citresewn.pack.cits.CITArmor;
+
+import java.lang.ref.WeakReference;
+import java.util.function.Supplier;
+
+@Mixin(ItemStack.class)
+public class ItemStackMixin implements CITArmor.Cached {
+    private WeakReference<CITArmor> citresewn_cachedCITArmor = new WeakReference<>(null);
+    private long citresewn_cacheTimeCITArmor = 0;
+
+    @Override
+    public CITArmor citresewn_getCachedCITArmor(Supplier<CITArmor> realtime) {
+        if (System.currentTimeMillis() - citresewn_cacheTimeCITArmor >= CITResewnConfig.INSTANCE().cache_ms) {
+            citresewn_cachedCITArmor = new WeakReference<>(realtime.get());
+            citresewn_cacheTimeCITArmor = System.currentTimeMillis();
+        }
+
+        return citresewn_cachedCITArmor.get();
+    }
+}
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java
index 2214598..3728362 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ElytraFeatureRendererMixin.java
@@ -47,7 +47,7 @@ public class ElytraFeatureRendererMixin {
         ItemStack itemStack = this.elytraItemCached.get();
         LivingEntity livingEntity = this.livingEntityCached.get();
         if (itemStack != null && itemStack.isOf(Items.ELYTRA) && livingEntity != null) {
-            Identifier elytraTexture = CITResewn.INSTANCE.activeCITs.getElytraTexture(itemStack, livingEntity.world, livingEntity);
+            Identifier elytraTexture = CITResewn.INSTANCE.activeCITs.getElytraTextureCached(itemStack, livingEntity.world, livingEntity);
             this.elytraItemCached = new WeakReference<>(null);
             this.livingEntityCached = new WeakReference<>(null);
             if (elytraTexture != null)
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ItemStackMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ItemStackMixin.java
new file mode 100644
index 0000000..aaf4037
--- /dev/null
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/citelytra/ItemStackMixin.java
@@ -0,0 +1,25 @@
+package shcm.shsupercm.fabric.citresewn.mixin.citelytra;
+
+import net.minecraft.item.ItemStack;
+import org.spongepowered.asm.mixin.Mixin;
+import shcm.shsupercm.fabric.citresewn.config.CITResewnConfig;
+import shcm.shsupercm.fabric.citresewn.pack.cits.CITElytra;
+
+import java.lang.ref.WeakReference;
+import java.util.function.Supplier;
+
+@Mixin(ItemStack.class)
+public class ItemStackMixin implements CITElytra.Cached {
+    private WeakReference<CITElytra> citresewn_cachedCITElytra = new WeakReference<>(null);
+    private long citresewn_cacheTimeCITElytra = 0;
+
+    @Override
+    public CITElytra citresewn_getCachedCITElytra(Supplier<CITElytra> realtime) {
+        if (System.currentTimeMillis() - citresewn_cacheTimeCITElytra >= CITResewnConfig.INSTANCE().cache_ms) {
+            citresewn_cachedCITElytra = new WeakReference<>(realtime.get());
+            citresewn_cacheTimeCITElytra = System.currentTimeMillis();
+        }
+
+        return citresewn_cachedCITElytra.get();
+    }
+}
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java
index 07a4f0d..c3aba64 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemRendererMixin.java
@@ -20,7 +20,7 @@ public class ItemRendererMixin {
         if (!CITResewnConfig.INSTANCE().enabled || CITResewn.INSTANCE.activeCITs == null)
             return;
 
-        BakedModel citModel = CITResewn.INSTANCE.activeCITs.getItemModel(stack, world == null ? MinecraftClient.getInstance().world : world, entity, seed);
+        BakedModel citModel = CITResewn.INSTANCE.activeCITs.getItemModelCached(stack, world == null ? MinecraftClient.getInstance().world : world, entity, seed);
         if (citModel != null)
             cir.setReturnValue(citModel);
     }
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemStackMixin.java b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemStackMixin.java
new file mode 100644
index 0000000..f739fc8
--- /dev/null
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/mixin/cititem/ItemStackMixin.java
@@ -0,0 +1,25 @@
+package shcm.shsupercm.fabric.citresewn.mixin.cititem;
+
+import net.minecraft.item.ItemStack;
+import org.spongepowered.asm.mixin.Mixin;
+import shcm.shsupercm.fabric.citresewn.config.CITResewnConfig;
+import shcm.shsupercm.fabric.citresewn.pack.cits.CITItem;
+
+import java.lang.ref.WeakReference;
+import java.util.function.Supplier;
+
+@Mixin(ItemStack.class)
+public class ItemStackMixin implements CITItem.Cached {
+    private WeakReference<CITItem> citresewn_cachedCITItem = new WeakReference<>(null);
+    private long citresewn_cacheTimeCITItem = 0;
+
+    @Override
+    public CITItem citresewn_getCachedCITItem(Supplier<CITItem> realtime) {
+        if (System.currentTimeMillis() - citresewn_cacheTimeCITItem >= CITResewnConfig.INSTANCE().cache_ms) {
+            citresewn_cachedCITItem = new WeakReference<>(realtime.get());
+            citresewn_cacheTimeCITItem = System.currentTimeMillis();
+        }
+
+        return citresewn_cachedCITItem.get();
+    }
+}
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITArmor.java b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITArmor.java
index 49a7d7b..e919a86 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITArmor.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITArmor.java
@@ -9,6 +9,7 @@ import shcm.shsupercm.fabric.citresewn.ex.CITParseException;
 import shcm.shsupercm.fabric.citresewn.pack.CITPack;
 
 import java.util.*;
+import java.util.function.Supplier;
 
 public class CITArmor extends CIT {
     public final Map<String, Identifier> textures = new HashMap<>();
@@ -34,4 +35,8 @@ public class CITArmor extends CIT {
             throw new CITParseException(pack.resourcePack, identifier, (e.getClass() == Exception.class ? "" : e.getClass().getSimpleName() + ": ") + e.getMessage());
         }
     }
+
+    public interface Cached {
+        CITArmor citresewn_getCachedCITArmor(Supplier<CITArmor> realtime);
+    }
 }
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITElytra.java b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITElytra.java
index aa092b3..5c9aa1a 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITElytra.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITElytra.java
@@ -6,6 +6,7 @@ import shcm.shsupercm.fabric.citresewn.ex.CITParseException;
 import shcm.shsupercm.fabric.citresewn.pack.CITPack;
 
 import java.util.Properties;
+import java.util.function.Supplier;
 
 public class CITElytra extends CIT {
     public final Identifier textureIdentifier;
@@ -20,4 +21,8 @@ public class CITElytra extends CIT {
             throw new CITParseException(pack.resourcePack, identifier, (e.getClass() == Exception.class ? "" : e.getClass().getSimpleName() + ": ") + e.getMessage());
         }
     }
+
+    public interface Cached {
+        CITElytra citresewn_getCachedCITElytra(Supplier<CITElytra> realtime);
+    }
 }
diff --git a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java
index cfddf09..690c521 100644
--- a/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java
+++ b/src/main/java/shcm/shsupercm/fabric/citresewn/pack/cits/CITItem.java
@@ -21,6 +21,7 @@ import net.minecraft.resource.ResourceType;
 import net.minecraft.util.Hand;
 import net.minecraft.util.Identifier;
 import net.minecraft.util.registry.Registry;
+import net.minecraft.world.World;
 import org.apache.commons.io.IOUtils;
 import shcm.shsupercm.fabric.citresewn.CITResewn;
 import shcm.shsupercm.fabric.citresewn.ex.CITLoadException;
@@ -314,8 +315,15 @@ public class CITItem extends CIT {
         return null;
     }
 
-    public BakedModel getItemModel(ItemStack stack, Hand hand, ClientWorld world, LivingEntity entity, int seed) {
-        return bakedSubModels.apply(this.bakedModel, stack, world, entity, seed);
+    public BakedModel getItemModel(ItemStack stack, ClientWorld world, LivingEntity entity, int seed) {
+        // get sub items or bakedModel if no sub item matches @Nullable
+        BakedModel bakedModel = bakedSubModels.apply(this.bakedModel, stack, world, entity, seed);
+
+        // apply model overrides
+        if (bakedModel != null && bakedModel.getOverrides() != null)
+            bakedModel = bakedModel.getOverrides().apply(bakedModel, stack, world, entity, seed);
+
+        return bakedModel;
     }
 
     public static class CITOverrideList extends ModelOverrideList {
@@ -338,4 +346,8 @@ public class CITItem extends CIT {
                     , bakedModel);
         }
     }
+
+    public interface Cached {
+        CITItem citresewn_getCachedCITItem(Supplier<CITItem> realtime);
+    }
 }