Add LongArrayMultiStateCounter.combineValues

This method combines values into a smaller array
by using an index map.

Bug: 244623253
Test: atest com.android.internal.os.LongArrayMultiStateCounterTest
Change-Id: Ic35d30cc1f8765cd4e21366943da8291355a0047

3 files changed, 79 insertions, 0 deletions

diff --git a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
index 50dcca64a146..664aeee6e299 100644
--- a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
+++ b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
@@ -98,6 +98,18 @@ public final class LongArrayMultiStateCounter implements Parcelable {
             native_getValues(mNativeObject, array);
         }
 
+        /**
+         * Combines contained values into a smaller array by aggregating them
+         * according to an index map.
+         */
+        public boolean combineValues(long[] array, int[] indexMap) {
+            if (indexMap.length != mLength) {
+                throw new IllegalArgumentException(
+                        "Wrong index map size " + indexMap.length + ", expected " + mLength);
+            }
+            return native_combineValues(mNativeObject, array, indexMap);
+        }
+
         @Override
         public String toString() {
             final long[] array = new long[mLength];
@@ -116,6 +128,10 @@ public final class LongArrayMultiStateCounter implements Parcelable {
 
         @FastNative
         private native void native_getValues(long nativeObject, long[] array);
+
+        @FastNative
+        private native boolean native_combineValues(long nativeObject, long[] array,
+                int[] indexMap);
     }
 
     private static final NativeAllocationRegistry sRegistry =
diff --git a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
index 5b946d5c8d3a..a95b6e37f5de 100644
--- a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
+++ b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
@@ -244,6 +244,38 @@ static void native_getValues_LongArrayContainer(JNIEnv *env, jobject self, jlong
     std::copy(vector->data(), vector->data() + vector->size(), scopedArray.get());
 }
 
+static jboolean native_combineValues_LongArrayContainer(JNIEnv *env, jobject self, jlong nativePtr,
+                                                        jlongArray jarray, jintArray jindexMap) {
+    std::vector<uint64_t> *vector = reinterpret_cast<std::vector<uint64_t> *>(nativePtr);
+    ScopedLongArrayRW scopedArray(env, jarray);
+    ScopedIntArrayRO scopedIndexMap(env, jindexMap);
+
+    const uint64_t *data = vector->data();
+    uint64_t *array = reinterpret_cast<uint64_t *>(scopedArray.get());
+    const uint8_t size = scopedArray.size();
+
+    for (int i = 0; i < size; i++) {
+        array[i] = 0;
+    }
+
+    bool nonZero = false;
+    for (int i = 0; i < vector->size(); i++) {
+        jint index = scopedIndexMap[i];
+        if (index < 0 || index >= size) {
+            jniThrowExceptionFmt(env, "java/lang/IndexOutOfBoundsException",
+                                 "Index %d is out of bounds: [0, %d]", index, size - 1);
+            return false;
+        }
+
+        if (data[i] != 0L) {
+            array[index] += data[i];
+            nonZero = true;
+        }
+    }
+
+    return nonZero;
+}
+
 static const JNINativeMethod g_LongArrayContainer_methods[] = {
         // @CriticalNative
         {"native_init", "(I)J", (void *)native_init_LongArrayContainer},
@@ -253,6 +285,8 @@ static const JNINativeMethod g_LongArrayContainer_methods[] = {
         {"native_setValues", "(J[J)V", (void *)native_setValues_LongArrayContainer},
         // @FastNative
         {"native_getValues", "(J[J)V", (void *)native_getValues_LongArrayContainer},
+        // @FastNative
+        {"native_combineValues", "(J[J[I)Z", (void *)native_combineValues_LongArrayContainer},
 };
 
 int register_com_android_internal_os_LongArrayMultiStateCounter(JNIEnv *env) {
diff --git a/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java b/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java
index a22dd1c63ddf..516dee7dc9aa 100644
--- a/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/LongArrayMultiStateCounterTest.java
@@ -161,4 +161,33 @@ public class LongArrayMultiStateCounterTest {
         assertThrows(RuntimeException.class,
                 () -> LongArrayMultiStateCounter.CREATOR.createFromParcel(parcel));
     }
+
+    @Test
+    public void combineValues() {
+        long[] values = new long[] {0, 1, 2, 3, 42};
+        LongArrayMultiStateCounter.LongArrayContainer container =
+                new LongArrayMultiStateCounter.LongArrayContainer(values.length);
+        container.setValues(values);
+
+        long[] out = new long[3];
+        int[] indexes = {2, 1, 1, 0, 0};
+        boolean nonZero = container.combineValues(out, indexes);
+        assertThat(nonZero).isTrue();
+        assertThat(out).isEqualTo(new long[]{45, 3, 0});
+
+        // All zeros
+        container.setValues(new long[]{0, 0, 0, 0, 0});
+        nonZero = container.combineValues(out, indexes);
+        assertThat(nonZero).isFalse();
+        assertThat(out).isEqualTo(new long[]{0, 0, 0});
+
+        // Index out of range
+        IndexOutOfBoundsException e1 = assertThrows(
+                IndexOutOfBoundsException.class,
+                () -> container.combineValues(out, new int[]{0, 1, -1, 0, 0}));
+        assertThat(e1.getMessage()).isEqualTo("Index -1 is out of bounds: [0, 2]");
+        IndexOutOfBoundsException e2 = assertThrows(IndexOutOfBoundsException.class,
+                () -> container.combineValues(out, new int[]{0, 1, 4, 0, 0}));
+        assertThat(e2.getMessage()).isEqualTo("Index 4 is out of bounds: [0, 2]");
+    }
 }