2 files changed, 49 insertions, 57 deletions

diff --git a/java/src/com/android/intentresolver/inject/SystemServices.kt b/java/src/com/android/intentresolver/inject/SystemServices.kt
index 9e7c67b6..c09598e0 100644
--- a/java/src/com/android/intentresolver/inject/SystemServices.kt
+++ b/java/src/com/android/intentresolver/inject/SystemServices.kt
@@ -17,6 +17,7 @@ package com.android.intentresolver.inject
 
 import android.app.ActivityManager
 import android.app.admin.DevicePolicyManager
+import android.app.prediction.AppPredictionManager
 import android.content.ClipboardManager
 import android.content.ContentInterface
 import android.content.ContentResolver
@@ -26,7 +27,6 @@ import android.content.pm.ShortcutManager
 import android.os.UserManager
 import android.view.WindowManager
 import androidx.core.content.getSystemService
-import com.android.intentresolver.v2.data.repository.UserScopedContext
 import com.android.intentresolver.v2.data.repository.UserScopedService
 import com.android.intentresolver.v2.data.repository.UserScopedServiceImpl
 import dagger.Binds
@@ -91,10 +91,29 @@ class PackageManagerModule {
 
 @Module
 @InstallIn(SingletonComponent::class)
+class PredictionManagerModule {
+    @Provides
+    fun scopedPredictionManager(
+        @ApplicationContext ctx: Context,
+    ): UserScopedService<AppPredictionManager> {
+        return UserScopedServiceImpl(ctx, AppPredictionManager::class)
+    }
+}
+
+@Module
+@InstallIn(SingletonComponent::class)
 class ShortcutManagerModule {
     @Provides
-    fun shortcutManager(@ApplicationContext ctx: Context): ShortcutManager =
-        ctx.requireSystemService()
+    fun shortcutManager(@ApplicationContext ctx: Context): ShortcutManager {
+        return ctx.requireSystemService()
+    }
+
+    @Provides
+    fun scopedShortcutManager(
+        @ApplicationContext ctx: Context,
+    ): UserScopedService<ShortcutManager> {
+        return UserScopedServiceImpl(ctx, ShortcutManager::class)
+    }
 }
 
 @Module
@@ -104,7 +123,7 @@ class UserManagerModule {
     fun userManager(@ApplicationContext ctx: Context): UserManager = ctx.requireSystemService()
 
     @Provides
-    fun scopedUserManager(ctx: UserScopedContext): UserScopedService<UserManager> {
+    fun scopedUserManager(@ApplicationContext ctx: Context): UserScopedService<UserManager> {
         return UserScopedServiceImpl(ctx, UserManager::class)
     }
 }
diff --git a/java/src/com/android/intentresolver/v2/data/repository/UserScopedService.kt b/java/src/com/android/intentresolver/v2/data/repository/UserScopedService.kt
index 507979a0..65a48a55 100644
--- a/java/src/com/android/intentresolver/v2/data/repository/UserScopedService.kt
+++ b/java/src/com/android/intentresolver/v2/data/repository/UserScopedService.kt
@@ -18,77 +18,50 @@ package com.android.intentresolver.v2.data.repository
 
 import android.content.Context
 import android.os.UserHandle
-import android.util.LruCache
 import androidx.core.content.getSystemService
-import javax.inject.Inject
+import dagger.hilt.android.qualifiers.ApplicationContext
 import kotlin.reflect.KClass
 
 /**
- * Provides cached instances of a [system service][Context.getSystemService] created with
+ * Provides instances of a [system service][Context.getSystemService] created with
  * [the context of a specified user][Context.createContextAsUser].
  *
- * System services which have only `@UserHandleAware` APIs operate on the user id available from
+ * Some services which have only `@UserHandleAware` APIs operate on the user id available from
  * [Context.getUser], the context used to retrieve the service. This utility helps adapt a per-user
  * API model to work in multi-user manner.
  *
  * Example usage:
  * ```
- * val usageStats = userScopedService<UsageStatsManager>(context)
+ *     @Provides
+ *     fun scopedUserManager(@ApplicationContext ctx: Context): UserScopedService<UserManager> {
+ *         return UserScopedServiceImpl(ctx, UserManager::class)
+ *     }
  *
- * fun getStatsForUser(
- *     user: User,
- *     from: Long,
- *     to: Long
- * ): UsageStats {
- *     return usageStats.forUser(user)
- *        .queryUsageStats(INTERVAL_BEST, from, to)
- * }
+ *     class MyUserHelper @Inject constructor(
+ *         private val userMgr: UserScopedService<UserManager>,
+ *     ) {
+ *         fun isPrivateProfile(user: UserHandle): UserManager {
+ *             return userMgr.forUser(user).isPrivateProfile()
+ *         }
+ *     }
  * ```
  */
-interface UserScopedService<T> {
+fun interface UserScopedService<T> {
+    /** Create a service instance for the given user. */
     fun forUser(user: UserHandle): T
 }
 
-/**
- * Provides cached Context instances each distinct per-User.
- *
- * @see [UserScopedService]
- */
-class UserScopedContext @Inject constructor(private val applicationContext: Context) {
-    private val contextCacheSizeLimit = 8
-
-    private val instances =
-        object : LruCache<UserHandle, Context>(contextCacheSizeLimit) {
-            override fun create(key: UserHandle): Context {
-                return applicationContext.createContextAsUser(key, 0)
-            }
-        }
-
-    fun forUser(user: UserHandle): Context {
-        synchronized(this) {
-            return if (applicationContext.user == user) {
-                applicationContext
+class UserScopedServiceImpl<T : Any>(
+    @ApplicationContext private val context: Context,
+    private val serviceType: KClass<T>,
+) : UserScopedService<T> {
+    override fun forUser(user: UserHandle): T {
+        val context =
+            if (context.user == user) {
+                context
             } else {
-                return instances[user]
+                context.createContextAsUser(user, 0)
             }
-        }
-    }
-}
-
-/** Returns a cache of service instances, distinct by user */
-class UserScopedServiceImpl<T : Any>(contexts: UserScopedContext, serviceType: KClass<T>) :
-    UserScopedService<T> {
-    private val instances =
-        object : LruCache<UserHandle, T>(8) {
-            override fun create(key: UserHandle): T {
-                val context = contexts.forUser(key)
-                return requireNotNull(context.getSystemService(serviceType.java))
-            }
-        }
-
-    override fun forUser(user: UserHandle): T {
-        synchronized(this) {
-            return instances[user]
-        }
+        return requireNotNull(context.getSystemService(serviceType.java))
     }
 }