Extract shortcuts loading logic from ChooserActivity

Extract shortcut loading logic from ChooserActivity into a new class
mostly as-is. Major changes:
- run the logic on a background executor and deliver the result on the
  main thread;
- replace dependencies from ChooserListAdapter with the data it
  provided.
A number of tests thap previusly used
ChooserListAdapter#addServiceResults to provide shortcut results into
the view are updated and re-enabled.

Re-introduction of ag/20236439

Fix: 259462188
Test: manual tests
Test: atest IntentResolverUnitTests
Change-Id: I2555c3e486b9a443c5101bbda33b5b214c959b0f

3 files changed, 499 insertions, 267 deletions

diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index 558dfcf7..d5a0c32c 100644
--- a/java/src/com/android/intentresolver/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/ChooserActivity.java
@@ -45,12 +45,10 @@ import android.content.IntentSender;
 import android.content.IntentSender.SendIntentException;
 import android.content.SharedPreferences;
 import android.content.pm.ActivityInfo;
-import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.ResolveInfo;
 import android.content.pm.ShortcutInfo;
-import android.content.pm.ShortcutManager;
 import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.database.Cursor;
@@ -60,7 +58,6 @@ import android.graphics.Insets;
 import android.graphics.drawable.Drawable;
 import android.metrics.LogMaker;
 import android.net.Uri;
-import android.os.AsyncTask;
 import android.os.Bundle;
 import android.os.Environment;
 import android.os.Handler;
@@ -110,6 +107,7 @@ import com.android.intentresolver.model.AbstractResolverComparator;
 import com.android.intentresolver.model.AppPredictionServiceResolverComparator;
 import com.android.intentresolver.model.ResolverRankerServiceResolverComparator;
 import com.android.intentresolver.shortcuts.AppPredictorFactory;
+import com.android.intentresolver.shortcuts.ShortcutLoader;
 import com.android.intentresolver.widget.ResolverDrawerLayout;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.config.sysui.SystemUiDeviceConfigFlags;
