Use ImmutableMap for faster lookups of key/value indices in FastMap

src/main/java/malte0811/ferritecore/FastMap.java

@@ -11,15 +11,21 @@ public class FastMap<Value> {
     private final List<Key<?>> keys;
     private final List<Property<?>> rawKeys;
     private final List<Value> values;
+    private final Map<Property<?>, Integer> toKeyIndex;
 
     public FastMap(Collection<Property<?>> properties, Map<Map<Property<?>, Comparable<?>>, Value> valuesMap) {
-        List<Key<?>> keys = new ArrayList<>(properties.size());
+        this.rawKeys = ImmutableList.copyOf(properties);
+        List<Key<?>> keys = new ArrayList<>(rawKeys.size());
         int factorUpTo = 1;
-        for (Property<?> prop : properties) {
+        ImmutableMap.Builder<Property<?>, Integer> toKeyIndex = ImmutableMap.builder();
+        for (Property<?> prop : rawKeys) {
+            toKeyIndex.put(prop, keys.size());
             keys.add(new Key<>(prop, factorUpTo));
             factorUpTo *= prop.getAllowedValues().size();
         }
-        this.keys = keys;
+        this.keys = ImmutableList.copyOf(keys);
+        this.toKeyIndex = toKeyIndex.build();
+
         List<Value> valuesList = new ArrayList<>(factorUpTo);
         for (int i = 0; i < factorUpTo; ++i) {
             valuesList.add(null);
@@ -27,8 +33,7 @@ public class FastMap<Value> {
         for (Map.Entry<Map<Property<?>, Comparable<?>>, Value> state : valuesMap.entrySet()) {
             valuesList.set(getIndexOf(state.getKey()), state.getValue());
         }
-        this.values = valuesList;
-        this.rawKeys = ImmutableList.copyOf(properties);
+        this.values = ImmutableList.copyOf(valuesList);
     }
 
     @Nullable
@@ -48,14 +53,12 @@ public class FastMap<Value> {
     @Nullable
     private <T extends Comparable<T>>
     Key<T> getKeyFor(Property<T> prop) {
-        // It might be possible to speed this up by sorting the keys by their hash code and using a binary search,
-        // however I do not think that it would actually be faster in practice.
-        for (Key<?> key : keys) {
-            if (key.getProperty().equals(prop)) {
-                return (Key<T>) key;
-            }
+        Integer index = toKeyIndex.get(prop);
+        if (index == null) {
+            return null;
+        } else {
+            return (Key<T>) keys.get(index);
         }
-        return null;
     }
 
     public int getIndexOf(Map<Property<?>, Comparable<?>> state) {
@@ -97,11 +100,17 @@ public class FastMap<Value> {
         private final Property<T> property;
         private final List<T> values;
         private final int mapFactor;
+        private final Map<Comparable<?>, Integer> toValueIndex;
 
         private Key(Property<T> property, int mapFactor) {
             this.property = property;
-            this.values = new ArrayList<>(property.getAllowedValues());
+            this.values = ImmutableList.copyOf(property.getAllowedValues());
             this.mapFactor = mapFactor;
+            ImmutableMap.Builder<Comparable<?>, Integer> toValueIndex = ImmutableMap.builder();
+            for (int i = 0; i < this.values.size(); i++) {
+                toValueIndex.put(this.values.get(i), i);
+            }
+            this.toValueIndex = toValueIndex.build();
         }
 
         public int toPartialMapIndex(Comparable<?> value) {
@@ -109,8 +118,8 @@ public class FastMap<Value> {
         }
 
         private int getInternalIndex(Comparable<?> value) {
-            int result = values.indexOf(value);
-            if (result >= 0) {
+            Integer result = toValueIndex.get(value);
+            if (result != null) {
                 return result;
             } else {
                 throw new IllegalStateException("Unknown value: "+value+" in "+property);

src/main/java/malte0811/ferritecore/mixin/fastmap/FastMapStateHolderMixin.java

@@ -27,7 +27,8 @@ public abstract class FastMapStateHolderMixin<O, S> implements FastMapStateHolde
             method = "with",
             at = @At(
                     value = "INVOKE",
-                    target = "Lcom/google/common/collect/Table;get(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;"
+                    target = "Lcom/google/common/collect/Table;get(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;",
+                    remap = false
             )
     )
     public Object getNeighborFromFastMap(Table<?, ?, ?> ignore, Object rowKey, Object columnKey) {

summary.md

@@ -39,14 +39,10 @@ with a similar coefficient. A `FastMap` in this case is simply an `ArrayList` us
 multi-dimensional array. Finding a neighbor state can be done by a few integer modulo,
 multiplication and addition operations.
 
-Saved memory: Around 600 MB (the `FastMap`s are around 4 MB total)  
+Saved memory: Around 600 MB (the `FastMap`s are around 7 MB total)  
 CPU impact: hard to prove, but most likely near zero  
 Side: both  
-Mixin folder: `fastmap`  
-Notes: There is a chance that this is slower for blocks with a very high number of
-properties, as the property to be modified is found using a linear search. This could be
-improved by using a binary search (probably sorted by hash), but it is not clear that this
-will ever be faster in practice.
+Mixin subpackage: `fastmap`  
 
 ### BlockState property storage
 Each blockstate stores its properties as an `ImmutableMap<Property<?>, Comparable<?>>`,
@@ -63,7 +59,7 @@ old one with the new implementation where possible.
 Saved memory: Around 170 MB  
 CPU impact: unclear (see second paragraph)  
 Side: both  
-Mixin folder: `nopropertymap`  
+Mixin subpackage: `nopropertymap`  
 
 ### Multipart model predicate caching
 Each multipart model stores a number of predicates to determine which parts to show under
@@ -85,4 +81,4 @@ Saved memory: 300-400 MB (relative to the state after the first change, so 100 M
 compared to a "clean" instance)  
 CPU impact: Some impact in model loading (but less allocations), zero while playing  
 Side: client  
-Mixin folder: `predicates`  
+Mixin subpackage: `predicates`