Add getPackageStates to replace forAllPackageStates

Rather than pre-optimize the filtering path, the FilteredSnapshot
now just returns a Map<String, PackageState>.

Any optimization can come in the form of a custom Map implementation
which caches on iteration, but that can be saved for when there's a
use case for that.

This leaves the old API intact for now due to a bulid system quirk
where the dependent CL can't be merged if the API in the prebuilt
doesn't exist.

Test: atest PackageManagerLocalSnapshotTest

For now a simple Map makes the API easier to use.

Bug: 246609797

Change-Id: Ifc68d4753290e2922672859b1a8a15b1ec0795cf
diff --git a/services/api/current.txt b/services/api/current.txt
index 834ed2f..d05431d 100644
--- a/services/api/current.txt
+++ b/services/api/current.txt
@@ -59,6 +59,7 @@
     method public void close();
     method public void forAllPackageStates(@NonNull java.util.function.Consumer<com.android.server.pm.pkg.PackageState>);
     method @Nullable public com.android.server.pm.pkg.PackageState getPackageState(@NonNull String);
+    method @NonNull public java.util.Map<java.lang.String,com.android.server.pm.pkg.PackageState> getPackageStates();
   }
 
   public static interface PackageManagerLocal.UnfilteredSnapshot extends java.lang.AutoCloseable {
diff --git a/services/core/java/com/android/server/pm/PackageManagerLocal.java b/services/core/java/com/android/server/pm/PackageManagerLocal.java
index 21c2f2c..fad61b8 100644
--- a/services/core/java/com/android/server/pm/PackageManagerLocal.java
+++ b/services/core/java/com/android/server/pm/PackageManagerLocal.java
@@ -167,14 +167,16 @@
         PackageState getPackageState(@NonNull String packageName);
 
         /**
-         * Iterates on all states. This should only be used when either the target package name
-         * is not known or the large majority of the states are expected to be used.
-         *
+         * Returns a map of all {@link PackageState PackageStates} on the device.
+         * <p>
          * This will cause app visibility filtering to be invoked on each state on the device,
-         * which can be expensive.
+         * which can be expensive. Prefer {@link #getPackageState(String)} if possible.
          *
-         * @param consumer Block to accept each state as it becomes available post-filtering.
+         * @return Mapping of package name to {@link PackageState}.
          */
+        @NonNull
+        Map<String, PackageState> getPackageStates();
+
         void forAllPackageStates(@NonNull Consumer<PackageState> consumer);
 
         @Override
diff --git a/services/core/java/com/android/server/pm/local/PackageManagerLocalImpl.java b/services/core/java/com/android/server/pm/local/PackageManagerLocalImpl.java
index 4ff0d59..024b63e 100644
--- a/services/core/java/com/android/server/pm/local/PackageManagerLocalImpl.java
+++ b/services/core/java/com/android/server/pm/local/PackageManagerLocalImpl.java
@@ -22,6 +22,7 @@
 import android.annotation.UserIdInt;
 import android.os.Binder;
 import android.os.UserHandle;
+import android.util.ArrayMap;
 
 import com.android.server.pm.Computer;
 import com.android.server.pm.PackageManagerLocal;
@@ -30,7 +31,6 @@
 import com.android.server.pm.snapshot.PackageDataSnapshot;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -143,7 +143,7 @@
         private final int mUserId;
 
         @Nullable
-        private ArrayList<PackageState> mFilteredPackageStates;
+        private Map<String, PackageState> mFilteredPackageStates;
 
         @Nullable
         private final UnfilteredSnapshotImpl mParentSnapshot;
@@ -179,25 +179,36 @@
             return mSnapshot.getPackageStateFiltered(packageName, mCallingUid, mUserId);
         }
 
+        @NonNull
         @Override
