Precompute cumulative weights for weighted block-model selection

src/main/java/me/jellysquid/mods/sodium/mixin/features/model/WeightedBakedModelMixin.java

@@ -15,6 +15,9 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 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.callback.CallbackInfo;
 
 import java.util.*;
 
@@ -28,13 +31,33 @@ public class WeightedBakedModelMixin implements UnwrappableBakedModel {
     @Final
     private int totalWeight;
 
+    /**
+     * The running totals of the weights of the entries in {@link #list}. Neither the list nor the weights change after
+     * the model has been constructed.
+     */
+    @Unique
+    private int[] embeddium$cumulativeWeights;
+
+    @Inject(method = "<init>", at = @At("RETURN"))
+    private void computeCumulativeWeights(CallbackInfo ci) {
+        int[] cumulativeWeights = new int[this.list.size()];
+        int sum = 0;
+
+        for (int i = 0; i < cumulativeWeights.length; i++) {
+            sum += this.list.get(i).getWeight().asInt();
+            cumulativeWeights[i] = sum;
+        }
+
+        this.embeddium$cumulativeWeights = cumulativeWeights;
+    }
+
     /**
      * @author JellySquid
      * @reason Avoid excessive object allocations
      */
     @Overwrite(remap = false)
     public List<BakedQuad> getQuads(@Nullable BlockState state, @Nullable Direction face, RandomSource random, ModelData modelData, RenderType renderLayer) {
-        WeightedEntry.Wrapper<BakedModel> quad = getAt(this.list, Math.abs((int) random.nextLong()) % this.totalWeight);
+        WeightedEntry.Wrapper<BakedModel> quad = this.getAt(Math.abs((int) random.nextLong()) % this.totalWeight);
 
         if (quad != null) {
             return quad.getData()
@@ -50,7 +73,7 @@ public List<BakedQuad> getQuads(@Nullable BlockState state, @Nullable Direction
      */
     @Overwrite(remap = false)
     public ChunkRenderTypeSet getRenderTypes(@NotNull BlockState state, @NotNull RandomSource rand, @NotNull ModelData data) {
-        WeightedEntry.Wrapper<BakedModel> quad = getAt(this.list, Math.abs((int) rand.nextLong()) % this.totalWeight);
+        WeightedEntry.Wrapper<BakedModel> quad = this.getAt(Math.abs((int) rand.nextLong()) % this.totalWeight);
 
         if (quad != null) {
             return quad.getData().getRenderTypes(state, rand, data);
@@ -59,28 +82,27 @@ public ChunkRenderTypeSet getRenderTypes(@NotNull BlockState state, @NotNull Ran
         return ChunkRenderTypeSet.none();
     }
 
+    /**
+     * Finds the entry which the given weighted index selects. Vanilla walks the list, subtracting each entry's weight
+     * from the index, and stops at the first entry for which the result is negative. Weights are never negative, so
+     * that is the first entry whose running total is greater than the index, and the entries need not be read.
+     */
     @Unique
-    private static <T extends WeightedEntry> T getAt(List<T> pool, int totalWeight) {
-        int i = 0;
-        int len = pool.size();
-
-        T weighted;
+    private WeightedEntry.Wrapper<BakedModel> getAt(int weightedIndex) {
+        int[] cumulativeWeights = this.embeddium$cumulativeWeights;
 
-        do {
-            if (i >= len) {
-                return null;
+        for (int i = 0; i < cumulativeWeights.length; i++) {
+            if (cumulativeWeights[i] > weightedIndex) {
+                return this.list.get(i);
             }
+        }
 
-            weighted = pool.get(i++);
-            totalWeight -= weighted.getWeight().asInt();
-        } while (totalWeight >= 0);
-
-        return weighted;
+        return null;
     }
 
     @Override
     public @Nullable BakedModel embeddium$getInnerModel(RandomSource rand) {
-        WeightedEntry.Wrapper<BakedModel> quad = getAt(this.list, Math.abs((int) rand.nextLong()) % this.totalWeight);
+        WeightedEntry.Wrapper<BakedModel> quad = this.getAt(Math.abs((int) rand.nextLong()) % this.totalWeight);
 
         if (quad != null && quad.getData().getClass() == SimpleBakedModel.class) {
             return quad.getData();