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

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

@@ -23,10 +23,12 @@
  * questions.
  */
 
 package java.util;
 
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
 /**
  * A Red-Black tree based {@link NavigableMap} implementation.
  * The map is sorted according to the {@linkplain Comparable natural

@@ -943,10 +945,64 @@
      */
     public SortedMap<K,V> tailMap(K fromKey) {
         return tailMap(fromKey, true);
     }
 
+    @Override
+    public void forEach(BiConsumer<? super K, ? super V> action) {
+        Objects.requireNonNull(action);
+        int expectedModCount = modCount;
+        TreeMap.Entry<K,V> e, p, pl;
+
+        if ((e = getFirstEntry()) != null ) {
+            do {
+                action.accept(e.key, e.value);
+                if ((p = e.right) != null) {
+                    // go right for successors
+                    while ((pl = p.left) != null)
+                        // then bear left for closest successor
+                        p = pl;
+                } else {
+                    // find root or parent where we are not the right child
+                    while ((p = e.parent) != null && e == p.right)
+                        e = p;
+                }
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
+            } while ((e = p) != null);
+        }
+    }
+
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        Objects.requireNonNull(function);
+        int expectedModCount = modCount;
+        TreeMap.Entry<K,V> e, p, pl;
+
+        if ((e = getFirstEntry()) != null ) {
+            do {
+                e.value = Objects.requireNonNull(function.apply(e.key, e.value));
+                if ((p = e.right) != null) {
+                    // go right for successors
+                    while ((pl = p.left) != null)
+                        // then bear left for closest successor
+                        p = pl;
+                } else {
+                    // find root or parent where we are not the right child
+                    while ((p = e.parent) != null && e == p.right)
+                        e = p;
+                }
+
+                if (expectedModCount != modCount) {
+                    throw new ConcurrentModificationException();
+                }
+            } while ((e = p) != null);
+        }
+    }
+
     // View class support
 
     class Values extends AbstractCollection<V> {
         public Iterator<V> iterator() {
             return new ValueIterator(getFirstEntry());