8017088: Map/HashMap.compute() incorrect with key mapping to null value

Reviewed-by: dl, dholmes, plevart

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

@@ -1749,25 +1749,25 @@ public class HashMap<K,V>
                 V oldValue = pEntry.value;
                 if (oldValue != null) {
                     V newValue = remappingFunction.apply(key, oldValue);
-                    if (newValue == null) { // remove mapping
-                        modCount++;
-                        size--;
-                        tb.deleteTreeNode(p);
-                        pEntry.recordRemoval(this);
-                        if (tb.root == null || tb.first == null) {
-                            // assert tb.root == null && tb.first == null :
-                            //     "TreeBin.first and root should both be null";
-                            // TreeBin is now empty, we should blank this bin
-                            table[i] = null;
-                        }
-                    } else {
-                        pEntry.value = newValue;
-                        pEntry.recordAccess(this);
+                if (newValue == null) { // remove mapping
+                    modCount++;
+                    size--;
+                    tb.deleteTreeNode(p);
+                    pEntry.recordRemoval(this);
+                    if (tb.root == null || tb.first == null) {
+                        // assert tb.root == null && tb.first == null :
+                        //     "TreeBin.first and root should both be null";
+                        // TreeBin is now empty, we should blank this bin
+                        table[i] = null;
                     }
-                    return newValue;
+                } else {
+                    pEntry.value = newValue;
+                    pEntry.recordAccess(this);
                 }
+                return newValue;
             }
         }
+        }
         return null;
     }
 