-        public void forAllPackageStates(@NonNull Consumer<PackageState> consumer) {
+        public Map<String, PackageState> getPackageStates() {
             checkClosed();
 
             if (mFilteredPackageStates == null) {
                 var packageStates = mSnapshot.getPackageStates();
-                var filteredPackageStates = new ArrayList<PackageState>();
+                var filteredPackageStates = new ArrayMap<String, PackageState>();
                 for (int index = 0, size = packageStates.size(); index < size; index++) {
                     var packageState = packageStates.valueAt(index);
                     if (!mSnapshot.shouldFilterApplication(packageState, mCallingUid, mUserId)) {
-                        filteredPackageStates.add(packageState);
+                        filteredPackageStates.put(packageStates.keyAt(index), packageState);
                     }
                 }
-                mFilteredPackageStates = filteredPackageStates;
+                mFilteredPackageStates = Collections.unmodifiableMap(filteredPackageStates);
             }
 
-            for (int index = 0, size = mFilteredPackageStates.size(); index < size; index++) {
-                var packageState = mFilteredPackageStates.get(index);
-                consumer.accept(packageState);
+            return mFilteredPackageStates;
+        }
+
+        @Override
+        public void forAllPackageStates(@NonNull Consumer<PackageState> consumer) {
+            checkClosed();
+
+            var packageStates = mSnapshot.getPackageStates();
+            for (int index = 0, size = packageStates.size(); index < size; index++) {
+                var packageState = packageStates.valueAt(index);
+                if (!mSnapshot.shouldFilterApplication(packageState, mCallingUid, mUserId)) {
+                    consumer.accept(packageState);
+                }
             }
         }
     }
diff --git a/services/tests/PackageManagerServiceTests/unit/src/com/android/server/pm/test/PackageManagerLocalSnapshotTest.kt b/services/tests/PackageManagerServiceTests/unit/src/com/android/server/pm/test/PackageManagerLocalSnapshotTest.kt
index faa2352..5f26d6f 100644
--- a/services/tests/PackageManagerServiceTests/unit/src/com/android/server/pm/test/PackageManagerLocalSnapshotTest.kt
+++ b/services/tests/PackageManagerServiceTests/unit/src/com/android/server/pm/test/PackageManagerLocalSnapshotTest.kt
@@ -53,6 +53,11 @@
         snapshot.use {
             val packageStates = it.packageStates
 
+            // Check for unmodifiable
+            assertFailsWith(UnsupportedOperationException::class) {
+                it.packageStates.clear()
+            }
+
             // Check contents
             assertThat(packageStates).containsExactly(
                 packageStateAll.packageName, packageStateAll,
@@ -78,9 +83,14 @@
             assertThat(filteredOne.getPackageState(packageStateUser10.packageName)).isNull()
 
             filteredThree.use {
-                val statesList = mutableListOf<PackageState>()
-                assertThat(it.forAllPackageStates { statesList += it })
-                assertThat(statesList).containsExactly(packageStateAll, packageStateUser10)
+                // Check for unmodifiable
+                assertFailsWith(UnsupportedOperationException::class) {
+                    it.packageStates.clear()
+                }
+                assertThat(it.packageStates).containsExactly(
+                    packageStateAll.packageName, packageStateAll,
+                    packageStateUser10.packageName, packageStateUser10,
+                )
             }
 
             // Call after child close, parent open fails
@@ -96,7 +106,7 @@
 
         // Call after close fails
         assertClosedFailure { snapshot.packageStates }
-        assertClosedFailure { filteredOne.forAllPackageStates {} }
+        assertClosedFailure { filteredOne.packageStates }
         assertClosedFailure {
             filteredTwo.getPackageState(packageStateAll.packageName)
         }
@@ -116,9 +126,15 @@
                 .isEqualTo(packageStateUser0)
             assertThat(it.getPackageState(packageStateUser10.packageName)).isNull()
 
-            val statesList = mutableListOf<PackageState>()
-            assertThat(it.forAllPackageStates { statesList += it })
-            assertThat(statesList).containsExactly(packageStateAll, packageStateUser0)
+            // Check for unmodifiable
+            assertFailsWith(UnsupportedOperationException::class) {
+                it.packageStates.clear()
+            }
+
+            assertThat(it.packageStates).containsExactly(
+                packageStateAll.packageName, packageStateAll,
+                packageStateUser0.packageName, packageStateUser0,
+            )
         }
 
         // Call after close fails