< prev index next >

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

Print this page
rev 16828 : implements default methods in TreeMap v0
rev 16829 : TreeMap
rev 16830 : TreeMap

@@ -27,10 +27,11 @@
 
 import java.io.Serializable;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * A Red-Black tree based {@link NavigableMap} implementation.
  * The map is sorted according to the {@linkplain Comparable natural
  * ordering} of its keys, or by a {@link Comparator} provided at map

@@ -529,17 +530,35 @@
      * @throws NullPointerException if the specified key is null
      *         and this map uses natural ordering, or its comparator
      *         does not permit null keys
      */
     public V put(K key, V value) {
-        Entry<K,V> t = root;
-        if (t == null) {
-            compare(key, key); // type (and possibly null) check
+        return put(key, value, true);
+    }
 
+    private void addEntry(K key, V value, int cmp, Entry<K, V> parent) {
+        Entry<K,V> e = new Entry<>(key, value, parent);
+        if (cmp < 0)
+            parent.left = e;
+        else
+            parent.right = e;
+        fixAfterInsertion(e);
+        size++;
+        modCount++;
+    }
+
+    private void addEntryToEmptyMap(K key, V value) {
+        compare(key, key); // type (and possibly null) check
             root = new Entry<>(key, value, null);
             size = 1;
             modCount++;
+    }
+
+    private V put(K key, V value, boolean replaceOld) {
+        Entry<K,V> t = root;
+        if (t == null) {
+            addEntryToEmptyMap(key, value);
             return null;
         }
         int cmp;
         Entry<K,V> parent;
         // split comparator and comparable paths

@@ -550,12 +569,76 @@
                 cmp = cpr.compare(key, t.key);
                 if (cmp < 0)
                     t = t.left;
                 else if (cmp > 0)
                     t = t.right;
+                else {
+                    V oldValue = t.value;
+                    if(replaceOld) {
+                        t.value = value;
+                    }
+                    return oldValue;
+                }
+            } while (t != null);
+        }
+        else {
+            if (key == null)
+                throw new NullPointerException();
+            @SuppressWarnings("unchecked")
+            Comparable<? super K> k = (Comparable<? super K>) key;
+            do {
+                parent = t;
+                cmp = k.compareTo(t.key);
+                if (cmp < 0)
+                    t = t.left;
+                else if (cmp > 0)
+                    t = t.right;
+                else {
+                    V oldValue = t.value;
+                    if(replaceOld) {
+                        t.value = value;
+                    }
+                    return oldValue;
+                }
+            } while (t != null);
+        }
+        addEntry(key, value, cmp, parent);
+        return null;
+    }
+
+    @Override
+    public V putIfAbsent(K key, V value) {
+        return put(key, value, false);
+    }
+
+    @Override
+    public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction) {
+        Objects.requireNonNull(mappingFunction);
+        V newValue;
+        Entry<K,V> t = root;
+        if (t == null) {
+            if ((newValue = mappingFunction.apply(key)) != null) {
+                addEntryToEmptyMap(key, newValue);
+                return newValue;
+            } else {
+                return null;
+            }
+        }
+        int cmp;
+        Entry<K,V> parent;
+        // split comparator and comparable paths
+        Comparator<? super K> cpr = comparator;
+        if (cpr != null) {
+            do {
+                parent = t;
+                cmp = cpr.compare(key, t.key);
+                if (cmp < 0)
+                    t = t.left;
+                else if (cmp > 0)
+                    t = t.right;
                 else
-                    return t.setValue(value);
+                    return t.value;
             } while (t != null);
         }
         else {
             if (key == null)
                 throw new NullPointerException();

@@ -567,21 +650,93 @@
                 if (cmp < 0)
                     t = t.left;
                 else if (cmp > 0)
                     t = t.right;
                 else
-                    return t.setValue(value);
+                    return t.value;
             } while (t != null);
         }
-        Entry<K,V> e = new Entry<>(key, value, parent);
+        if ((newValue = mappingFunction.apply(key)) != null) {
+            addEntry(key, newValue, cmp, parent);
+            return newValue;
+        }
+        return null;
+    }
+
+    @Override
+    public V computeIfPresent(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+        Objects.requireNonNull(remappingFunction);
+        Entry<K,V> oldEntry = getEntry(key);
+        if (oldEntry != null && oldEntry.value != null) {
+            return remapValue(oldEntry, key, remappingFunction);
+        } else {
+            return null;
+        }
+    }
+
+    private V remapValue(Entry<K, V> t, K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+        V newValue = remappingFunction.apply(key, t.value);
+        if (newValue == null) {
+            deleteEntry(t);
+            return null;
+        } else {
+            // replace old mapping
+            t.value = newValue;
+            return newValue;
+        }
+    }
+
+    @Override
+    public V compute(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+        Objects.requireNonNull(remappingFunction);
+        V newValue;
+        Entry<K,V> t = root;
+        if (t == null) {
+            newValue = remappingFunction.apply(key, null);
+            if (newValue != null) {
+                addEntryToEmptyMap(key, newValue);
+                return newValue;
+            } else {
+                return null;
+            }
+        }
+        int cmp;
+        Entry<K,V> parent;
+        // split comparator and comparable paths
+        Comparator<? super K> cpr = comparator;
+        if (cpr != null) {
+            do {
+                parent = t;
+                cmp = cpr.compare(key, t.key);
         if (cmp < 0)
-            parent.left = e;
+                    t = t.left;
+                else if (cmp > 0)
+                    t = t.right;
         else
-            parent.right = e;
-        fixAfterInsertion(e);
-        size++;
-        modCount++;
+                    return remapValue(t, key, remappingFunction);
+            } while (t != null);
+        }
+        else {
+            if (key == null)
+                throw new NullPointerException();
+            @SuppressWarnings("unchecked")
+            Comparable<? super K> k = (Comparable<? super K>) key;
+            do {
+                parent = t;
+                cmp = k.compareTo(t.key);
+                if (cmp < 0)
+                    t = t.left;
+                else if (cmp > 0)
+                    t = t.right;
+                else
+                    return remapValue(t, key, remappingFunction);
+            } while (t != null);
+        }
+        if ((newValue = remappingFunction.apply(key, null)) != null) {
+            addEntry(key, newValue, cmp, parent);
+            return newValue;
+        }
         return null;
     }
 
     /**
      * Removes the mapping for this key from this TreeMap if present.
< prev index next >