From e709b3dc38be4fefa412a71a24e8155d1e68fba9 Mon Sep 17 00:00:00 2001
From: Skye <me@skye.vg>
Date: Thu, 21 Nov 2024 18:51:43 +0900
Subject: [PATCH] Cache iota sizes to improve performance

---
 .../hexcasting/api/casting/iota/Iota.java     |  4 +++
 .../hexcasting/api/casting/iota/IotaType.java | 34 ++++---------------
 .../hexcasting/api/casting/iota/ListIota.java | 23 +++++++++++++
 3 files changed, 33 insertions(+), 28 deletions(-)

diff --git a/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/Iota.java b/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/Iota.java
index 7f024abc94..cad6baee82 100644
--- a/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/Iota.java
+++ b/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/Iota.java
@@ -96,6 +96,10 @@ public int size() {
         return 1;
     }
 
+    public int depth() {
+        return 1;
+    }
+
     public Component display() {
         return this.type.display(this.serialize());
     }
diff --git a/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/IotaType.java b/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/IotaType.java
index 4bcf9b4b1f..c0167ae5df 100644
--- a/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/IotaType.java
+++ b/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/IotaType.java
@@ -15,8 +15,6 @@
 import net.minecraft.util.FormattedCharSequence;
 
 import javax.annotation.Nullable;
-import java.util.ArrayDeque;
-import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 
@@ -81,33 +79,13 @@ public static boolean isTooLargeToSerialize(Iterable<Iota> examinee) {
     }
 
     private static boolean isTooLargeToSerialize(Iterable<Iota> examinee, int startingCount) {
-        // We don't recurse here, just a work queue (or work stack, if we liked.)
-        // Each element is a found sub-iota, and how deep it is.
-        //
-        // TODO: is it worth trying to cache the depth and size statically on a SpellList.
-        var listsToExamine = new ArrayDeque<>(Collections.singleton(new Pair<>(examinee, 0)));
-        int totalEltsFound = startingCount; // count the first list
-        while (!listsToExamine.isEmpty()) {
-            var iotaPair = listsToExamine.removeFirst();
-            var sublist = iotaPair.getFirst();
-            int depth = iotaPair.getSecond();
-            for (var iota : sublist) {
-                totalEltsFound += iota.size();
-                if (totalEltsFound >= HexIotaTypes.MAX_SERIALIZATION_TOTAL) {
-                    return true; // too bad
-                }
-                var subIotas = iota.subIotas();
-                if (subIotas != null) {
-                    if (depth + 1 >= HexIotaTypes.MAX_SERIALIZATION_DEPTH) {
-                        return true;
-                    }
-
-                    listsToExamine.addLast(new Pair<>(subIotas, depth + 1));
-                }
-            }
+        int totalSize = startingCount;
+        for (Iota iota : examinee) {
+            if (iota.depth() >= HexIotaTypes.MAX_SERIALIZATION_DEPTH)
+                return true;
+            totalSize += iota.size();
         }
-        // we made it!
-        return false;
+        return totalSize >= HexIotaTypes.MAX_SERIALIZATION_TOTAL;
     }
 
     /**
diff --git a/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/ListIota.java b/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/ListIota.java
index 251c4e161d..d3eb9455bf 100644
--- a/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/ListIota.java
+++ b/Common/src/main/java/at/petrak/hexcasting/api/casting/iota/ListIota.java
@@ -15,12 +15,25 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import static java.lang.Math.max;
+
 /**
  * This is a <i>wrapper</i> for {@link SpellList}.
  */
 public class ListIota extends Iota {
+    private final int depth;
+    private final int size;
+
     public ListIota(@NotNull SpellList list) {
         super(HexIotaTypes.LIST, list);
+        int maxChildDepth = 0;
+        int totalSize = 1;
+        for (Iota iota : list) {
+            totalSize += iota.size();
+            maxChildDepth = max(maxChildDepth, iota.depth());
+        }
+        depth = maxChildDepth + 1;
+        size = totalSize;
     }
 
     public ListIota(@NotNull List<Iota> list) {
@@ -78,6 +91,16 @@ public boolean toleratesOther(Iota that) {
         return this.getList();
     }
 
+    @Override
+    public int size() {
+        return size;
+    }
+
+    @Override
+    public int depth() {
+        return depth;
+    }
+
     public static IotaType<ListIota> TYPE = new IotaType<>() {
         @Nullable
         @Override