From 6a97c5d9f5dea1a7e7914df4b0730d25324b74ac Mon Sep 17 00:00:00 2001
From: Kweku Adams <kwekua@google.com>
Date: Fri, 10 Sep 2021 13:28:24 -0700
Subject: [PATCH] Improve replaceAll.

The default implementation of Map.replaceAll() iterates through the
entrySet. Switch to a simple for loop iterating through the array
since the iterator is an inefficient way to access the array contents.

Benchmark results:

Before:

android.util.ArrayMapPerfTest#testReplaceAll_Large:
	replaceAll_Large_min (ns): 888977
	replaceAll_Large_median (ns): 904642
	replaceAll_Large_mean (ns): 914596
	replaceAll_Large_standardDeviation: 23960
android.util.ArrayMapPerfTest#testReplaceAll_Small:
	replaceAll_Small_min (ns): 177767
	replaceAll_Small_median (ns): 179955
	replaceAll_Small_mean (ns): 180331
	replaceAll_Small_standardDeviation: 2865

After:

android.util.ArrayMapPerfTest#testReplaceAll_Large:
	replaceAll_Large_min (ns): 557385
	replaceAll_Large_median (ns): 586495
	replaceAll_Large_mean (ns): 583712
	replaceAll_Large_standardDeviation: 20356
android.util.ArrayMapPerfTest#testReplaceAll_Small:
	replaceAll_Small_min (ns): 108051
	replaceAll_Small_median (ns): 110802
	replaceAll_Small_mean (ns): 109987
	replaceAll_Small_standardDeviation: 1334

Bug: 194098491
Test: atest CorePerfTests:ArrayMapPerfTest (see results above)
Test: atest CtsUtilTestCases:ArrayMapTest
Test: atest FrameworksCoreTests:ArrayMapTest
Change-Id: If87ae5f3aa4b184b95c16aed5f2fb6f31d11668a
---
 .../src/android/util/ArrayMapPerfTest.java    | 31 +++++++++++++++++++
 core/java/android/util/ArrayMap.java          | 31 +++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/apct-tests/perftests/core/src/android/util/ArrayMapPerfTest.java b/apct-tests/perftests/core/src/android/util/ArrayMapPerfTest.java
index b93a6ea18d8db..cf1512381c545 100644
--- a/apct-tests/perftests/core/src/android/util/ArrayMapPerfTest.java
+++ b/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);
+            }
+        }
+    }
 }
diff --git a/core/java/android/util/ArrayMap.java b/core/java/android/util/ArrayMap.java
index f50b51bbdb356..5153aee4bcca1 100644
--- a/core/java/android/util/ArrayMap.java
+++ b/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