Add useful utils to CollectionUtils

Test: new tests in this patch
Change-Id: I9cca999f9235d4d6508533230afcd8ee0c1c54c1
diff --git a/staticlibs/framework/com/android/net/module/util/CollectionUtils.java b/staticlibs/framework/com/android/net/module/util/CollectionUtils.java
index 7cac90d..f08880c 100644
--- a/staticlibs/framework/com/android/net/module/util/CollectionUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/CollectionUtils.java
@@ -18,12 +18,15 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.util.ArrayMap;
+import android.util.Pair;
 import android.util.SparseArray;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 /**
@@ -291,4 +294,100 @@
             @NonNull final Predicate<? super T> condition) {
         return -1 != indexOf(haystack, condition);
     }
+
+    /**
+     * Standard map function, but returns a new modifiable ArrayList
+     *
+     * This returns a new list that contains, for each element of the source collection, its
+     * image through the passed transform.
+     * Elements in the source can be null if the transform accepts null inputs.
+     * Elements in the output can be null if the transform ever returns null.
+     * This function never returns null. If the source collection is empty, it returns the
+     * empty list.
+     * Contract : this method calls the transform function exactly once for each element in the
+     * list, in iteration order.
+     *
+     * @param source the source collection
+     * @param transform the function to transform the elements
+     * @param <T> type of source elements
+     * @param <R> type of destination elements
+     * @return an unmodifiable list of transformed elements
+     */
+    @NonNull
+    public static <T, R> ArrayList<R> map(@NonNull final Collection<T> source,
+            @NonNull final Function<? super T, ? extends R> transform) {
+        final ArrayList<R> dest = new ArrayList<>(source.size());
+        for (final T e : source) {
+            dest.add(transform.apply(e));
+        }
+        return dest;
+    }
+
+    /**
+     * Standard zip function, but returns a new modifiable ArrayList
+     *
+     * This returns a list of pairs containing, at each position, a pair of the element from the
+     * first list at that index and the element from the second list at that index.
+     * Both lists must be the same size. They may contain null.
+     *
+     * The easiest way to visualize what's happening is to think of two lists being laid out next
+     * to each other and stitched together with a zipper.
+     *
+     * Contract : this method will read each element of each list exactly once, in some unspecified
+     * order. If it throws, it will not read any element.
+     *
+     * @param first the first list of elements
+     * @param second the second list of elements
+     * @param <T> the type of first elements
+     * @param <R> the type of second elements
+     * @return the zipped list
+     */
+    @NonNull
+    public static <T, R> ArrayList<Pair<T, R>> zip(@NonNull final List<T> first,
+            @NonNull final List<R> second) {
+        final int size = first.size();
+        if (size != second.size()) {
+            throw new IllegalArgumentException("zip : collections must be the same size");
+        }
+        final ArrayList<Pair<T, R>> dest = new ArrayList<>(size);
+        for (int i = 0; i < size; ++i) {
+            dest.add(new Pair<>(first.get(i), second.get(i)));
+        }
+        return dest;
+    }
+
+    /**
+     * Returns a new ArrayMap that associates each key with the value at the same index.
+     *
+     * Both lists must be the same size.
+     * Both keys and values may contain null.
+     * Keys may not contain the same value twice.
+     *
+     * Contract : this method will read each element of each list exactly once, but does not
+     * specify the order, except if it throws in which case the number of reads is undefined.
+     *
+     * @param keys The list of keys
+     * @param values The list of values
+     * @param <T> The type of keys
+     * @param <R> The type of values
+     * @return The associated map
+     */
+    @NonNull
+    public static <T, R> ArrayMap<T, R> assoc(
+            @NonNull final List<T> keys, @NonNull final List<R> values) {
+        final int size = keys.size();
+        if (size != values.size()) {
+            throw new IllegalArgumentException("assoc : collections must be the same size");
+        }
+        final ArrayMap<T, R> dest = new ArrayMap<>(size);
+        for (int i = 0; i < size; ++i) {
+            final T key = keys.get(i);
+            if (dest.containsKey(key)) {
+                throw new IllegalArgumentException(
+                        "assoc : keys may not contain the same value twice");
+            }
+            dest.put(key, values.get(i));
+        }
+        return dest;
+    }
 }
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/CollectionUtilsTest.kt b/staticlibs/tests/unit/src/com/android/net/module/util/CollectionUtilsTest.kt
index 0f00d0b..9fb025b 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/CollectionUtilsTest.kt
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/CollectionUtilsTest.kt
@@ -22,6 +22,7 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
 import kotlin.test.assertNull
 import kotlin.test.assertSame
@@ -133,4 +134,52 @@
         assertSame(CollectionUtils.findLast(listMulti) { it == "E" }, listMulti[7])
         assertNull(CollectionUtils.findLast(listMulti) { it == "F" })
     }
+
+    @Test
+    fun testMap() {
+        val listAE = listOf("A", "B", "C", "D", "E", null)
+        assertEquals(listAE.map { "-$it-" }, CollectionUtils.map(listAE) { "-$it-" })
+    }
+
+    @Test
+    fun testZip() {
+        val listAE = listOf("A", "B", "C", "D", "E")
+        val list15 = listOf(1, 2, 3, 4, 5)
+        // Normal #zip returns kotlin.Pair, not android.util.Pair
+        assertEquals(list15.zip(listAE).map { android.util.Pair(it.first, it.second) },
+                CollectionUtils.zip(list15, listAE))
+        val listNull = listOf("A", null, "B", "C", "D")
+        assertEquals(list15.zip(listNull).map { android.util.Pair(it.first, it.second) },
+                CollectionUtils.zip(list15, listNull))
+        assertEquals(emptyList<android.util.Pair<Int, Int>>(),
+                CollectionUtils.zip(emptyList<Int>(), emptyList<Int>()))
+        assertFailsWith<IllegalArgumentException> {
+            // Different size
+            CollectionUtils.zip(listOf(1, 2), list15)
+        }
+    }
+
+    @Test
+    fun testAssoc() {
+        val listADA = listOf("A", "B", "C", "D", "A")
+        val list15 = listOf(1, 2, 3, 4, 5)
+        assertEquals(list15.zip(listADA).toMap(), CollectionUtils.assoc(list15, listADA))
+
+        // Null key is fine
+        val assoc = CollectionUtils.assoc(listOf(1, 2, null), listOf("A", "B", "C"))
+        assertEquals("C", assoc[null])
+
+        assertFailsWith<IllegalArgumentException> {
+            // Same key multiple times
+            CollectionUtils.assoc(listOf("A", "B", "A"), listOf(1, 2, 3))
+        }
+        assertFailsWith<IllegalArgumentException> {
+            // Same key multiple times, but it's null
+            CollectionUtils.assoc(listOf(null, "B", null), listOf(1, 2, 3))
+        }
+        assertFailsWith<IllegalArgumentException> {
+            // Different size
+            CollectionUtils.assoc(listOf(1, 2), list15)
+        }
+    }
 }