8016446: Improve forEach/replaceAll for Map, HashMap, Hashtable,...

8016446: Improve forEach/replaceAll for Map, HashMap, Hashtable, IdentityHashMap, WeakHashMap, TreeMap, ConcurrentMap
Reviewed-by: forax, mduigou, psandoz
Contributed-by: NMike Duigou &lt;mike.duigou@oracle.com&gt;, Remi Forax <forax@univ-mlv.fr>

src/share/classes/java/util/HashMap.java

@@ -29,6 +29,7 @@ import java.io.*;
 import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -1299,10 +1300,112 @@ public class HashMap<K,V>
      */
     public V remove(Object key) {
         Entry<K,V> e = removeEntryForKey(key);
-        return (e == null ? null : e.value);
+       return (e == null ? null : e.value);
+   }
+
+   // optimized implementations of default methods in Map
+
+    @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        Objects.requireNonNull(action);
+        final int expectedModCount = modCount;
+        if (nullKeyEntry != null) {
+            forEachNullKey(expectedModCount, action);
+        }
+        Object[] tab = this.table;
+        for (int index = 0; index < tab.length; index++) {
+            Object item = tab[index];
+            if (item == null) {
+                continue;
+            }
+            if (item instanceof HashMap.TreeBin) {
+                eachTreeNode(expectedModCount, ((TreeBin)item).first, action);
+                continue;
+            }
+            @SuppressWarnings("unchecked")
+            Entry<K, V> entry = (Entry<K, V>)item;
+            while (entry != null) {
+                action.accept(entry.key, entry.value);
+                entry = (Entry<K, V>)entry.next;
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
+            }
+        }
+    }
+
+    private void eachTreeNode(int expectedModCount, TreeNode<K, V> node, BiConsumer<? super K, ? super V> action) {
+        while (node != null) {
+            @SuppressWarnings("unchecked")
+            Entry<K, V> entry = (Entry<K, V>)node.entry;
+            action.accept(entry.key, entry.value);
+            node = (TreeNode<K, V>)entry.next;
+
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
+    }
+
+    private void forEachNullKey(int expectedModCount, BiConsumer<? super K, ? super V> action) {
+        action.accept(null, nullKeyEntry.value);
+
+        if (expectedModCount != modCount) {
+            throw new ConcurrentModificationException();
+        }
+    }
+
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);
+        final int expectedModCount = modCount;
+        if (nullKeyEntry != null) {
+            replaceforNullKey(expectedModCount, function);
+        }
+        Object[] tab = this.table;
+        for (int index = 0; index < tab.length; index++) {
+            Object item = tab[index];
+            if (item == null) {
+                continue;
+            }
+            if (item instanceof HashMap.TreeBin) {
+                replaceEachTreeNode(expectedModCount, ((TreeBin)item).first, function);
+                continue;
+            }
+            @SuppressWarnings("unchecked")
+            Entry<K, V> entry = (Entry<K, V>)item;
+            while (entry != null) {
+                entry.value = function.apply(entry.key, entry.value);
+                entry = (Entry<K, V>)entry.next;
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
+            }
+        }
+    }
+
+    private void replaceEachTreeNode(int expectedModCount, TreeNode<K, V> node, BiFunction<? super K, ? super V, ? extends V> function) {
+        while (node != null) {
+            @SuppressWarnings("unchecked")
+            Entry<K, V> entry = (Entry<K, V>)node.entry;
+            entry.value = function.apply(entry.key, entry.value);
+            node = (TreeNode<K, V>)entry.next;
+
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
     }
 
-    // optimized implementations of default methods in Map
+    private void replaceforNullKey(int expectedModCount, BiFunction<? super K, ? super V, ? extends V> function) {
+        nullKeyEntry.value = function.apply(null, nullKeyEntry.value);
+
+        if (expectedModCount != modCount) {
+            throw new ConcurrentModificationException();
+        }
+    }
 
     @Override
     public V putIfAbsent(K key, V value) {
@@ -2297,12 +2400,12 @@ public class HashMap<K,V>
             if (e == null)
                 throw new NoSuchElementException();
 
-            if (e instanceof Entry) {
-                retVal = (Entry<K,V>)e;
-                next = ((Entry<K,V>)e).next;
-            } else { // TreeBin
+            if (e instanceof TreeNode) { // TreeBin
                 retVal = (Entry<K,V>)((TreeNode)e).entry;
                 next = retVal.next;
+            } else {
+                retVal = (Entry<K,V>)e;
+                next = ((Entry<K,V>)e).next;
             }
 
             if (next == null) { // Move to next bin

src/share/classes/java/util/Hashtable.java

@@ -932,19 +932,39 @@ public class Hashtable<K,V>
     public synchronized void forEach(BiConsumer<? super K, ? super V> action) {
         Objects.requireNonNull(action);     // explicit check required in case
                                             // table is empty.
-        Entry<?,?>[] tab = table;
-        for (Entry<?,?> entry : tab) {
+        final int expectedModCount = modCount;
+
+        Entry<?, ?>[] tab = table;
+        for (Entry<?, ?> entry : tab) {
             while (entry != null) {
                 action.accept((K)entry.key, (V)entry.value);
                 entry = entry.next;
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
             }
         }
     }
 
     @Override
-    public synchronized void replaceAll(
-            BiFunction<? super K, ? super V, ? extends V> function) {
-        Map.super.replaceAll(function);
+    public synchronized void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);     // explicit check required in case
+                                              // table is empty.
+        final int expectedModCount = modCount;
+
+        Entry<K, V>[] tab = (Entry<K, V>[])table;
+        for (Entry<K, V> entry : tab) {
+            while (entry != null) {
+                entry.value = Objects.requireNonNull(
+                    function.apply(entry.key, entry.value));
+                entry = entry.next;
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
+            }
+        }
     }
 
     @Override
@@ -1058,7 +1078,7 @@ public class Hashtable<K,V>
     }
 
     @Override
-    public V computeIfPresent(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+    public synchronized V computeIfPresent(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
         Objects.requireNonNull(remappingFunction);
 
         Entry<?,?> tab[] = table;
@@ -1087,7 +1107,7 @@ public class Hashtable<K,V>
     }
 
     @Override
-    public V compute(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+    public synchronized V compute(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
         Objects.requireNonNull(remappingFunction);
 
         Entry<?,?> tab[] = table;
@@ -1122,7 +1142,7 @@ public class Hashtable<K,V>
     }
 
     @Override
-    public V merge(K key, V value, BiFunction<? super V, ? super V, ? extends V> remappingFunction) {
+    public synchronized V merge(K key, V value, BiFunction<? super V, ? super V, ? extends V> remappingFunction) {
         Objects.requireNonNull(remappingFunction);
 
         Entry<?,?> tab[] = table;

src/share/classes/java/util/IdentityHashMap.java

@@ -27,6 +27,8 @@ package java.util;
 
 import java.io.*;
 import java.lang.reflect.Array;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
 /**
@@ -1337,6 +1339,42 @@ public class IdentityHashMap<K,V>
         tab[i + 1] = value;
     }
 
+    @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        Objects.requireNonNull(action);
+        int expectedModCount = modCount;
+
+        Object[] t = table;
+        for (int index = 0; index < t.length; index += 2) {
+            Object k = t[index];
+            if (k != null) {
+                action.accept((K) unmaskNull(k), (V) t[index + 1]);
+            }
+
+            if (modCount != expectedModCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
+    }
+
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);
+        int expectedModCount = modCount;
+
+        Object[] t = table;
+        for (int index = 0; index < t.length; index += 2) {
+            Object k = t[index];
+            if (k != null) {
+                t[index + 1] = function.apply((K) unmaskNull(k), (V) t[index + 1]);
+            }
+
+            if (modCount != expectedModCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
+    }
+
     /**
      * Similar form as array-based Spliterators, but skips blank elements,
      * and guestimates size as decreasing by half per split.

src/share/classes/java/util/LinkedHashMap.java

@@ -25,6 +25,8 @@
 
 package java.util;
 import java.io.*;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 
 /**
  * <p>Hash table and linked list implementation of the <tt>Map</tt> interface,
@@ -296,6 +298,32 @@ public class LinkedHashMap<K,V>
         header.before = header.after = header;
     }
 
+    @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        Objects.requireNonNull(action);
+        int expectedModCount = modCount;
+        for (Entry<K, V> entry = header.after; entry != header; entry = entry.after) {
+            action.accept(entry.key, entry.value);
+
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
+    }
+
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);
+        int expectedModCount = modCount;
+        for (Entry<K, V> entry = header.after; entry != header; entry = entry.after) {
+            entry.value = function.apply(entry.key, entry.value);
+
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
+    }
+
     /**
      * LinkedHashMap entry.
      */

src/share/classes/java/util/Map.java

@@ -545,6 +545,7 @@ public interface Map<K,V> {
                 k = entry.getKey();
                 v = entry.getValue();
             } catch(IllegalStateException ise) {
+                // this usually means the entry is no longer in the map.
                 throw new ConcurrentModificationException(ise);
             }
             action.accept(k, v);
@@ -599,9 +600,19 @@ public interface Map<K,V> {
                 k = entry.getKey();
                 v = entry.getValue();
             } catch(IllegalStateException ise) {
+                // this usually means the entry is no longer in the map.
+                throw new ConcurrentModificationException(ise);
+            }
+
+            // ise thrown from function is not a cme.
+            v = function.apply(k, v);
+
+            try {
+                entry.setValue(v);
+            } catch(IllegalStateException ise) {
+                // this usually means the entry is no longer in the map.
                 throw new ConcurrentModificationException(ise);
             }
-            entry.setValue(function.apply(k, v));
         }
     }
 

src/share/classes/java/util/TreeMap.java

@@ -25,6 +25,8 @@
 
 package java.util;
 
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
 /**
@@ -945,6 +947,33 @@ public class TreeMap<K,V>
         return tailMap(fromKey, true);
     }
 
+    @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        Objects.requireNonNull(action);
+        int expectedModCount = modCount;
+        for (Entry<K, V> e = getFirstEntry(); e != null; e = successor(e)) {
+            action.accept(e.key, e.value);
+
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
+    }
+
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);
+        int expectedModCount = modCount;
+
+        for (Entry<K, V> e = getFirstEntry(); e != null; e = successor(e)) {
+            e.value = Objects.requireNonNull(function.apply(e.key, e.value));
+
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+        }
+    }
+
     // View class support
 
     class Values extends AbstractCollection<V> {

src/share/classes/java/util/WeakHashMap.java

@@ -28,6 +28,8 @@ package java.util;
 import java.lang.ref.WeakReference;
 import java.lang.ref.ReferenceQueue;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
 
@@ -1036,6 +1038,48 @@ public class WeakHashMap<K,V>
         }
     }
 
+    @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        Objects.requireNonNull(action);
+        int expectedModCount = modCount;
+
+        Entry<K, V>[] tab = getTable();
+        for (Entry<K, V> entry : tab) {
+            while (entry != null) {
+                Object key = entry.get();
+                if (key != null) {
+                    action.accept((K)WeakHashMap.unmaskNull(key), entry.value);
+                }
+                entry = entry.next;
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
+            }
+        }
+    }
+
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);
+        int expectedModCount = modCount;
+
+        Entry<K, V>[] tab = getTable();;
+        for (Entry<K, V> entry : tab) {
+            while (entry != null) {
+                Object key = entry.get();
+                if (key != null) {
+                    entry.value = function.apply((K)WeakHashMap.unmaskNull(key), entry.value);
+                }
+                entry = entry.next;
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
+            }
+        }
+    }
+
     /**
      * Similar form as other hash Spliterators, but skips dead
      * elements.

src/share/classes/java/util/concurrent/ConcurrentMap.java

@@ -35,6 +35,8 @@
 
 package java.util.concurrent;
 import java.util.Map;
+import java.util.Objects;
+import java.util.function.BiFunction;
 
 /**
  * A {@link java.util.Map} providing additional atomic
@@ -183,4 +185,26 @@ public interface ConcurrentMap<K, V> extends Map<K, V> {
      *         or value prevents it from being stored in this map
      */
     V replace(K key, V value);
+
+    /**
+     * {@inheritDoc}
+     *
+     * @implNote This implementation assumes that the ConcurrentMap cannot
+     * contain null values and get() returning null unambiguously means the key
+     * is absent. Implementations which support null values
+     * <strong>must</strong> override this default implementation.
+     */
+    @Override
+    default void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);
+        forEach((k,v) -> {
+            while(!replace(k, v, function.apply(k, v))) {
+                // v changed or k is gone
+                if( (v = get(k)) == null) {
+                    // k is no longer in the map.
+                    break;
+                }
+            }
+        });
+    }
 }