Merge "Improve replaceAll."

apct-tests/perftests/core/src/android/util/ArrayMapPerfTest.java

@@ -27,6 +27,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 
 @RunWith(AndroidJUnit4.class)
 @LargeTest
@@ -69,4 +70,34 @@ public class ArrayMapPerfTest {
             }
         }
     }
+
+    @Test
+    public void testReplaceAll_Small() {
+        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
+        BiFunction<Integer, Integer, Integer> function = (k, v) -> 2 * v;
+        while (state.keepRunning()) {
+            for (int i = 0; i < NUM_ITERATIONS; ++i) {
+                ArrayMap<Integer, Integer> map = new ArrayMap<>();
+                for (int j = 0; j < SET_SIZE_SMALL; j++) {
+                    map.put(j, j);
+                }
+                map.replaceAll(function);
+            }
+        }
+    }
+
+    @Test
+    public void testReplaceAll_Large() {
+        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
+        BiFunction<Integer, Integer, Integer> function = (k, v) -> 2 * v;
+        while (state.keepRunning()) {
+            for (int i = 0; i < NUM_ITERATIONS; ++i) {
+                ArrayMap<Integer, Integer> map = new ArrayMap<>();
+                for (int j = 0; j < SET_SIZE_LARGE; j++) {
+                    map.put(j, j);
+                }
+                map.replaceAll(function);
+            }
+        }
+    }
 }

core/java/android/util/ArrayMap.java

@@ -28,6 +28,7 @@ import java.util.ConcurrentModificationException;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 
 /**
  * ArrayMap is a generic key->value mapping data structure that is
@@ -1023,6 +1024,36 @@ public final class ArrayMap<K, V> implements Map<K, V> {
         return MapCollections.removeAllHelper(this, collection);
     }
 
+    /**
+     * Replaces each entry's value with the result of invoking the given function on that entry
+     * until all entries have been processed or the function throws an exception. Exceptions thrown
+     * by the function are relayed to the caller. This implementation overrides
+     * the default implementation to avoid iterating using the {@link #entrySet()} and iterates in
+     * the key-value order consistent with {@link #keyAt(int)} and {@link #valueAt(int)}.
+     *
+     * @param function The function to apply to each entry
+     */
+    @Override
+    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
+        if (function == null) {
+            throw new NullPointerException("function must not be null");
+        }
+
+        final int size = mSize;
+        try {
+            for (int i = 0; i < size; ++i) {
+                final int valIndex = (i << 1) + 1;
+                //noinspection unchecked
+                mArray[valIndex] = function.apply((K) mArray[i << 1], (V) mArray[valIndex]);
+            }
+        } catch (ArrayIndexOutOfBoundsException e) {
+            throw new ConcurrentModificationException();
+        }
+        if (size != mSize) {
+            throw new ConcurrentModificationException();
+        }
+    }
+
     /**
      * Remove all keys in the array map that do <b>not</b> exist in the given collection.
      * @param collection The collection whose contents are to be used to determine which