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

rev 7312 : 8016446: Add override forEach/replaceAll to HashMap, Hashtable, IdentityHashMap, WeakHashMap, TreeMap
Reviewed-by: duke

@@ -27,10 +27,11 @@
 
 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;
 
 /**

@@ -1303,10 +1304,72 @@
     }
 
     // optimized implementations of default methods in Map
 
     @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        Objects.requireNonNull(action);
+        int expectedModCount = modCount;
+
+        // handle null key entry
+        if(null != nullKeyEntry) {
+            action.accept(null, nullKeyEntry.value);
+        }
+
+        Object[] t = table;
+        for(Object next : t) {
+            if (next instanceof HashMap.TreeBin) {
+                // Point to the first TreeNode
+                TreeNode<K,V> node = ((TreeBin) next).first;
+
+                while(node != null) {
+                    action.accept(node.entry.key, node.entry.value);
+                    node = (TreeNode<K,V>) node.entry.next;
+                }
+            } else while(next != null) {
+                Entry<K,V> entry = (Entry<K,V>) next;
+                action.accept(entry.key, entry.value);
+                next = 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;
+
+        // handle null key entry
+        if(null != nullKeyEntry) {
+            nullKeyEntry.value = function.apply(null, nullKeyEntry.value);
+        }
+
+        Object[] t = table;
+        for(Object next : t) {
+            if (next instanceof HashMap.TreeBin) {
+                // Point to the first TreeNode
+                TreeNode<K,V> node = ((TreeBin) next).first;
+
+                while(node != null) {
+                    node.entry.value = function.apply(node.entry.key, node.entry.value);
+                    node = (TreeNode<K,V>) node.entry.next;
+                }
+            } else while(next != null) {
+                Entry<K,V> entry = (Entry<K,V>) next;
+                entry.value = function.apply(entry.key, entry.value);
+                next = entry.next;
+            }
+        }
+
+        if (expectedModCount != modCount)
+            throw new ConcurrentModificationException();
+    }
+
+    @Override
     public V putIfAbsent(K key, V value) {
         if (table == EMPTY_TABLE) {
             inflateTable(threshold);
         }
         if (key == null) {

@@ -2295,16 +2358,16 @@
             Entry<K,V> retVal;
 
             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
                 findNextBin();
             }