Add TrieMap.computeIfAbsent()

This is another logical operation -- combining lookup with insert. This patch adds explicit support for it, so it happens in a single pass (most of the time). Note that concurrent activity may lead to the value being computed multiple times. Nevertheless fixes #147. Signed-off-by: Robert Varga <[email protected]>
PANTHEONtech · Jan 27, 2025 · 1d9e15f · 1d9e15f
1 parent 0842609
commit 1d9e15f
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 1 deletion.
diff --git a/triemap/src/main/java/tech/pantheon/triemap/CNode.java b/triemap/src/main/java/tech/pantheon/triemap/CNode.java
@@ -23,6 +23,7 @@
 import static tech.pantheon.triemap.Result.RESTART;
 
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Function;
 import org.eclipse.jdt.annotation.NonNull;
 import org.eclipse.jdt.annotation.Nullable;
 
@@ -91,6 +92,47 @@ static <K, V> MainNode<K, V> dual(final SNode<K, V> first, final @NonNull K key,
         }
     }
 
+    @Nullable Object computeIfAbsent(final MutableTrieMap<K, V> ct, final Gen startGen, final int hc,
+            final @NonNull K key, final @NonNull Function<? super K, ? extends V> fn, final int lev,
+            final INode<K, V> parent) {
+        final int idx = hc >>> lev & 0x1f;
+        final int flag = 1 << idx;
+        final int mask = flag - 1;
+        final int pos = Integer.bitCount(bitmap & mask);
+
+        if ((bitmap & flag) == 0) {
+            final var val = fn.apply(key);
+            return val == null || insert(ct, parent, pos, flag, key, val, hc) ? val : RESTART;
+        }
+
+        // 1a) insert below
+        final var cnAtPos = array[pos];
+        if (cnAtPos instanceof INode<K, V> in) {
+            // try to renew if needed and enter next level
+            return startGen != in.gen && !renew(ct, parent, startGen)
+                ? RESTART : in.computeIfAbsent(ct, startGen, hc, key, fn, lev + LEVEL_BITS, parent);
+        } else if (cnAtPos instanceof SNode<K, V> sn) {
+            return computeIfAbsent(ct, parent, pos, sn, key, fn, hc, lev);
+        } else {
+            throw invalidElement(cnAtPos);
+        }
+    }
+
+    private @Nullable Object computeIfAbsent(final MutableTrieMap<K, V> ct, final INode<K, V> in, final int pos,
+            final SNode<K, V> sn, final @NonNull K key, final @NonNull Function<? super K, ? extends V> fn,
+            final int hc, final int lev) {
+        if (sn.matches(hc, key)) {
+            return sn.value();
+        }
+        final var val = fn.apply(key);
+        if (val == null) {
+            return null;
+        }
+
+        final var rn = gen == in.gen ? this : renewed(ct, gen);
+        return in.gcasWrite(ct, rn.updatedAt(pos, new INode<>(in, sn, key, val, hc, lev), gen)) ? val : RESTART;
+    }
+
     @Nullable Object lookup(final TrieMap<K, V> ct, final Gen startGen, final int hc, final @NonNull K key,
             final int lev, final INode<K, V> parent) {
         // 1) a multinode

diff --git a/triemap/src/main/java/tech/pantheon/triemap/INode.java b/triemap/src/main/java/tech/pantheon/triemap/INode.java
@@ -22,6 +22,7 @@
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.VarHandle;
+import java.util.function.Function;
 import org.eclipse.jdt.annotation.NonNull;
 import org.eclipse.jdt.annotation.NonNullByDefault;
 import org.eclipse.jdt.annotation.Nullable;
@@ -254,6 +255,22 @@ public int elementSize(final ImmutableTrieMap<K, V> ct) {
         }
     }
 