@@ -1779,7 +1779,7 @@ public class HashMap<K,V>
         if (key == null) {
             V oldValue = nullKeyEntry == null ? null : nullKeyEntry.value;
             V newValue = remappingFunction.apply(key, oldValue);
-            if (newValue != oldValue) {
+            if (newValue != oldValue || (oldValue == null && nullKeyEntry != null)) {
                 if (newValue == null) {
                     removeNullKey();
                 } else {
@@ -1803,7 +1803,7 @@ public class HashMap<K,V>
                 if (e.hash == hash && Objects.equals(e.key, key)) {
                     V oldValue = e.value;
                     V newValue = remappingFunction.apply(key, oldValue);
-                    if (newValue != oldValue) {
+                    if (newValue != oldValue || oldValue == null) {
                         if (newValue == null) {
                             modCount++;
                             size--;
@@ -1829,7 +1829,7 @@ public class HashMap<K,V>
             TreeNode p = tb.getTreeNode(hash, key);
             V oldValue = p == null ? null : (V)p.entry.value;
             V newValue = remappingFunction.apply(key, oldValue);
-            if (newValue != oldValue) {
+            if (newValue != oldValue || (oldValue == null && p != null)) {
                 if (newValue == null) {
                     Entry<K,V> pEntry = (Entry<K,V>)p.entry;
                     modCount++;

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

@@ -640,7 +640,7 @@ public interface Map<K,V> {
      * @param key key with which the specified value is to be associated
      * @param value value to be associated with the specified key
      * @return the previous value associated with the specified key, or
-     *         {@code 1} if there was no mapping for the key.
+     *         {@code null} if there was no mapping for the key.
      *         (A {@code null} return can also indicate that the map
      *         previously associated {@code null} with the key,
      *         if the implementation supports null values.)
@@ -994,20 +994,40 @@ public interface Map<K,V> {
         V oldValue = get(key);
         for (;;) {
             V newValue = remappingFunction.apply(key, oldValue);
-            if (oldValue != null) {
-                if (newValue != null) {
-                    if (replace(key, oldValue, newValue))
-                        return newValue;
-                } else if (remove(key, oldValue)) {
+            if (newValue == null) {
+                // delete mapping
+                if(oldValue != null || containsKey(key)) {
+                    // something to remove
+                    if (remove(key, oldValue)) {
+                        // removed the old value as expected
+                        return null;
+                    }
+
+                    // some other value replaced old value. try again.
+                    oldValue = get(key);
+                } else {
+                    // nothing to do. Leave things as they were.
                     return null;
                 }
-                oldValue = get(key);
             } else {
-                if (newValue != null) {
-                    if ((oldValue = putIfAbsent(key, newValue)) == null)
+                // add or replace old mapping
+                if (oldValue != null) {
+                    // replace
+                    if (replace(key, oldValue, newValue)) {
+                        // replaced as expected.
                         return newValue;
+                    }
+
+                    // some other value replaced old value. try again.
+                    oldValue = get(key);
                 } else {
-                    return null;
+                    // add (replace if oldValue was null)
+                    if ((oldValue = putIfAbsent(key, newValue)) == null) {
+                        // replaced
+                        return newValue;
+                    }
+
+                    // some other value replaced old value. try again.
                 }
             }
         }

test/java/util/Map/Defaults.java

@@ -271,14 +271,21 @@ public class Defaults {
 
     @Test(dataProvider = "Map<IntegerEnum,String> rw=true keys=withNull values=withNull")
     public void testComputeIfPresentNulls(String description, Map<IntegerEnum, String> map) {
-        assertTrue(map.containsKey(null));
-        assertNull(map.get(null));
+        assertTrue(map.containsKey(null), description + ": null key absent");
+        assertNull(map.get(null), description + ": value not null");
         assertSame(map.computeIfPresent(null, (k, v) -> {
-            fail();
+            fail(description + ": null value is not deemed present");
             return EXTRA_VALUE;
         }), null, description);
         assertTrue(map.containsKey(null));
-        assertSame(map.get(null), null, description);
+        assertNull(map.get(null), description);
+        assertNull(map.remove(EXTRA_KEY), description + ": unexpected mapping");
+        assertNull(map.put(EXTRA_KEY, null), description + ": unexpected value");
+        assertSame(map.computeIfPresent(EXTRA_KEY, (k, v) -> {
+            fail(description + ": null value is not deemed present");
+            return EXTRA_VALUE;
+        }), null, description);
+        assertNull(map.get(EXTRA_KEY), description + ": null mapping gone");
     }
 
     @Test(dataProvider = "Map<IntegerEnum,String> rw=true keys=all values=all")
@@ -306,6 +313,12 @@ public class Defaults {
     public void testComputeNulls(String description, Map<IntegerEnum, String> map) {
         assertTrue(map.containsKey(null), "null key absent");
         assertNull(map.get(null), "value not null");
+        assertSame(map.compute(null, (k, v) -> {
+            assertNull(k);
+            assertNull(v);
+            return null;
+        }), null, description);
+        assertFalse(map.containsKey(null), description + ": null key present.");
         assertSame(map.compute(null, (k, v) -> {
             assertSame(k, null);
             assertNull(v);
@@ -313,13 +326,47 @@ public class Defaults {
         }), EXTRA_VALUE, description);
         assertTrue(map.containsKey(null));
         assertSame(map.get(null), EXTRA_VALUE, description);
-        assertSame(map.remove(null), EXTRA_VALUE, "removed value not expected");
-        assertFalse(map.containsKey(null), "null key present");
+        assertSame(map.remove(null), EXTRA_VALUE, description + ": removed value not expected");
+        // no mapping before and after
+        assertFalse(map.containsKey(null), description + ": null key present");
         assertSame(map.compute(null, (k, v) -> {
-            assertSame(k, null);
+            assertNull(k);
+            assertNull(v);
+            return null;
+        }), null, description + ": expected null result" );
+        assertFalse(map.containsKey(null), description + ": null key present");
+        // compute with map not containing value
+        assertNull(map.remove(EXTRA_KEY),  description + ": unexpected mapping");
+        assertFalse(map.containsKey(EXTRA_KEY),  description + ": key present");
+        assertSame(map.compute(EXTRA_KEY, (k, v) -> {
+            assertSame(k, EXTRA_KEY);
+            assertNull(v);
+            return null;
+        }), null, description);
+        assertFalse(map.containsKey(EXTRA_KEY),  description + ": null key present");
+        // ensure removal.
+        assertNull(map.put(EXTRA_KEY, EXTRA_VALUE));
+        assertSame(map.compute(EXTRA_KEY, (k, v) -> {
+            assertSame(k, EXTRA_KEY);
+            assertSame(v, EXTRA_VALUE);
+            return null;
+        }), null, description + ": null resulted expected");
+        assertFalse(map.containsKey(EXTRA_KEY),  description + ": null key present");
+       // compute with map containing null value
+        assertNull(map.put(EXTRA_KEY, null),  description + ": unexpected value");
+        assertSame(map.compute(EXTRA_KEY, (k, v) -> {
+            assertSame(k, EXTRA_KEY);
             assertNull(v);
             return null;
         }), null, description);
+        assertFalse(map.containsKey(EXTRA_KEY),  description + ": null key present");
+        assertNull(map.put(EXTRA_KEY, null),  description + ": unexpected value");
+        assertSame(map.compute(EXTRA_KEY, (k, v) -> {
+            assertSame(k, EXTRA_KEY);
+            assertNull(v);
+            return EXTRA_VALUE;
+        }), EXTRA_VALUE, description);
+        assertTrue(map.containsKey(EXTRA_KEY), "null key present");
     }
 
     @Test(dataProvider = "Map<IntegerEnum,String> rw=true keys=all values=all")