Slight cleanup of FastMap code

src/main/java/malte0811/ferritecore/FastMapEntry.java

@@ -1,49 +0,0 @@
-package malte0811.ferritecore;
-
-import net.minecraft.state.Property;
-import org.apache.commons.lang3.tuple.Pair;
-
-import javax.annotation.Nonnull;
-import java.util.*;
-
-public class FastMapEntry extends AbstractMap<Property<?>, Comparable<?>> {
-    private final FastMap<?> baseMap;
-    private final int indexInBaseMap;
-
-    public FastMapEntry(FastMap<?> baseMap, int indexInBaseMap) {
-        this.baseMap = baseMap;
-        this.indexInBaseMap = indexInBaseMap;
-    }
-
-    @Nonnull
-    @Override
-    public Set<Entry<Property<?>, Comparable<?>>> entrySet() {
-        return new EntrySet();
-    }
-
-    private class EntrySet extends AbstractSet<Map.Entry<Property<?>, Comparable<?>>> {
-
-        @Nonnull
-        @Override
-        public Iterator<Entry<Property<?>, Comparable<?>>> iterator() {
-            Iterator<Property<?>> baseIterator = baseMap.getProperties().iterator();
-            return new Iterator<Entry<Property<?>, Comparable<?>>>() {
-                @Override
-                public boolean hasNext() {
-                    return baseIterator.hasNext();
-                }
-
-                @Override
-                public Entry<Property<?>, Comparable<?>> next() {
-                    Property<?> nextProp = baseIterator.next();
-                    return Pair.of(nextProp, baseMap.getValue(indexInBaseMap, nextProp));
-                }
-            };
-        }
-
-        @Override
-        public int size() {
-            return baseMap.getProperties().size();
-        }
-    }
-}

src/main/java/malte0811/ferritecore/ducks/FastMapStateHolder.java

@@ -1,7 +1,7 @@
 package malte0811.ferritecore.ducks;
 
 import com.google.common.collect.ImmutableMap;