+    @Nullable Object computeIfAbsent(final MutableTrieMap<K, V> ct, final Gen startGen, final int hc,
+            final @NonNull K key, final @NonNull Function<? super K, ? extends V> fn, final int lev,
+            final INode<K, V> parent) {
+        final var m = gcasRead(ct);
+        if (m instanceof CNode<K, V> cn) {
+            return cn.computeIfAbsent(ct, startGen, hc, key, fn, lev, this);
+        } else if (m instanceof TNode) {
+            clean(ct, parent, lev);
+            return RESTART;
+        } else if (m instanceof LNode<K, V> ln) {
+            return ln.entries.computeIfAbsent(ct, this, ln, key, fn);
+        } else {
+            throw invalidElement(m);
+        }
+    }
+
     /**
      * Inserts a key value pair, overwriting the old pair if the keys match.
      *

diff --git a/triemap/src/main/java/tech/pantheon/triemap/LNodeEntries.java b/triemap/src/main/java/tech/pantheon/triemap/LNodeEntries.java
@@ -19,6 +19,7 @@
 import static tech.pantheon.triemap.PresencePredicate.PRESENT;
 import static tech.pantheon.triemap.Result.RESTART;
 
+import java.util.function.Function;
 import org.eclipse.jdt.annotation.NonNull;
 import org.eclipse.jdt.annotation.Nullable;
 
@@ -86,6 +87,17 @@ static <K,V> LNodeEntries<K, V> of(final @NonNull K k1, final @NonNull V v1,
         return entry != null ? entry.value() : null;
     }
 
+    @Nullable Object computeIfAbsent(final MutableTrieMap<K, V> ct, final INode<K, V> in, final LNode<K, V> ln,
+            final @NonNull K key, final @NonNull Function<? super K, ? extends V> fn) {
+        final var entry = findEntry(key);
+        if (entry != null) {
+            return entry.value();
+        }
+
+        final var val = fn.apply(key);
+        return val == null || in.gcasWrite(ct, toInserted(ln, key, val)) ? val : RESTART;
+    }
+
     final boolean insert(final MutableTrieMap<K, V> ct, final INode<K, V> in, final LNode<K, V> ln,
             final @NonNull K key, final @NonNull V val) {
         final var entry = findEntry(key);

diff --git a/triemap/src/main/java/tech/pantheon/triemap/MutableTrieMap.java b/triemap/src/main/java/tech/pantheon/triemap/MutableTrieMap.java
@@ -23,6 +23,7 @@
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.VarHandle;
+import java.util.function.Function;
 import org.eclipse.jdt.annotation.NonNull;
 import org.eclipse.jdt.annotation.Nullable;
 
@@ -138,6 +139,24 @@ public V replace(final K key, final V value) {
         return (V) res;
     }
 
+    @Override
+    @SuppressWarnings("unchecked")
+    public V computeIfAbsent(final K key, final Function<? super K, ? extends V> mappingFunction) {
+        final var k = requireNonNull(key);
+        final int hc = computeHash(key);
+        final var fn = requireNonNull(mappingFunction);
+
+        // Keep looping as long as RESTART is being returned
+        Object res;
+        do {
+            // Keep looping as long as we do not get a reply
+            final var r = readRoot();
+            res = r.computeIfAbsent(this, r.gen, hc, k, fn, 0, null);
+        } while (res == RESTART);
+
+        return (V) res;
+    }
+
     @Override
     public int size() {
         return immutableSnapshot().size();

diff --git a/triemap/src/main/java/tech/pantheon/triemap/TrieMap.java b/triemap/src/main/java/tech/pantheon/triemap/TrieMap.java
@@ -22,6 +22,7 @@
 import java.util.AbstractMap;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 
 /**
  * This is a port of Scala's TrieMap class from the Scala Collections library. This implementation does not support
@@ -144,6 +145,9 @@ public final V get(final Object key) {
     @Override
     public abstract V replace(K key, V value);
 
+    @Override
+    public abstract V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction);
+
     @Override
     public abstract int size();
 

diff --git a/triemap/src/test/java/tech/pantheon/triemap/TestConcurrentMapComputeIfAbsent.java b/triemap/src/test/java/tech/pantheon/triemap/TestConcurrentMapComputeIfAbsent.java
@@ -16,6 +16,7 @@
 package tech.pantheon.triemap;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -38,7 +39,7 @@ void testConcurrentMapComputeIfAbsentDoesNotComputeIfPresent() {
         for (int i = 0; i < COUNT; i++) {
             map.put(i, Integer.toString(i));
             assertEquals(Integer.toString(i), map.computeIfAbsent(i,
-                    (ignored) -> fail("Should not have called function")));
+                    ignored -> fail("Should not have called function")));
         }
     }
 
@@ -63,4 +64,11 @@ void testConflictingHash() {
         // Check with equivalent key
         assertSame(v3, map.computeIfAbsent(k3dup, k -> v3));
     }
+
+    @Test
+    void testComputeNull() {
+        final var map = TrieMap.create();
+        assertNull(map.computeIfAbsent("key", k -> null));
+        assertEquals("{}", map.toString());
+    }
 }