@@ -136,6 +134,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 /**
@@ -187,8 +186,8 @@ public class ChooserActivity extends ResolverActivity implements
     // `ShortcutSelectionLogic` which packs the appropriate elements into the final `TargetInfo`.
     // That flow should be refactored so that `ChooserActivity` isn't responsible for holding their
     // intermediate data, and then these members can be removed.
-    private Map<ChooserTarget, AppTarget> mDirectShareAppTargetCache;
-    private Map<ChooserTarget, ShortcutInfo> mDirectShareShortcutInfoCache;
+    private final Map<ChooserTarget, AppTarget> mDirectShareAppTargetCache = new HashMap<>();
+    private final Map<ChooserTarget, ShortcutInfo> mDirectShareShortcutInfoCache = new HashMap<>();
 
     public static final int TARGET_TYPE_DEFAULT = 0;
     public static final int TARGET_TYPE_CHOOSER_TARGET = 1;
@@ -279,8 +278,6 @@ public class ChooserActivity extends ResolverActivity implements
 
     private View mContentView = null;
 
-    private final ShortcutToChooserTargetConverter mShortcutToChooserTargetConverter =
-            new ShortcutToChooserTargetConverter();
     private final SparseArray<ProfileRecord> mProfileRecords = new SparseArray<>();
 
     private void setupPreDrawForSharedElementTransition(View v) {
@@ -432,11 +429,13 @@ public class ChooserActivity extends ResolverActivity implements
         mShouldDisplayLandscape =
                 shouldDisplayLandscape(getResources().getConfiguration().orientation);
         setRetainInOnStop(intent.getBooleanExtra(EXTRA_PRIVATE_RETAIN_IN_ON_STOP, false));
+        IntentFilter targetIntentFilter = getTargetIntentFilter(target);
         createProfileRecords(
                 new AppPredictorFactory(
-                        this,
+                        getApplicationContext(),
                         target.getStringExtra(Intent.EXTRA_TEXT),
-                        getTargetIntentFilter(target)));
+                        targetIntentFilter),
+                targetIntentFilter);
 
         mPreviewCoordinator = new ChooserContentPreviewCoordinator(
                 mBackgroundThreadPoolExecutor,
@@ -497,7 +496,6 @@ public class ChooserActivity extends ResolverActivity implements
                         getTargetIntent(), getContentResolver(), this::isImageType),
                 target.getAction()
         );
-        mDirectShareShortcutInfoCache = new HashMap<>();
 
         setEnterSharedElementCallback(new SharedElementCallback() {
             @Override
@@ -518,20 +516,31 @@ public class ChooserActivity extends ResolverActivity implements
         return R.style.Theme_DeviceDefault_Chooser;
     }
 
-    private void createProfileRecords(AppPredictorFactory factory) {
+    private void createProfileRecords(
+            AppPredictorFactory factory, IntentFilter targetIntentFilter) {
         UserHandle mainUserHandle = getPersonalProfileUserHandle();
-        createProfileRecord(mainUserHandle, factory);
+        createProfileRecord(mainUserHandle, targetIntentFilter, factory);
 
         UserHandle workUserHandle = getWorkProfileUserHandle();
         if (workUserHandle != null) {
-            createProfileRecord(workUserHandle, factory);
+            createProfileRecord(workUserHandle, targetIntentFilter, factory);
         }
     }
 
-    private void createProfileRecord(UserHandle userHandle, AppPredictorFactory factory) {
+    private void createProfileRecord(
+            UserHandle userHandle, IntentFilter targetIntentFilter, AppPredictorFactory factory) {
         AppPredictor appPredictor = factory.create(userHandle);
+        ShortcutLoader shortcutLoader = ActivityManager.isLowRamDeviceStatic()
+                    ? null
+                    : createShortcutLoader(
+                            getApplicationContext(),
+                            appPredictor,
+                            userHandle,
+                            targetIntentFilter,
+                            shortcutsResult -> onShortcutsLoaded(userHandle, shortcutsResult));
         mProfileRecords.put(
-                userHandle.getIdentifier(), new ProfileRecord(appPredictor));
+                userHandle.getIdentifier(),
+                new ProfileRecord(appPredictor, shortcutLoader));
     }
 
     @Nullable
@@ -539,50 +548,19 @@ public class ChooserActivity extends ResolverActivity implements
         return mProfileRecords.get(userHandle.getIdentifier(), null);
     }
 
-    private void setupAppPredictorForUser(
-            UserHandle userHandle, AppPredictor.Callback appPredictorCallback) {
-        AppPredictor appPredictor = getAppPredictor(userHandle);
-        if (appPredictor == null) {
-            return;
-        }
-        mDirectShareAppTargetCache = new HashMap<>();
-        appPredictor.registerPredictionUpdates(this.getMainExecutor(), appPredictorCallback);
-    }
-
-    private AppPredictor.Callback createAppPredictorCallback(
-            ChooserListAdapter chooserListAdapter) {
-        return resultList -> {
-            if (isFinishing() || isDestroyed()) {
-                return;
-            }
-            if (chooserListAdapter.getCount() == 0) {
-                return;
-            }
-            if (resultList.isEmpty()
-                    && shouldQueryShortcutManager(chooserListAdapter.getUserHandle())) {
-                // APS may be disabled, so try querying targets ourselves.
-                queryDirectShareTargets(chooserListAdapter, true);
-                return;
-            }
-            final List<ShortcutManager.ShareShortcutInfo> shareShortcutInfos =
-                    new ArrayList<>();
-
-            List<AppTarget> shortcutResults = new ArrayList<>();
-            for (AppTarget appTarget : resultList) {
-                if (appTarget.getShortcutInfo() == null) {
-                    continue;
-                }
-                shortcutResults.add(appTarget);
-            }
-            resultList = shortcutResults;
-            for (AppTarget appTarget : resultList) {
-                shareShortcutInfos.add(new ShortcutManager.ShareShortcutInfo(
-                        appTarget.getShortcutInfo(),
-                        new ComponentName(
-                                appTarget.getPackageName(), appTarget.getClassName())));
-            }
-            sendShareShortcutInfoList(shareShortcutInfos, chooserListAdapter, resultList);
-        };
+    @VisibleForTesting
+    protected ShortcutLoader createShortcutLoader(
+            Context context,
+            AppPredictor appPredictor,
+            UserHandle userHandle,
+            IntentFilter targetIntentFilter,
+            Consumer<ShortcutLoader.Result> callback) {
+        return new ShortcutLoader(
+                context,
+                appPredictor,
+                userHandle,
+                targetIntentFilter,
+                callback);
     }
 
     static SharedPreferences getPinnedSharedPrefs(Context context) {
@@ -1482,147 +1460,6 @@ public class ChooserActivity extends ResolverActivity implements
         }
     }
 
-    @VisibleForTesting
-    protected void queryDirectShareTargets(
-                ChooserListAdapter adapter, boolean skipAppPredictionService) {
-        ProfileRecord record = getProfileRecord(adapter.getUserHandle());
-        if (record == null) {
-            return;
-        }
-
-        record.loadingStartTime = SystemClock.elapsedRealtime();
-
-        UserHandle userHandle = adapter.getUserHandle();
-        if (!skipAppPredictionService) {
-            if (record.appPredictor != null) {
-                record.appPredictor.requestPredictionUpdate();
-                return;
-            }
-        }
-        // Default to just querying ShortcutManager if AppPredictor not present.
-        final IntentFilter filter = getTargetIntentFilter();
-        if (filter == null) {
-            return;
-        }
-
-        AsyncTask.execute(() -> {
-            Context selectedProfileContext = createContextAsUser(userHandle, 0 /* flags */);
-            ShortcutManager sm = (ShortcutManager) selectedProfileContext
-                    .getSystemService(Context.SHORTCUT_SERVICE);
-            List<ShortcutManager.ShareShortcutInfo> resultList = sm.getShareTargets(filter);
-            sendShareShortcutInfoList(resultList, adapter, null);
-        });
-    }
-
-    /**
-     * Returns {@code false} if {@code userHandle} is the work profile and it's either
-     * in quiet mode or not running.
-     */
-    private boolean shouldQueryShortcutManager(UserHandle userHandle) {
-        if (!shouldShowTabs()) {
-            return true;
-        }
-        if (!getWorkProfileUserHandle().equals(userHandle)) {
-            return true;
-        }
-        if (!isUserRunning(userHandle)) {
-            return false;
-        }
-        if (!isUserUnlocked(userHandle)) {
-            return false;
-        }
-        if (isQuietModeEnabled(userHandle)) {
-            return false;
-        }
-        return true;
-    }
-
-    private void sendShareShortcutInfoList(
-                List<ShortcutManager.ShareShortcutInfo> resultList,
-                ChooserListAdapter chooserListAdapter,
-                @Nullable List<AppTarget> appTargets) {
-        if (appTargets != null && appTargets.size() != resultList.size()) {
-            throw new RuntimeException("resultList and appTargets must have the same size."
-                    + " resultList.size()=" + resultList.size()
-                    + " appTargets.size()=" + appTargets.size());
-        }
-        UserHandle userHandle = chooserListAdapter.getUserHandle();
-        Context selectedProfileContext = createContextAsUser(userHandle, 0 /* flags */);
-        for (int i = resultList.size() - 1; i >= 0; i--) {
-            final String packageName = resultList.get(i).getTargetComponent().getPackageName();
-            if (!isPackageEnabled(selectedProfileContext, packageName)) {
-                resultList.remove(i);
-                if (appTargets != null) {
-                    appTargets.remove(i);
-                }
-            }
-        }
-
-        // If |appTargets| is not null, results are from AppPredictionService and already sorted.
-        final int shortcutType = (appTargets == null ? TARGET_TYPE_SHORTCUTS_FROM_SHORTCUT_MANAGER :
-                TARGET_TYPE_SHORTCUTS_FROM_PREDICTION_SERVICE);
-
-        // Match ShareShortcutInfos with DisplayResolveInfos to be able to use the old code path
-        // for direct share targets. After ShareSheet is refactored we should use the
-        // ShareShortcutInfos directly.
-        List<ServiceResultInfo> resultRecords = new ArrayList<>();
-        for (DisplayResolveInfo displayResolveInfo : chooserListAdapter.getDisplayResolveInfos()) {
-            List<ShortcutManager.ShareShortcutInfo> matchingShortcuts =
-                    filterShortcutsByTargetComponentName(
-                            resultList, displayResolveInfo.getResolvedComponentName());
-            if (matchingShortcuts.isEmpty()) {
-                continue;
-            }
-
-            List<ChooserTarget> chooserTargets = mShortcutToChooserTargetConverter
-                    .convertToChooserTarget(
-                        matchingShortcuts,
-                        resultList,
-                        appTargets,
-                        mDirectShareAppTargetCache,
-                        mDirectShareShortcutInfoCache);
-
-            ServiceResultInfo resultRecord = new ServiceResultInfo(
-                    displayResolveInfo, chooserTargets);
-            resultRecords.add(resultRecord);
-        }
-
-        runOnUiThread(() -> {
-            if (!isDestroyed()) {
-                onShortcutsLoaded(chooserListAdapter, shortcutType, resultRecords);
-            }
-        });
-    }
-
-    private List<ShortcutManager.ShareShortcutInfo> filterShortcutsByTargetComponentName(
-            List<ShortcutManager.ShareShortcutInfo> allShortcuts, ComponentName requiredTarget) {
-        List<ShortcutManager.ShareShortcutInfo> matchingShortcuts = new ArrayList<>();
-        for (ShortcutManager.ShareShortcutInfo shortcut : allShortcuts) {
-            if (requiredTarget.equals(shortcut.getTargetComponent())) {
-                matchingShortcuts.add(shortcut);
-            }
-        }
-        return matchingShortcuts;
-    }
-
-    private boolean isPackageEnabled(Context context, String packageName) {
-        if (TextUtils.isEmpty(packageName)) {
-            return false;
-        }
-        ApplicationInfo appInfo;
-        try {
-            appInfo = context.getPackageManager().getApplicationInfo(packageName, 0);
-        } catch (NameNotFoundException e) {
-            return false;
-        }
-
-        if (appInfo != null && appInfo.enabled
-                && (appInfo.flags & ApplicationInfo.FLAG_SUSPENDED) == 0) {
-            return true;
-        }
-        return false;
-    }
-
     private void logDirectShareTargetReceived(int logCategory, UserHandle forUser) {
         ProfileRecord profileRecord = getProfileRecord(forUser);
         if (profileRecord == null) {
@@ -1656,7 +1493,7 @@ public class ChooserActivity extends ResolverActivity implements
                     Log.d(TAG, "Action to be updated is " + targetIntent.getAction());
                 }
             } else if (DEBUG) {
-                Log.d(TAG, "Can not log Chooser Counts of null ResovleInfo");
+                Log.d(TAG, "Can not log Chooser Counts of null ResolveInfo");
             }
         }
         mIsSuccessfullySelected = true;