-import malte0811.ferritecore.FastMap;
+import malte0811.ferritecore.fastmap.FastMap;
 import net.minecraft.state.Property;
 
 public interface FastMapStateHolder<S> {

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

@@ -1,4 +1,4 @@
-package malte0811.ferritecore;
+package malte0811.ferritecore.fastmap;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -8,19 +8,22 @@ import javax.annotation.Nullable;
 import java.util.*;
 
 public class FastMap<Value> {
-    private final List<Key<?>> keys;
+    private final List<FastMapKey<?>> keys;
     private final List<Property<?>> rawKeys;
     private final List<Value> values;
+    // It might be possible to get rid of this (and the equivalent map for values) by sorting the key vectors by
+    // property name (natural order for values) and using a binary search above a given size, but choosing that size
+    // would likely be more effort than it's worth
     private final Map<Property<?>, Integer> toKeyIndex;
 
     public FastMap(Collection<Property<?>> properties, Map<Map<Property<?>, Comparable<?>>, Value> valuesMap) {
         this.rawKeys = ImmutableList.copyOf(properties);
-        List<Key<?>> keys = new ArrayList<>(rawKeys.size());
+        List<FastMapKey<?>> keys = new ArrayList<>(rawKeys.size());
         int factorUpTo = 1;
         ImmutableMap.Builder<Property<?>, Integer> toKeyIndex = ImmutableMap.builder();
         for (Property<?> prop : rawKeys) {
             toKeyIndex.put(prop, keys.size());
-            keys.add(new Key<>(prop, factorUpTo));
+            keys.add(new FastMapKey<>(prop, factorUpTo));
             factorUpTo *= prop.getAllowedValues().size();
         }
         this.keys = ImmutableList.copyOf(keys);
@@ -39,7 +42,7 @@ public class FastMap<Value> {
     @Nullable
     public <T extends Comparable<T>>
     Value with(int last, Property<T> prop, T value) {
-        final Key<T> keyToChange = getKeyFor(prop);
+        final FastMapKey<T> keyToChange = getKeyFor(prop);
         if (keyToChange == null) {
             return null;
         }
@@ -50,20 +53,9 @@ public class FastMap<Value> {
         return values.get(newIndex);
     }
 
-    @Nullable
-    private <T extends Comparable<T>>
-    Key<T> getKeyFor(Property<T> prop) {
-        Integer index = toKeyIndex.get(prop);
-        if (index == null) {
-            return null;
-        } else {
-            return (Key<T>) keys.get(index);
-        }
-    }
-
     public int getIndexOf(Map<Property<?>, Comparable<?>> state) {
         int id = 0;
-        for (Key<?> k : keys) {
+        for (FastMapKey<?> k : keys) {
             id += k.toPartialMapIndex(state.get(k.getProperty()));
         }
         return id;
@@ -72,14 +64,14 @@ public class FastMap<Value> {
     @Nullable
     public <T extends Comparable<T>>
     T getValue(int stateIndex, Property<T> property) {
-        final Key<T> propId = getKeyFor(property);
+        final FastMapKey<T> propId = getKeyFor(property);
         if (propId == null) {
             return null;
         }
         return propId.getValue(stateIndex);
     }
 
-    public Collection<Property<?>> getProperties() {
+    public List<Property<?>> getProperties() {
         return rawKeys;
     }
 
@@ -96,55 +88,22 @@ public class FastMap<Value> {
         return with(globalTableIndex, rowKey, (T) columnKey);
     }
 
-    private static class Key<T extends Comparable<T>> {
-        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 = 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) {
-            return mapFactor * getInternalIndex(value);
-        }
-
-        private int getInternalIndex(Comparable<?> value) {
-            Integer result = toValueIndex.get(value);
-            if (result != null) {
-                return result;
-            } else {
-                throw new IllegalStateException("Unknown value: "+value+" in "+property);
-            }
-        }
-
-        public T getValue(int mapIndex) {
-            int index = (mapIndex / mapFactor) % values.size();
-            return values.get(index);
-        }
+    public int numProperties() {
+        return keys.size();
+    }
 
-        public int replaceIn(int mapIndex, T newValue) {
-            final int lowerData = mapIndex % mapFactor;
-            final int upperFactor = mapFactor * values.size();
-            final int upperData = mapIndex - mapIndex % upperFactor;
-            int internalIndex = getInternalIndex(newValue);
-            if (internalIndex < 0) {
-                return -1;
-            } else {
-                return lowerData + mapFactor * internalIndex + upperData;
-            }
-        }
+    FastMapKey<?> getKey(int keyIndex) {
+        return keys.get(keyIndex);
+    }
 
-        public Property<T> getProperty() {
-            return property;
+    @Nullable
+    private <T extends Comparable<T>>
+    FastMapKey<T> getKeyFor(Property<T> prop) {
+        Integer index = toKeyIndex.get(prop);
+        if (index == null) {
+            return null;
+        } else {
+            return (FastMapKey<T>) getKey(index);
         }
     }
 }

src/main/java/malte0811/ferritecore/fastmap/FastMapKey.java

@@ -0,0 +1,60 @@
+package malte0811.ferritecore.fastmap;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import net.minecraft.state.Property;
+
+import java.util.List;
+import java.util.Map;
+
+class FastMapKey<T extends Comparable<T>> {
+    private final Property<T> property;
+    private final List<T> values;
+    private final int mapFactor;
+    private final Map<Comparable<?>, Integer> toValueIndex;
+
+    FastMapKey(Property<T> property, int mapFactor) {
+        this.property = property;
+        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();
+    }
+
+    T getValue(int mapIndex) {
+        int index = (mapIndex / mapFactor) % values.size();
+        return values.get(index);
+    }
+
+    int replaceIn(int mapIndex, T newValue) {
+        final int lowerData = mapIndex % mapFactor;
+        final int upperFactor = mapFactor * values.size();
+        final int upperData = mapIndex - mapIndex % upperFactor;
+        int internalIndex = getInternalIndex(newValue);
+        if (internalIndex < 0) {
+            return -1;
+        } else {
+            return lowerData + mapFactor * internalIndex + upperData;
+        }
+    }
+
+    Property<T> getProperty() {
+        return property;
+    }
+
+    int toPartialMapIndex(Comparable<?> value) {
+        return mapFactor * getInternalIndex(value);
+    }
+
+    private int getInternalIndex(Comparable<?> value) {
+        Integer result = toValueIndex.get(value);
+        if (result != null) {
+            return result;
+        } else {
+            throw new IllegalStateException("Unknown value: " + value + " in " + property);
+        }
+    }
+}

src/main/java/malte0811/ferritecore/fastmap/FastSubMap.java

@@ -0,0 +1,37 @@
+package malte0811.ferritecore.fastmap;
+
+import net.minecraft.state.Property;
+
+import javax.annotation.Nonnull;
+import java.util.AbstractMap;
+import java.util.Set;
+
+public class FastSubMap extends AbstractMap<Property<?>, Comparable<?>> {
+    final FastMap<?> baseMap;
+    final int indexInBaseMap;
+
+    public FastSubMap(FastMap<?> baseMap, int indexInBaseMap) {
+        this.baseMap = baseMap;
+        this.indexInBaseMap = indexInBaseMap;
+    }
+
+    @Nonnull
+    @Override
+    public Set<Entry<Property<?>, Comparable<?>>> entrySet() {
+        return new FastSubMapEntrySet(this);
+    }
+
+    @Override
+    public Comparable<?> get(Object key) {
+        if (!(key instanceof Property<?>)) {
+            return null;
+        }
+        Property<?> prop = (Property<?>) key;
+        return baseMap.getValue(indexInBaseMap, prop);
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+        return get(key) != null;
+    }
+}

src/main/java/malte0811/ferritecore/fastmap/FastSubMapEntrySet.java

@@ -0,0 +1,45 @@
+package malte0811.ferritecore.fastmap;
+
+import net.minecraft.state.Property;
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nonnull;
+import java.util.AbstractSet;
+import java.util.Iterator;
+import java.util.Map;
+
+class FastSubMapEntrySet extends AbstractSet<Map.Entry<Property<?>, Comparable<?>>> {
+    private final FastMap<?> fastMap;
+    private final int mapIndex;
+
+    public FastSubMapEntrySet(FastSubMap fastSubMap) {
+        this.fastMap = fastSubMap.baseMap;
+        this.mapIndex = fastSubMap.indexInBaseMap;
+    }
+
+    @Nonnull
+    @Override
+    public Iterator<Map.Entry<Property<?>, Comparable<?>>> iterator() {
+        return new Iterator<Map.Entry<Property<?>, Comparable<?>>>() {
+            private int iteratorIndex = 0;
+
+            @Override
+            public boolean hasNext() {
+                return iteratorIndex < fastMap.numProperties();
+            }
+
+            @Override
+            public Map.Entry<Property<?>, Comparable<?>> next() {
+                Property<?> nextProp = fastMap.getProperties().get(iteratorIndex);
+                Comparable<?> nextValue = fastMap.getKey(iteratorIndex).getValue(mapIndex);
+                ++iteratorIndex;
+                return Pair.of(nextProp, nextValue);
+            }
+        };
+    }
+
+    @Override
+    public int size() {
+        return fastMap.numProperties();
+    }
+}

src/main/java/malte0811/ferritecore/impl/StateHolderImpl.java

@@ -1,8 +1,8 @@
 package malte0811.ferritecore.impl;
 
-import malte0811.ferritecore.FastMap;
 import malte0811.ferritecore.ducks.FastMapStateHolder;
 import malte0811.ferritecore.ducks.NoPropertyStateHolder;
+import malte0811.ferritecore.fastmap.FastMap;
 import net.minecraft.state.Property;
 
 import java.util.Map;

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

@@ -2,8 +2,8 @@ package malte0811.ferritecore.mixin.fastmap;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Table;
-import malte0811.ferritecore.FastMap;
 import malte0811.ferritecore.ducks.FastMapStateHolder;
+import malte0811.ferritecore.fastmap.FastMap;
 import malte0811.ferritecore.impl.StateHolderImpl;
 import net.minecraft.state.Property;
 import net.minecraft.state.StateHolder;

src/main/java/malte0811/ferritecore/mixin/nopropertymap/NoPropertyStateHolderMixin.java

@@ -1,10 +1,10 @@
 package malte0811.ferritecore.mixin.nopropertymap;
 
 import com.google.common.collect.ImmutableMap;
-import malte0811.ferritecore.FastMap;
-import malte0811.ferritecore.FastMapEntry;
 import malte0811.ferritecore.ducks.FastMapStateHolder;
 import malte0811.ferritecore.ducks.NoPropertyStateHolder;
+import malte0811.ferritecore.fastmap.FastMap;
+import malte0811.ferritecore.fastmap.FastSubMap;
 import net.minecraft.state.Property;
 import net.minecraft.state.StateHolder;
 import org.spongepowered.asm.mixin.Mixin;
@@ -81,7 +81,7 @@ public abstract class NoPropertyStateHolderMixin implements NoPropertyStateHolde
         final FastMap<?> globalTable = ((FastMapStateHolder<?>) this).getStateMap();
         if (globalTable != null) {
             final int globalIndex = ((FastMapStateHolder<?>) this).getStateIndex();
-            return new FastMapEntry(globalTable, globalIndex);
+            return new FastSubMap(globalTable, globalIndex);
         } else {
             return getValues();
         }