Add replace to map-like collections

This was added to Map in API 24 which users of ArrayMap are able to use, but there's no reason ArrayMap, SimpleArrayMap, or the primitive-specialized collection users shouldn't have access on all API levels.

Bug: 116778381
Test: ./gradlew :collection:build
Change-Id: I304cce70a50786f52b843ecd0ca14d1008b7a004

7 files changed, 121 insertions, 0 deletions

diff --git a/collection/api/current.txt b/collection/api/current.txt
index 70490b8a4f2..5d2f7c4065e 100644
--- a/collection/api/current.txt
+++ b/collection/api/current.txt
@@ -93,6 +93,7 @@ package androidx.collection {
     method public E putIfAbsent(long, E);
     method public void remove(long);
     method public void removeAt(int);
+    method public E replace(long, E);
     method public void setValueAt(int, E);
     method public int size();
     method public E valueAt(int);
@@ -138,6 +139,7 @@ package androidx.collection {
     method public V putIfAbsent(K, V);
     method public V remove(java.lang.Object);
     method public V removeAt(int);
+    method public V replace(K, V);
     method public V setValueAt(int, V);
     method public int size();
     method public V valueAt(int);
@@ -164,6 +166,7 @@ package androidx.collection {
     method public void remove(int);
     method public void removeAt(int);
     method public void removeAtRange(int, int);
+    method public E replace(int, E);
     method public void setValueAt(int, E);
     method public int size();
     method public E valueAt(int);
diff --git a/collection/src/main/java/androidx/collection/LongSparseArray.java b/collection/src/main/java/androidx/collection/LongSparseArray.java
index 451b7f78eda..41400ed6bbc 100644
--- a/collection/src/main/java/androidx/collection/LongSparseArray.java
+++ b/collection/src/main/java/androidx/collection/LongSparseArray.java
@@ -152,6 +152,23 @@ public class LongSparseArray<E> implements Cloneable {
         }
     }
 
+    /**
+     * Replace the mapping for {@code key} only if it is already mapped to a value.
+     * @param key The key of the mapping to replace.
+     * @param value The value to store for the given key.
+     * @return Returns the previous mapped value or null.
+     */
+    @Nullable
+    public E replace(long key, E value) {
+        int index = indexOfKey(key);
+        if (index >= 0) {
+            E oldValue = (E) mValues[index];
+            mValues[index] = value;
+            return oldValue;
+        }
+        return null;
+    }
+
     private void gc() {
         // Log.e("SparseArray", "gc start with " + mSize);
 
diff --git a/collection/src/main/java/androidx/collection/SimpleArrayMap.java b/collection/src/main/java/androidx/collection/SimpleArrayMap.java
index 73e765a2e7c..816e6346656 100644
--- a/collection/src/main/java/androidx/collection/SimpleArrayMap.java
+++ b/collection/src/main/java/androidx/collection/SimpleArrayMap.java
@@ -609,6 +609,21 @@ public class SimpleArrayMap<K, V> {
     }
 
     /**
+     * Replace the mapping for {@code key} only if it is already mapped to a value.
+     * @param key The key of the mapping to replace.
+     * @param value The value to store for the given key.
+     * @return Returns the previous mapped value or null.
+     */
+    @Nullable
+    public V replace(K key, V value) {
+        int index = indexOfKey(key);
+        if (index >= 0) {
+            return setValueAt(index, value);
+        }
+        return null;
+    }
+
+    /**
      * Return the number of items in this array map.
      */
     public int size() {
diff --git a/collection/src/main/java/androidx/collection/SparseArrayCompat.java b/collection/src/main/java/androidx/collection/SparseArrayCompat.java
index d17eca4cd04..b956bb0e58c 100644
--- a/collection/src/main/java/androidx/collection/SparseArrayCompat.java
+++ b/collection/src/main/java/androidx/collection/SparseArrayCompat.java
@@ -169,6 +169,23 @@ public class SparseArrayCompat<E> implements Cloneable {
         }
     }
 
+    /**
+     * Replace the mapping for {@code key} only if it is already mapped to a value.
+     * @param key The key of the mapping to replace.
+     * @param value The value to store for the given key.
+     * @return Returns the previous mapped value or null.
+     */
+    @Nullable
+    public E replace(int key, E value) {
+        int index = indexOfKey(key);
+        if (index >= 0) {
+            E oldValue = (E) mValues[index];
+            mValues[index] = value;
+            return oldValue;
+        }
+        return null;
+    }
+
     private void gc() {
         // Log.e("SparseArray", "gc start with " + mSize);
 
diff --git a/collection/src/test/java/androidx/collection/LongSparseArrayTest.java b/collection/src/test/java/androidx/collection/LongSparseArrayTest.java
index c8098bc90a1..33fc679f741 100644
--- a/collection/src/test/java/androidx/collection/LongSparseArrayTest.java
+++ b/collection/src/test/java/androidx/collection/LongSparseArrayTest.java
@@ -83,6 +83,29 @@ public class LongSparseArrayTest {
     }
 
     @Test
+    public void replaceWhenAbsentDoesNotStore() {
+        LongSparseArray<String> map = new LongSparseArray<>();
+        assertNull(map.replace(1L, "1"));
+        assertFalse(map.containsKey(1L));
+    }
+
+    @Test
+    public void replaceStoresAndReturnsOldValue() {
+        LongSparseArray<String> map = new LongSparseArray<>();
+        map.put(1L, "1");
+        assertEquals("1", map.replace(1L, "2"));
+        assertEquals("2", map.get(1L));
+    }
+
+    @Test
+    public void replaceStoresAndReturnsNullWhenMappedToNull() {
+        LongSparseArray<String> map = new LongSparseArray<>();
+        map.put(1L, null);
+        assertNull(map.replace(1L, "1"));
+        assertEquals("1", map.get(1L));
+    }
+
+    @Test
     public void isEmpty() {
         LongSparseArray<String> LongSparseArray = new LongSparseArray<>();
         assertTrue(LongSparseArray.isEmpty()); // Newly created LongSparseArray should be empty
diff --git a/collection/src/test/java/androidx/collection/SimpleArrayMapTest.java b/collection/src/test/java/androidx/collection/SimpleArrayMapTest.java
index 7326832deeb..6bb86da7eea 100644
--- a/collection/src/test/java/androidx/collection/SimpleArrayMapTest.java
+++ b/collection/src/test/java/androidx/collection/SimpleArrayMapTest.java
@@ -86,6 +86,29 @@ public class SimpleArrayMapTest {
         assertNull(map.putIfAbsent("one", "2"));
     }
 
+    @Test
+    public void replaceWhenAbsentDoesNotStore() {
+        SimpleArrayMap<String, String> map = new SimpleArrayMap<>();
+        assertNull(map.replace("one", "1"));
+        assertFalse(map.containsKey("one"));
+    }
+
+    @Test
+    public void replaceStoresAndReturnsOldValue() {
+        SimpleArrayMap<String, String> map = new SimpleArrayMap<>();
+        map.put("one", "1");
+        assertEquals("1", map.replace("one", "2"));
+        assertEquals("2", map.get("one"));
+    }
+
+    @Test
+    public void replaceStoresAndReturnsNullWhenMappedToNull() {
+        SimpleArrayMap<String, String> map = new SimpleArrayMap<>();
+        map.put("one", null);
+        assertNull(map.replace("one", "1"));
+        assertEquals("1", map.get("one"));
+    }
+
     /**
      * Attempt to generate a ConcurrentModificationException in ArrayMap.
      */
diff --git a/collection/src/test/java/androidx/collection/SparseArrayCompatTest.java b/collection/src/test/java/androidx/collection/SparseArrayCompatTest.java
index 6d4ab3a6f6d..556694fe401 100644
--- a/collection/src/test/java/androidx/collection/SparseArrayCompatTest.java
+++ b/collection/src/test/java/androidx/collection/SparseArrayCompatTest.java
@@ -83,6 +83,29 @@ public class SparseArrayCompatTest {
     }
 
     @Test
+    public void replaceWhenAbsentDoesNotStore() {
+        SparseArrayCompat<String> map = new SparseArrayCompat<>();
+        assertNull(map.replace(1, "1"));
+        assertFalse(map.containsKey(1));
+    }
+
+    @Test
+    public void replaceStoresAndReturnsOldValue() {
+        SparseArrayCompat<String> map = new SparseArrayCompat<>();
+        map.put(1, "1");
+        assertEquals("1", map.replace(1, "2"));
+        assertEquals("2", map.get(1));
+    }
+
+    @Test
+    public void replaceStoresAndReturnsNullWhenMappedToNull() {
+        SparseArrayCompat<String> map = new SparseArrayCompat<>();
+        map.put(1, null);
+        assertNull(map.replace(1, "1"));
+        assertEquals("1", map.get(1));
+    }
+
+    @Test
     public void isEmpty() throws Exception {
         SparseArrayCompat<String> sparseArrayCompat = new SparseArrayCompat<>();
         assertTrue(sparseArrayCompat.isEmpty()); // Newly created SparseArrayCompat should be empty