@@ -1701,9 +1538,6 @@ public class ChooserActivity extends ResolverActivity implements
         if (directShareAppPredictor == null) {
             return;
         }
-        if (!targetInfo.isChooserTargetInfo()) {
-            return;
-        }
         AppTarget appTarget = targetInfo.getDirectShareAppTarget();
         if (appTarget != null) {
             // This is a direct share click that was provided by the APS
@@ -1820,11 +1654,6 @@ public class ChooserActivity extends ResolverActivity implements
         ChooserListAdapter chooserListAdapter = createChooserListAdapter(context, payloadIntents,
                 initialIntents, rList, filterLastUsed,
                 createListController(userHandle));
-        if (!ActivityManager.isLowRamDeviceStatic()) {
-            AppPredictor.Callback appPredictorCallback =
-                    createAppPredictorCallback(chooserListAdapter);
-            setupAppPredictorForUser(userHandle, appPredictorCallback);
-        }
         return new ChooserGridAdapter(chooserListAdapter);
     }
 
@@ -2111,42 +1940,41 @@ public class ChooserActivity extends ResolverActivity implements
     }
 
     private void maybeQueryAdditionalPostProcessingTargets(ChooserListAdapter chooserListAdapter) {
-        // don't support direct share on low ram devices
-        if (ActivityManager.isLowRamDeviceStatic()) {
+        UserHandle userHandle = chooserListAdapter.getUserHandle();
+        ProfileRecord record = getProfileRecord(userHandle);
+        if (record == null) {
             return;
         }
-
-        // no need to query direct share for work profile when its locked or disabled
-        if (!shouldQueryShortcutManager(chooserListAdapter.getUserHandle())) {
+        if (record.shortcutLoader == null) {
             return;
         }
-
-        if (DEBUG) {
-            Log.d(TAG, "querying direct share targets from ShortcutManager");
-        }
-
-        queryDirectShareTargets(chooserListAdapter, false);
+        record.loadingStartTime = SystemClock.elapsedRealtime();
+        record.shortcutLoader.queryShortcuts(chooserListAdapter.getDisplayResolveInfos());
     }
 
-    @VisibleForTesting
     @MainThread
-    protected void onShortcutsLoaded(
-            ChooserListAdapter adapter, int targetType, List<ServiceResultInfo> resultInfos) {
-        UserHandle userHandle = adapter.getUserHandle();
+    private void onShortcutsLoaded(
+            UserHandle userHandle, ShortcutLoader.Result shortcutsResult) {
         if (DEBUG) {
             Log.d(TAG, "onShortcutsLoaded for user: " + userHandle);
         }
-        for (ServiceResultInfo resultInfo : resultInfos) {
-            if (resultInfo.resultTargets != null) {
+        mDirectShareShortcutInfoCache.putAll(shortcutsResult.directShareShortcutInfoCache);
+        mDirectShareAppTargetCache.putAll(shortcutsResult.directShareAppTargetCache);
+        ChooserListAdapter adapter =
+                mChooserMultiProfilePagerAdapter.getListAdapterForUserHandle(userHandle);
+        if (adapter != null) {
+            for (ShortcutLoader.ShortcutResultInfo resultInfo : shortcutsResult.shortcutsByApp) {
                 adapter.addServiceResults(
-                        resultInfo.originalTarget,
-                        resultInfo.resultTargets,
-                        targetType,
-                        emptyIfNull(mDirectShareShortcutInfoCache),
-                        emptyIfNull(mDirectShareAppTargetCache));
+                        resultInfo.appTarget,
+                        resultInfo.shortcuts,
+                        shortcutsResult.isFromAppPredictor
+                                ? TARGET_TYPE_SHORTCUTS_FROM_PREDICTION_SERVICE
+                                : TARGET_TYPE_SHORTCUTS_FROM_SHORTCUT_MANAGER,
+                        mDirectShareShortcutInfoCache,
+                        mDirectShareAppTargetCache);
             }
+            adapter.completeServiceTargetLoading();
         }
-        adapter.completeServiceTargetLoading();
 
         logDirectShareTargetReceived(
                 MetricsEvent.ACTION_DIRECT_SHARE_TARGETS_LOADED_SHORTCUT_MANAGER,
@@ -2156,24 +1984,6 @@ public class ChooserActivity extends ResolverActivity implements
         getChooserActivityLogger().logSharesheetDirectLoadComplete();
     }
 
-    @VisibleForTesting
-    protected boolean isUserRunning(UserHandle userHandle) {
-        UserManager userManager = getSystemService(UserManager.class);
-        return userManager.isUserRunning(userHandle);
-    }
-
-    @VisibleForTesting
-    protected boolean isUserUnlocked(UserHandle userHandle) {
-        UserManager userManager = getSystemService(UserManager.class);
-        return userManager.isUserUnlocked(userHandle);
-    }
-
-    @VisibleForTesting
-    protected boolean isQuietModeEnabled(UserHandle userHandle) {
-        UserManager userManager = getSystemService(UserManager.class);
-        return userManager.isQuietModeEnabled(userHandle);
-    }
-
     private void setupScrollListener() {
         if (mResolverDrawerLayout == null) {
             return;
@@ -3211,16 +3021,6 @@ public class ChooserActivity extends ResolverActivity implements
         }
     }
 
-    static class ServiceResultInfo {
-        public final DisplayResolveInfo originalTarget;
-        public final List<ChooserTarget> resultTargets;
-
-        ServiceResultInfo(DisplayResolveInfo ot, List<ChooserTarget> rt) {
-            originalTarget = ot;
-            resultTargets = rt;
-        }
-    }
-
     static class ChooserTargetRankingInfo {
         public final List<AppTarget> scores;
         public final UserHandle userHandle;
@@ -3396,22 +3196,28 @@ public class ChooserActivity extends ResolverActivity implements
         getChooserActivityLogger().logSharesheetProfileChanged();
     }
 
-    private static <K, V> Map<K, V> emptyIfNull(@Nullable Map<K, V> map) {
-        return map == null ? Collections.emptyMap() : map;
-    }
-
     private static class ProfileRecord {
         /** The {@link AppPredictor} for this profile, if any. */
         @Nullable
         public final AppPredictor appPredictor;
-
+        /**
+         * null if we should not load shortcuts.
+         */
+        @Nullable
+        public final ShortcutLoader shortcutLoader;
         public long loadingStartTime;
 
-        ProfileRecord(@Nullable AppPredictor appPredictor) {
+        private ProfileRecord(
+                @Nullable AppPredictor appPredictor,
+                @Nullable ShortcutLoader shortcutLoader) {
             this.appPredictor = appPredictor;
+            this.shortcutLoader = shortcutLoader;
         }
 
         public void destroy() {
+            if (shortcutLoader != null) {
+                shortcutLoader.destroy();
+            }
             if (appPredictor != null) {
                 appPredictor.destroy();
             }
diff --git a/java/src/com/android/intentresolver/shortcuts/ShortcutLoader.java b/java/src/com/android/intentresolver/shortcuts/ShortcutLoader.java
new file mode 100644
index 00000000..1cfa2c8d
--- /dev/null
+++ b/java/src/com/android/intentresolver/shortcuts/ShortcutLoader.java
@@ -0,0 +1,426 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.intentresolver.shortcuts;
+
+import android.app.ActivityManager;
+import android.app.prediction.AppPredictor;
+import android.app.prediction.AppTarget;
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.IntentFilter;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.content.pm.PackageManager.ApplicationInfoFlags;
+import android.content.pm.PackageManager.NameNotFoundException;
+import android.content.pm.ShortcutInfo;
+import android.content.pm.ShortcutManager;
+import android.os.AsyncTask;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.service.chooser.ChooserTarget;
+import android.text.TextUtils;
+import android.util.Log;
+
+import androidx.annotation.MainThread;
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+import androidx.annotation.WorkerThread;
+
+import com.android.intentresolver.chooser.DisplayResolveInfo;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+/**
+ * Encapsulates shortcuts loading logic from either AppPredictor or ShortcutManager.
+ * <p>
+ *     A ShortcutLoader instance can be viewed as a per-profile singleton hot stream of shortcut
+ * updates. The shortcut loading is triggered by the {@link #queryShortcuts(DisplayResolveInfo[])},
+ * the processing will happen on the {@link #mBackgroundExecutor} and the result is delivered
+ * through the {@link #mCallback} on the {@link #mCallbackExecutor}, the main thread.
+ * </p>
+ * <p>
+ *    The current version does not improve on the legacy in a way that it does not guarantee that
+ * each invocation of the {@link #queryShortcuts(DisplayResolveInfo[])} will be matched by an
+ * invocation of the callback (there are early terminations of the flow). Also, the fetched
+ * shortcuts would be matched against the last known input, i.e. two invocations of
+ * {@link #queryShortcuts(DisplayResolveInfo[])} may result in two callbacks where shortcuts are
+ * processed against the latest input.
+ * </p>
+ */
+public class ShortcutLoader {
+    private static final String TAG = "ChooserActivity";
+
+    private static final Request NO_REQUEST = new Request(new DisplayResolveInfo[0]);
+
+    private final Context mContext;
+    @Nullable
+    private final AppPredictorProxy mAppPredictor;
+    private final UserHandle mUserHandle;
+    @Nullable
+    private final IntentFilter mTargetIntentFilter;
+    private final Executor mBackgroundExecutor;
+    private final Executor mCallbackExecutor;
+    private final boolean mIsPersonalProfile;
+    private final ShortcutToChooserTargetConverter mShortcutToChooserTargetConverter =
+            new ShortcutToChooserTargetConverter();
+    private final UserManager mUserManager;
+    private final AtomicReference<Consumer<Result>> mCallback = new AtomicReference<>();
+    private final AtomicReference<Request> mActiveRequest = new AtomicReference<>(NO_REQUEST);
+
+    @Nullable
+    private final AppPredictor.Callback mAppPredictorCallback;
+
+    @MainThread
+    public ShortcutLoader(
+            Context context,
+            @Nullable AppPredictor appPredictor,
+            UserHandle userHandle,
+            @Nullable IntentFilter targetIntentFilter,
+            Consumer<Result> callback) {
+        this(
+                context,
+                appPredictor == null ? null : new AppPredictorProxy(appPredictor),
+                userHandle,
+                userHandle.equals(UserHandle.of(ActivityManager.getCurrentUser())),
+                targetIntentFilter,
+                AsyncTask.SERIAL_EXECUTOR,
+                context.getMainExecutor(),
+                callback);
+    }
+
+    @VisibleForTesting
+    ShortcutLoader(
+            Context context,
+            @Nullable AppPredictorProxy appPredictor,
+            UserHandle userHandle,
+            boolean isPersonalProfile,
+            @Nullable IntentFilter targetIntentFilter,
+            Executor backgroundExecutor,
+            Executor callbackExecutor,
+            Consumer<Result> callback) {
+        mContext = context;
+        mAppPredictor = appPredictor;
+        mUserHandle = userHandle;
+        mTargetIntentFilter = targetIntentFilter;
+        mBackgroundExecutor = backgroundExecutor;
+        mCallbackExecutor = callbackExecutor;
+        mCallback.set(callback);
+        mIsPersonalProfile = isPersonalProfile;
+        mUserManager = (UserManager) mContext.getSystemService(Context.USER_SERVICE);
+
+        if (mAppPredictor != null) {
+            mAppPredictorCallback = createAppPredictorCallback();
+            mAppPredictor.registerPredictionUpdates(mCallbackExecutor, mAppPredictorCallback);
+        } else {
+            mAppPredictorCallback = null;
+        }
+    }
+
+    /**
+     * Unsubscribe from app predictor if one was provided.
+     */
+    @MainThread
+    public void destroy() {
+        if (mCallback.getAndSet(null) != null) {
+            if (mAppPredictor != null) {
+                mAppPredictor.unregisterPredictionUpdates(mAppPredictorCallback);
+            }
+        }
+    }
+
+    private boolean isDestroyed() {
+        return mCallback.get() == null;
+    }
+
+    /**
+     * Set new resolved targets. This will trigger shortcut loading.
+     * @param appTargets a collection of application targets a loaded set of shortcuts will be
+     *                   grouped against
+     */
+    @MainThread
+    public void queryShortcuts(DisplayResolveInfo[] appTargets) {
+        if (isDestroyed()) {
+            return;
+        }
+        mActiveRequest.set(new Request(appTargets));
+        mBackgroundExecutor.execute(this::loadShortcuts);
+    }
+
+    @WorkerThread
+    private void loadShortcuts() {
+        // no need to query direct share for work profile when its locked or disabled
+        if (!shouldQueryDirectShareTargets()) {
+            return;
+        }
+        Log.d(TAG, "querying direct share targets");
+        queryDirectShareTargets(false);
+    }
+
+    @WorkerThread
+    private void queryDirectShareTargets(boolean skipAppPredictionService) {
+        if (isDestroyed()) {
+            return;
+        }
+        if (!skipAppPredictionService && mAppPredictor != null) {
+            mAppPredictor.requestPredictionUpdate();
+            return;
+        }
+        // Default to just querying ShortcutManager if AppPredictor not present.
+        if (mTargetIntentFilter == null) {
+            return;
+        }
+
+        Context selectedProfileContext = mContext.createContextAsUser(mUserHandle, 0 /* flags */);
+        ShortcutManager sm = (ShortcutManager) selectedProfileContext
+                .getSystemService(Context.SHORTCUT_SERVICE);
+        List<ShortcutManager.ShareShortcutInfo> shortcuts =
+                sm.getShareTargets(mTargetIntentFilter);
+        sendShareShortcutInfoList(shortcuts, false, null);
+    }
+
+    private AppPredictor.Callback createAppPredictorCallback() {
+        return appPredictorTargets -> {
+            if (appPredictorTargets.isEmpty() && shouldQueryDirectShareTargets()) {
+                // APS may be disabled, so try querying targets ourselves.
+                queryDirectShareTargets(true);
+                return;
+            }
+
+            final List<ShortcutManager.ShareShortcutInfo> shortcuts = new ArrayList<>();
+            List<AppTarget> shortcutResults = new ArrayList<>();
+            for (AppTarget appTarget : appPredictorTargets) {
+                if (appTarget.getShortcutInfo() == null) {
+                    continue;
+                }
+                shortcutResults.add(appTarget);
+            }
+            appPredictorTargets = shortcutResults;
+            for (AppTarget appTarget : appPredictorTargets) {
+                shortcuts.add(new ShortcutManager.ShareShortcutInfo(
+                        appTarget.getShortcutInfo(),
+                        new ComponentName(appTarget.getPackageName(), appTarget.getClassName())));
+            }
+            sendShareShortcutInfoList(shortcuts, true, appPredictorTargets);
+        };
+    }
+
+    @WorkerThread
+    private void sendShareShortcutInfoList(
+            List<ShortcutManager.ShareShortcutInfo> shortcuts,
+            boolean isFromAppPredictor,
+            @Nullable List<AppTarget> appPredictorTargets) {
+        if (appPredictorTargets != null && appPredictorTargets.size() != shortcuts.size()) {
+            throw new RuntimeException("resultList and appTargets must have the same size."
+                    + " resultList.size()=" + shortcuts.size()
+                    + " appTargets.size()=" + appPredictorTargets.size());
+        }
+        Context selectedProfileContext = mContext.createContextAsUser(mUserHandle, 0 /* flags */);
+        for (int i = shortcuts.size() - 1; i >= 0; i--) {
+            final String packageName = shortcuts.get(i).getTargetComponent().getPackageName();
+            if (!isPackageEnabled(selectedProfileContext, packageName)) {
+                shortcuts.remove(i);
+                if (appPredictorTargets != null) {
+                    appPredictorTargets.remove(i);
+                }
+            }
+        }
+
+        HashMap<ChooserTarget, AppTarget> directShareAppTargetCache = new HashMap<>();
+        HashMap<ChooserTarget, ShortcutInfo> directShareShortcutInfoCache = new HashMap<>();
+        // Match ShareShortcutInfos with DisplayResolveInfos to be able to use the old code path
+        // for direct share targets. After ShareSheet is refactored we should use the
+        // ShareShortcutInfos directly.
+        final DisplayResolveInfo[] appTargets = mActiveRequest.get().appTargets;
+        List<ShortcutResultInfo> resultRecords = new ArrayList<>();
+        for (DisplayResolveInfo displayResolveInfo : appTargets) {
+            List<ShortcutManager.ShareShortcutInfo> matchingShortcuts =
+                    filterShortcutsByTargetComponentName(
+                            shortcuts, displayResolveInfo.getResolvedComponentName());
+            if (matchingShortcuts.isEmpty()) {
+                continue;
+            }
+
+            List<ChooserTarget> chooserTargets = mShortcutToChooserTargetConverter
+                    .convertToChooserTarget(
+                            matchingShortcuts,
+                            shortcuts,
+                            appPredictorTargets,
+                            directShareAppTargetCache,
+                            directShareShortcutInfoCache);
+
+            ShortcutResultInfo resultRecord =
+                    new ShortcutResultInfo(displayResolveInfo, chooserTargets);
+            resultRecords.add(resultRecord);
+        }
+
+        postReport(
+                new Result(
+                        isFromAppPredictor,
+                        appTargets,
+                        resultRecords.toArray(new ShortcutResultInfo[0]),
+                        directShareAppTargetCache,
+                        directShareShortcutInfoCache));
+    }
+
+    private void postReport(Result result) {
+        mCallbackExecutor.execute(() -> report(result));
+    }
+
+    @MainThread
+    private void report(Result result) {
+        Consumer<Result> callback = mCallback.get();
+        if (callback != null) {
+            callback.accept(result);
+        }
+    }
+
+    /**
+     * Returns {@code false} if {@code userHandle} is the work profile and it's either
+     * in quiet mode or not running.
+     */
+    private boolean shouldQueryDirectShareTargets() {
+        return mIsPersonalProfile || isProfileActive();
+    }
+
+    @VisibleForTesting
+    protected boolean isProfileActive() {
+        return mUserManager.isUserRunning(mUserHandle)
+                && mUserManager.isUserUnlocked(mUserHandle)
+                && !mUserManager.isQuietModeEnabled(mUserHandle);
+    }
+
+    private static boolean isPackageEnabled(Context context, String packageName) {
+        if (TextUtils.isEmpty(packageName)) {
+            return false;
+        }
+        ApplicationInfo appInfo;
+        try {
+            appInfo = context.getPackageManager().getApplicationInfo(
+                    packageName,
+                    ApplicationInfoFlags.of(PackageManager.GET_META_DATA));
+        } catch (NameNotFoundException e) {
+            return false;
+        }
+
+        return appInfo != null && appInfo.enabled
+                && (appInfo.flags & ApplicationInfo.FLAG_SUSPENDED) == 0;
+    }
+
+    private static List<ShortcutManager.ShareShortcutInfo> filterShortcutsByTargetComponentName(
+            List<ShortcutManager.ShareShortcutInfo> allShortcuts, ComponentName requiredTarget) {
+        List<ShortcutManager.ShareShortcutInfo> matchingShortcuts = new ArrayList<>();
+        for (ShortcutManager.ShareShortcutInfo shortcut : allShortcuts) {
+            if (requiredTarget.equals(shortcut.getTargetComponent())) {
+                matchingShortcuts.add(shortcut);
+            }
+        }
+        return matchingShortcuts;
+    }
+
+    private static class Request {
+        public final DisplayResolveInfo[] appTargets;
+
+        Request(DisplayResolveInfo[] targets) {
+            appTargets = targets;
+        }
+    }
+
+    /**
+     * Resolved shortcuts with corresponding app targets.
+     */
+    public static class Result {
+        public final boolean isFromAppPredictor;
+        /**
+         * Input app targets (see {@link ShortcutLoader#queryShortcuts(DisplayResolveInfo[])} the
+         * shortcuts were process against.
+         */
+        public final DisplayResolveInfo[] appTargets;
+        /**
+         * Shortcuts grouped by app target.
+         */
+        public final ShortcutResultInfo[] shortcutsByApp;
+        public final Map<ChooserTarget, AppTarget> directShareAppTargetCache;
+        public final Map<ChooserTarget, ShortcutInfo> directShareShortcutInfoCache;
+
+        @VisibleForTesting
+        public Result(
+                boolean isFromAppPredictor,
+                DisplayResolveInfo[] appTargets,
+                ShortcutResultInfo[] shortcutsByApp,
+                Map<ChooserTarget, AppTarget> directShareAppTargetCache,
+                Map<ChooserTarget, ShortcutInfo> directShareShortcutInfoCache) {
+            this.isFromAppPredictor = isFromAppPredictor;
+            this.appTargets = appTargets;
+            this.shortcutsByApp = shortcutsByApp;
+            this.directShareAppTargetCache = directShareAppTargetCache;
+            this.directShareShortcutInfoCache = directShareShortcutInfoCache;
+        }
+    }
+
+    /**
+     * Shortcuts grouped by app.
+     */
+    public static class ShortcutResultInfo {
+        public final DisplayResolveInfo appTarget;
+        public final List<ChooserTarget> shortcuts;
+
+        public ShortcutResultInfo(DisplayResolveInfo appTarget, List<ChooserTarget> shortcuts) {
+            this.appTarget = appTarget;
+            this.shortcuts = shortcuts;
+        }
+    }
+
+    /**
+     * A wrapper around AppPredictor to facilitate unit-testing.
+     */
+    @VisibleForTesting
+    public static class AppPredictorProxy {
+        private final AppPredictor mAppPredictor;
+
+        AppPredictorProxy(AppPredictor appPredictor) {
+            mAppPredictor = appPredictor;
+        }
+
+        /**
+         * {@link AppPredictor#registerPredictionUpdates}
+         */
+        public void registerPredictionUpdates(
+                Executor callbackExecutor, AppPredictor.Callback callback) {
+            mAppPredictor.registerPredictionUpdates(callbackExecutor, callback);
+        }
+
+        /**
+         * {@link AppPredictor#unregisterPredictionUpdates}
+         */
+        public void unregisterPredictionUpdates(AppPredictor.Callback callback) {
+            mAppPredictor.unregisterPredictionUpdates(callback);
+        }
+
+        /**
+         * {@link AppPredictor#requestPredictionUpdate}
+         */
+        public void requestPredictionUpdate() {
+            mAppPredictor.requestPredictionUpdate();
+        }
+    }
+}
diff --git a/java/src/com/android/intentresolver/ShortcutToChooserTargetConverter.java b/java/src/com/android/intentresolver/shortcuts/ShortcutToChooserTargetConverter.java
index ac4270d3..a37d6558 100644
--- a/java/src/com/android/intentresolver/ShortcutToChooserTargetConverter.java
+++ b/java/src/com/android/intentresolver/shortcuts/ShortcutToChooserTargetConverter.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.intentresolver;
+package com.android.intentresolver.shortcuts;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;