# HG changeset patch
# User mduigou
# Date 1371244514 25200
# Node ID 7813252a5f8c7fb8e8257582956a4639a765b80d
# Parent  ce8fbca80bbc0abb8d3ffa2fc71af17d1ac97be2
8016446: Improve forEach/replaceAll for Map, HashMap, Hashtable, IdentityHashMap, WeakHashMap, TreeMap
Reviewed-by: forax, duigou
Contributed-by: Mike Duigou <mike.duigou@oracle.com>, Remi Forax <forax@univ-mlv.fr>

diff --git a/src/share/classes/java/util/HashMap.java b/src/share/classes/java/util/HashMap.java
--- a/src/share/classes/java/util/HashMap.java
+++ b/src/share/classes/java/util/HashMap.java
@@ -29,6 +29,7 @@
 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,111 @@
      */
     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();
+                }
+            }
+        }
     }
 
-    // optimized implementations of default methods in Map
+    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) {
+                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();
+            }
+        }
+    }
+
+    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 +2399,12 @@
             if (e == null)
                 throw new NoSuchElementException();
 
-            if (e instanceof Entry) {
+            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;
-            } else { // TreeBin
-                retVal = (Entry<K,V>)((TreeNode)e).entry;
-                next = retVal.next;
             }
 
             if (next == null) { // Move to next bin
diff --git a/src/share/classes/java/util/Hashtable.java b/src/share/classes/java/util/Hashtable.java
--- a/src/share/classes/java/util/Hashtable.java
+++ b/src/share/classes/java/util/Hashtable.java
@@ -944,7 +944,16 @@
     @Override
     public synchronized void replaceAll(
             BiFunction<? super K, ? super V, ? extends V> function) {
-        Map.super.replaceAll(function);
+        Objects.requireNonNull(function);     // explicit check required in case
+                                              // table is empty.
+        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;
+            }
+        }
     }
 
     @Override
@@ -1058,7 +1067,7 @@
     }
 
     @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 +1096,7 @@
     }
 
     @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 +1131,7 @@
     }
 
     @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;
diff --git a/src/share/classes/java/util/IdentityHashMap.java b/src/share/classes/java/util/IdentityHashMap.java
--- a/src/share/classes/java/util/IdentityHashMap.java
+++ b/src/share/classes/java/util/IdentityHashMap.java
@@ -27,6 +27,8 @@
 
 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,38 @@
         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) {
+            if (t[index] != null) {
+                action.accept((K) unmaskNull(t[index]), (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) {
+            if (t[index] != null) {
+                t[index + 1] = function.apply((K) unmaskNull(t[index]), (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.
diff --git a/src/share/classes/java/util/LinkedHashMap.java b/src/share/classes/java/util/LinkedHashMap.java
--- a/src/share/classes/java/util/LinkedHashMap.java
+++ b/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,30 @@
         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) {
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+           action.accept(entry.key, entry.value);
+        }
+    }
+
+    @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) {
+            if (expectedModCount != modCount) {
+                throw new ConcurrentModificationException();
+            }
+           entry.value = function.apply(entry.key, entry.value);
+        }
+    }
+
     /**
      * LinkedHashMap entry.
      */
diff --git a/src/share/classes/java/util/Map.java b/src/share/classes/java/util/Map.java
--- a/src/share/classes/java/util/Map.java
+++ b/src/share/classes/java/util/Map.java
@@ -545,6 +545,7 @@
                 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 @@
                 k = entry.getKey();
                 v = entry.getValue();
             } catch(IllegalStateException ise) {
+                // this usually means the entry is no longer in the map.
                 throw new ConcurrentModificationException(ise);
             }
-            entry.setValue(function.apply(k, v));
+
+            // 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);
+            }
         }
     }
 
diff --git a/src/share/classes/java/util/TreeMap.java b/src/share/classes/java/util/TreeMap.java
--- a/src/share/classes/java/util/TreeMap.java
+++ b/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 @@
         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> {
diff --git a/src/share/classes/java/util/WeakHashMap.java b/src/share/classes/java/util/WeakHashMap.java
--- a/src/share/classes/java/util/WeakHashMap.java
+++ b/src/share/classes/java/util/WeakHashMap.java
@@ -28,6 +28,8 @@
 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 @@
         }
     }
 
+    @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.