Hash an ingredient key once per index update

loader-neoforge/src/main/java/org/cyclops/cyclopscore/ingredient/collection/IIngredientMapMutable.java

@@ -5,6 +5,7 @@
 import javax.annotation.Nullable;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.BiFunction;
 
 /**
  * A mutable mapping from ingredient component instances to values of any type.
@@ -90,6 +91,33 @@ default public int removeAll(Iterable<? extends T> instances, M matchCondition)
         return removed;
     }
 
+    /**
+     * Compute a new value for the given key instance, following {@link Map#compute} semantics.
+     *
+     * The remapping function is passed the key and the currently mapped value, or null when the key
+     * is absent. Returning null removes the mapping, returning anything else stores it.
+     *
+     * Implementations back onto a hash of the key instance, which for some component types is
+     * expensive. Doing this in one call rather than a get followed by a put lets them hash once.
+     *
+     * @param key An instance key.
+     * @param remappingFunction A function computing the new value from the key and the old value.
+     * @return The new value associated with the key, or null if none.
+     */
+    @Nullable
+    default V compute(T key, BiFunction<T, V, V> remappingFunction) {
+        V oldValue = get(key);
+        V newValue = remappingFunction.apply(key, oldValue);
+        if (newValue == null) {
+            if (oldValue != null) {
+                remove(key);
+            }
+        } else {
+            put(key, newValue);
+        }
+        return newValue;
+    }
+
     /**
      * Add all entries from the given map to this map.
      * @param map A map, that will not be changed, only read.

loader-neoforge/src/main/java/org/cyclops/cyclopscore/ingredient/collection/IngredientCollectionPrototypeMap.java

@@ -52,32 +52,35 @@ public IngredientCollectionPrototypeMap(IngredientComponent<T, M> component, boo
     public boolean add(T instance) {
         IIngredientMatcher<T, M> matcher = getComponent().getMatcher();
         T prototype = getPrototype(instance);
-        Long value = ingredients.get(prototype);
-        long existingValue = value == null ? 0 : value;
-        long newValue = Math.addExact(existingValue, matcher.getQuantity(instance));
-        if (newValue != 0) {
-            ingredients.put(prototype, newValue);
-        } else {
-            ingredients.remove(prototype);
-        }
+        long quantity = matcher.getQuantity(instance);
+        // Computed in one call so the prototype is hashed once rather than once to read and once
+        // to write. Hashing a prototype is the dominant cost here for component-bearing instances.
+        ingredients.compute(prototype, (key, value) -> {
+            long existingValue = value == null ? 0 : value;
+            long newValue = Math.addExact(existingValue, quantity);
+            return newValue == 0 ? null : newValue;
+        });
         return true;
     }
 
     @Override
     public boolean remove(T instance) {
         IIngredientMatcher<T, M> matcher = getComponent().getMatcher();
         T prototype = getPrototype(instance);
-        Long value = ingredients.get(prototype);
-        long existingValue = value == null ? 0 : value;
         long currentValue = matcher.getQuantity(instance);
-        if (currentValue == existingValue) {
-            ingredients.remove(prototype);
-            return true;
-        } else if (currentValue < existingValue || isNegativeQuantities()) {
-            ingredients.put(prototype, existingValue - currentValue);
-            return true;
-        }
-        return false;
+        boolean[] removed = new boolean[1];
+        ingredients.compute(prototype, (key, value) -> {
+            long existingValue = value == null ? 0 : value;
+            if (currentValue == existingValue) {
+                removed[0] = true;
+                return null;
+            } else if (currentValue < existingValue || isNegativeQuantities()) {
+                removed[0] = true;
+                return existingValue - currentValue;
+            }
+            return value;
+        });
+        return removed[0];
     }
 
     @Override

loader-neoforge/src/main/java/org/cyclops/cyclopscore/ingredient/collection/IngredientMapSingleClassified.java

@@ -12,6 +12,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.BiFunction;
 import java.util.function.Supplier;
 
 /**
@@ -112,6 +113,29 @@ public V remove(T key) {
         return null;
     }
 
+    @Nullable
+    @Override
+    public V compute(T key, BiFunction<T, V, V> remappingFunction) {
+        C classifier = getClassifier(key);
+        IIngredientMapMutable<T, M, V> map = this.classifiedMaps.get(classifier);
+        if (map == null) {
+            // Nothing is stored under this classifier, so only an insertion can come out of this
+            V newValue = remappingFunction.apply(key, null);
+            if (newValue != null) {
+                getOrCreateClassifiedCollection(classifier).put(key, newValue);
+                this.size++;
+            }
+            return newValue;
+        }
+        int sizeBefore = map.size();
+        V newValue = map.compute(key, remappingFunction);
+        this.size += map.size() - sizeBefore;
+        if (map.isEmpty()) {
+            this.classifiedMaps.remove(classifier);
+        }
+        return newValue;
+    }
+
     @Override
     public int size() {
         return this.size;

loader-neoforge/src/main/java/org/cyclops/cyclopscore/ingredient/collection/IngredientMapWrappedAdapter.java

@@ -9,6 +9,7 @@
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.function.BiFunction;
 
 /**
  * An abstract ingredient map adapter.
@@ -68,6 +69,14 @@ public V get(T key) {
         return this.collection.get(wrap(key));
     }
 
+    @Nullable
+    @Override
+    public V compute(T key, BiFunction<T, V, V> remappingFunction) {
+        // One wrapper, so the key is hashed once instead of once for the get and once for the put
+        return this.collection.compute(wrap(key),
+                (wrapper, value) -> remappingFunction.apply(wrapper.getInstance(), value));
+    }
+
     @Override
     public IngredientSet<T, M> keySet() {
         return new IngredientSet<>(this.getComponent(), this.collection.keySet());

loader-neoforge/src/test/java/org/cyclops/cyclopscore/ingredient/collection/TestIngredientMapCompute.java

@@ -0,0 +1,141 @@
+package org.cyclops.cyclopscore.ingredient.collection;
+
+import org.cyclops.cyclopscore.ingredient.ComplexStack;
+import org.cyclops.cyclopscore.ingredient.IngredientComponentStubs;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.function.Supplier;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * {@link IIngredientMapMutable#compute} across the map implementations.
+ *
+ * It exists so that a read followed by a write can hash the key once instead of twice, so the
+ * results have to stay identical to doing both separately.
+ *
+ * @author rubensworks
+ */
+@RunWith(Parameterized.class)
+public class TestIngredientMapCompute {
+
+    private static final ComplexStack A01 = new ComplexStack(ComplexStack.Group.A, 0, 1, null);
+    private static final ComplexStack A12 = new ComplexStack(ComplexStack.Group.A, 1, 2, null);
+    private static final ComplexStack B01 = new ComplexStack(ComplexStack.Group.B, 0, 1, null);
+
+    @Parameterized.Parameters
+    public static Collection<Object[]> data() {
+        return Arrays.<Object[]>asList(new Object[][]{
+                {(Supplier<IIngredientMapMutable<ComplexStack, Integer, String>>)
+                        () -> new IngredientHashMap<>(IngredientComponentStubs.COMPLEX)},
+                {(Supplier<IIngredientMapMutable<ComplexStack, Integer, String>>)
+                        () -> new IngredientTreeMap<>(IngredientComponentStubs.COMPLEX)},
+                {(Supplier<IIngredientMapMutable<ComplexStack, Integer, String>>)
+                        () -> new IngredientMapSingleClassified<>(IngredientComponentStubs.COMPLEX,
+                                () -> new IngredientHashMap<>(IngredientComponentStubs.COMPLEX),
+                                IngredientComponentStubs.COMPLEX.getCategoryTypes().get(0))},
+        });
+    }
+
+    private final Supplier<IIngredientMapMutable<ComplexStack, Integer, String>> factory;
+
+    public TestIngredientMapCompute(Supplier<IIngredientMapMutable<ComplexStack, Integer, String>> factory) {
+        this.factory = factory;
+    }
+
+    @Test
+    public void testInsertsWhenAbsent() {
+        IIngredientMapMutable<ComplexStack, Integer, String> map = this.factory.get();
+        assertThat(map.compute(A01, (key, value) -> {
+            assertThat(value, is(nullValue()));
+            return "first";
+        }), is("first"));
+        assertThat(map.get(A01), is("first"));
+        assertThat(map.size(), is(1));
+    }
+
+    @Test
+    public void testUpdatesWhenPresent() {
+        IIngredientMapMutable<ComplexStack, Integer, String> map = this.factory.get();
+        map.put(A01, "first");
+        assertThat(map.compute(A01, (key, value) -> value + "+second"), is("first+second"));
+        assertThat(map.get(A01), is("first+second"));
+        assertThat(map.size(), is(1));
+    }
+
+    @Test
+    public void testRemovesOnNull() {
+        IIngredientMapMutable<ComplexStack, Integer, String> map = this.factory.get();
+        map.put(A01, "first");
+        map.put(A12, "other");
+        assertThat(map.compute(A01, (key, value) -> null), is(nullValue()));
+        assertThat(map.get(A01), is(nullValue()));
+        assertThat(map.get(A12), is("other"));
+        assertThat(map.size(), is(1));
+    }
+
+    @Test
+    public void testNullOnAbsentIsNoOp() {
+        IIngredientMapMutable<ComplexStack, Integer, String> map = this.factory.get();
+        assertThat(map.compute(A01, (key, value) -> null), is(nullValue()));
+        assertThat(map.size(), is(0));
+        assertThat(map.isEmpty(), is(true));
+    }
+
+    @Test
+    public void testKeyIsPassedThrough() {
+        IIngredientMapMutable<ComplexStack, Integer, String> map = this.factory.get();
+        map.compute(A01, (key, value) -> {
+            assertThat(key, is(A01));
+            return "v";
+        });
+    }
+
+    /**
+     * The classified map keeps its own size and drops classifiers that run empty, so computing
+     * away the last entry of a classifier has to clean up just as a remove would.
+     */
+    @Test
+    public void testSizeTracksAcrossClassifiers() {
+        IIngredientMapMutable<ComplexStack, Integer, String> map = this.factory.get();
+        map.compute(A01, (key, value) -> "a");
+        map.compute(B01, (key, value) -> "b");
+        assertThat(map.size(), is(2));
+        map.compute(B01, (key, value) -> null);
+        assertThat(map.size(), is(1));
+        assertThat(map.get(A01), is("a"));
+        assertThat(map.get(B01), is(nullValue()));
+        map.compute(B01, (key, value) -> "b again");
+        assertThat(map.size(), is(2));
+        assertThat(map.get(B01), is("b again"));
+    }
+
+    @Test
+    public void testMatchesGetThenPut() {
+        IIngredientMapMutable<ComplexStack, Integer, String> computed = this.factory.get();
+        IIngredientMapMutable<ComplexStack, Integer, String> manual = this.factory.get();
+        ComplexStack[] keys = {A01, A12, B01, A01, B01, A12};
+        for (int i = 0; i < keys.length; i++) {
+            ComplexStack key = keys[i];
+            int step = i;
+            computed.compute(key, (k, value) -> step % 3 == 2 ? null : (value == null ? "v" + step : value + "v" + step));
+            String oldValue = manual.get(key);
+            String newValue = step % 3 == 2 ? null : (oldValue == null ? "v" + step : oldValue + "v" + step);
+            if (newValue == null) {
+                manual.remove(key);
+            } else {
+                manual.put(key, newValue);
+            }
+            assertThat(computed.size(), is(manual.size()));
+            for (ComplexStack probe : new ComplexStack[]{A01, A12, B01}) {
+                assertThat(computed.get(probe), is(manual.get(probe)));
+            }
+        }
+    }
+}