single lookup implementation of compute methods in maps (#259)

agrona/src/main/java/org/agrona/collections/Int2IntHashMap.java

@@ -20,6 +20,7 @@
 import java.util.*;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
+import java.util.function.IntBinaryOperator;
 import java.util.function.IntUnaryOperator;
 
 import static org.agrona.BitUtil.findNextPositivePowerOfTwo;
@@ -354,19 +355,114 @@ public void compact()
      */
     public int computeIfAbsent(final int key, final IntUnaryOperator mappingFunction)
     {
-        int value = get(key);
-        if (missingValue == value)
+        final int missingValue = this.missingValue;
+        final int[] entries = this.entries;
+        @DoNotSub final int mask = entries.length - 1;
+        @DoNotSub int index = Hashing.evenHash(key, mask);
+
+        int value;
+        while (missingValue != (value = entries[index + 1]))
         {
-            value = mappingFunction.applyAsInt(key);
-            if (missingValue != value)
+            if (key == entries[index])
             {
-                put(key, value);
+                break;
             }
+
+            index = next(index, mask);
         }
+        if (value == missingValue && (value = mappingFunction.applyAsInt(key)) != missingValue)
+        {
+            entries[index] = key;
+            entries[index + 1] = value;
+            size++;
+            increaseCapacity();
+        }
+
+        return value;
+    }
+
+    /**
+     * Primitive specialised version of {@link java.util.Map#computeIfPresent}
+     *
+     * @param key               to search on.
+     * @param remappingFunction to compute a value if a mapping is found.
+     * @return the updated value if a mapping was found, otherwise the missing value.
+     */
+    public int computeIfPresent(final int key, final IntBinaryOperator remappingFunction)
+    {
+        final int missingValue = this.missingValue;
+        final int[] entries = this.entries;
+        @DoNotSub final int mask = entries.length - 1;
+        @DoNotSub int index = Hashing.evenHash(key, mask);
+
+        int value;
+        while (missingValue != (value = entries[index + 1]))
+        {
+            if (key == entries[index])
+            {
+                break;
+            }
 
+            index = next(index, mask);
+        }
+        if (value != missingValue)
+        {
+            value = remappingFunction.applyAsInt(key, value);
+            entries[index + 1] = value;
+            if (value == missingValue)
+            {
+                size--;
+                compactChain(index);
+            }
+        }
         return value;
     }
 
+    /**
+     * Primitive specialised version of {@link java.util.Map#compute}
+     *
+     * @param key               to search on.
+     * @param remappingFunction to compute a value.
+     * @return the updated value.
+     */
+    public int compute(final int key, final IntBinaryOperator remappingFunction)
+    {
+        final int missingValue = this.missingValue;
+        final int[] entries = this.entries;
+        @DoNotSub final int mask = entries.length - 1;
+        @DoNotSub int index = Hashing.evenHash(key, mask);
+
+        int oldValue;
+        while (missingValue != (oldValue = entries[index + 1]))
+        {
+            if (key == entries[index])
+            {
+                break;
+            }
+
+            index = next(index, mask);
+        }
+        final int newValue = remappingFunction.applyAsInt(key, oldValue);
+        if (newValue != missingValue)
+        {
+            // add or replace old mapping
+            entries[index + 1] = newValue;
+            if (oldValue == missingValue)
+            {
+                entries[index] = key;
+                size++;
+                increaseCapacity();
+            }
+        }
+        else if (oldValue != missingValue)
+        {
+            entries[index + 1] = missingValue;
+            size--;
+            compactChain(index);
+        }
+        return newValue;
+    }
+
     // ---------------- Boxed Versions Below ----------------
 
     /**

agrona/src/main/java/org/agrona/collections/Int2ObjectHashMap.java

@@ -295,8 +295,8 @@ protected V getMapped(final int key)
     }
 
     /**
-     * Get a value for a given key, or if it does not exist then default the value via a
-     * {@link java.util.function.IntFunction} and put it in the map.
+     * Get a value for a given key, or if it does not exist then default the value
+     * via a {@link java.util.function.IntFunction} and put it in the map.
      * <p>
      * Primitive specialized version of {@link java.util.Map#computeIfAbsent}.
      *
@@ -306,23 +306,137 @@ protected V getMapped(final int key)
      */
     public V computeIfAbsent(final int key, final IntFunction<? extends V> mappingFunction)
     {
-        V value = getMapped(key);
-        if (null == value)
+        final int[] keys = this.keys;
+        final Object[] values = this.values;
+        @DoNotSub final int mask = values.length - 1;
+        @DoNotSub int index = Hashing.hash(key, mask);
+
+        Object mappedValue;
+        while (null != (mappedValue = values[index]))
         {
-            value = mappingFunction.apply(key);
-            if (null != value)
+            if (key == keys[index])
             {
-                put(key, value);
+                break;
             }
+
+            index = ++index & mask;
         }
-        else
+        V value = unmapNullValue(mappedValue);
+        if (value == null && (value = mappingFunction.apply(key)) != null)
         {
-            value = unmapNullValue(value);
+            values[index] = value;
+            if (mappedValue == null)
+            {
+                keys[index] = key;
+                if (++size > resizeThreshold)
+                {
+                    increaseCapacity();
+                }
+            }
         }
 
         return value;
     }
 
+    /**
+     * If the value for the specified key is present and non-null, attempts to compute a new
+     * mapping given the key and its current mapped value.
+     * <p>
+     * If the function returns {@code null}, the mapping is removed
+     * <p>
+     * Primitive specialized version of {@link java.util.Map#computeIfPresent}.
+     *
+     * @param key               to search on.
+     * @param remappingFunction to provide a value if the get returns missingValue.
+     * @return the new value associated with the specified key, or {@code null} if none
+     */
+    public V computeIfPresent(final int key, final IntObjToObjFunction<? super V, ? extends V> remappingFunction)
+    {
+        final int[] keys = this.keys;
+        final Object[] values = this.values;
+        @DoNotSub final int mask = values.length - 1;
+        @DoNotSub int index = Hashing.hash(key, mask);
+
+        Object mappedValue;
+        while (null != (mappedValue = values[index]))
+        {
+            if (key == keys[index])
+            {
+                break;
+            }
+
+            index = ++index & mask;
+        }
+        V value = unmapNullValue(mappedValue);
+        if (value != null)
+        {
+            value = remappingFunction.apply(key, value);
+            values[index] = value;
+            if (value == null)
+            {
+                --size;
+                compactChain(index);
+            }
+        }
+        return value;
+    }
+
+    /**
+     * Attempts to compute a mapping for the specified key and its current mapped
+     * value (or {@code null} if there is no current mapping).
+     * <p>
+     * If the function returns {@code null}, the mapping is removed (or remains
+     * absent if initially absent).
+     * <p>
+     * Primitive specialized version of {@link java.util.Map#compute}.
+     *
+     * @param key               to search on.
+     * @param remappingFunction to provide a value if the get returns missingValue.
+     * @return the new value associated with the specified key, or {@code null} if none
+     */
+    public V compute(final int key, final IntObjToObjFunction<? super V, ? extends V> remappingFunction)
+    {
+        final int[] keys = this.keys;
+        final Object[] values = this.values;
+        @DoNotSub final int mask = values.length - 1;
+        @DoNotSub int index = Hashing.hash(key, mask);
+
+        Object mappedvalue;
+        while (null != (mappedvalue = values[index]))
+        {
+            if (key == keys[index])
+            {
+                break;
+            }
+
+            index = ++index & mask;
+        }
+        final V oldValue = unmapNullValue(mappedvalue);
+        final V newValue = remappingFunction.apply(key, oldValue);
+
+        if (newValue != null)
+        {
+            // add or replace old mapping
+            values[index] = newValue;
+            if (mappedvalue == null)
+            {
+                keys[index] = key;
+                if (++size > resizeThreshold)
+                {
+                    increaseCapacity();
+                }
+            }
+        }
+        else if (mappedvalue != null || containsKey(key))
+        {
+            // delete mapping
+            values[index] = null;
+            size--;
+            compactChain(index);
+        }
+        return newValue;
+    }
+
     /**
      * {@inheritDoc}
      */

agrona/src/main/java/org/agrona/collections/IntObjToObjFunction.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2014-2022 Real Logic Limited.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.agrona.collections;
+
+/**
+ * This is an (int, Object) -> Object primitive specialisation of a BiFunction.
+ */
+@FunctionalInterface
+public interface
+    IntObjToObjFunction<T, R>
+{
+    /**
+     * Applies this function to the given arguments.
+     *
+     * @param i the second function argument
+     * @param t the first function argument
+     * @return the function result
+     */
+    R apply(int i, T t);
+}

agrona/src/main/java/org/agrona/collections/ObjIntToIntFunction.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2014-2022 Real Logic Limited.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.agrona.collections;
+
+/**
+ * This is an (Object, int) -> int primitive specialisation of a BiFunction.
+ */
+@FunctionalInterface
+public interface
+    ObjIntToIntFunction<T>
+{
+    /**
+     * Applies this function to the given arguments.
+     *
+     * @param t the first function argument
+     * @param i the second function argument
+     * @return the function result
+     */
+    int apply(T t, int i);
+}