Merge "Cache window and icon color of splash screen starting window" into sc-dev

3 files changed, 272 insertions, 91 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
index 1d37a128da0c..51a67e213144 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
@@ -22,7 +22,10 @@ import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import android.annotation.ColorInt;
 import android.annotation.NonNull;
 import android.app.ActivityThread;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
@@ -34,15 +37,19 @@ import android.graphics.drawable.AdaptiveIconDrawable;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.LayerDrawable;
+import android.net.Uri;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.Trace;
+import android.os.UserHandle;
+import android.util.ArrayMap;
 import android.util.Slog;
 import android.view.SurfaceControl;
 import android.view.View;
 import android.window.SplashScreenView;
 
 import com.android.internal.R;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.graphics.palette.Palette;
 import com.android.internal.graphics.palette.Quantizer;
 import com.android.internal.graphics.palette.VariationalKMeansQuantizer;
@@ -51,6 +58,8 @@ import com.android.wm.shell.common.TransactionPool;
 
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.IntSupplier;
+import java.util.function.Supplier;
 
 /**
  * Util class to create the view for a splash screen content.
@@ -76,9 +85,12 @@ public class SplashscreenContentDrawer {
     private int mBrandingImageWidth;
     private int mBrandingImageHeight;
     private int mMainWindowShiftLength;
+    private int mLastPackageContextConfigHash;
     private final TransactionPool mTransactionPool;
     private final SplashScreenWindowAttrs mTmpAttrs = new SplashScreenWindowAttrs();
     private final Handler mSplashscreenWorkerHandler;
+    @VisibleForTesting
+    final ColorCache mColorCache;
 
     SplashscreenContentDrawer(Context context, TransactionPool pool) {
         mContext = context;
@@ -92,6 +104,7 @@ public class SplashscreenContentDrawer {
                 new HandlerThread("wmshell.splashworker", THREAD_PRIORITY_TOP_APP_BOOST);
         shellSplashscreenWorkerThread.start();
         mSplashscreenWorkerHandler = shellSplashscreenWorkerThread.getThreadHandler();
+        mColorCache = new ColorCache(mContext, mSplashscreenWorkerHandler);
     }
 
     /**
@@ -183,14 +196,14 @@ public class SplashscreenContentDrawer {
         updateDensity();
 
         getWindowAttrs(context, mTmpAttrs);
-        final StartingWindowViewBuilder builder = new StartingWindowViewBuilder();
-        final int themeBGColor = peekWindowBGColor(context, this.mTmpAttrs);
+        mLastPackageContextConfigHash = context.getResources().getConfiguration().hashCode();
+        final int themeBGColor = mColorCache.getWindowColor(ai.packageName,
+                mLastPackageContextConfigHash, mTmpAttrs.mWindowBgColor, mTmpAttrs.mWindowBgResId,
+                () -> peekWindowBGColor(context, mTmpAttrs)).mBgColor;
         // TODO (b/173975965) Tracking the performance on improved splash screen.
-        return builder
-                .setContext(context)
+        return new StartingWindowViewBuilder(context, ai)
                 .setWindowBGColor(themeBGColor)
                 .makeEmptyView(emptyView)
-                .setActivityInfo(ai)
                 .build();
     }
 
@@ -232,50 +245,30 @@ public class SplashscreenContentDrawer {
     }
 
     private class StartingWindowViewBuilder {
-        private ActivityInfo mActivityInfo;
-        private Context mContext;
-        private boolean mEmptyView;
+        private final Context mContext;
+        private final ActivityInfo mActivityInfo;
 
-        // result
-        private boolean mBuildComplete = false;
-        private SplashScreenView mCachedResult;
+        private boolean mEmptyView;
         private int mThemeColor;
         private Drawable mFinalIconDrawable;
         private int mFinalIconSize = mIconSize;
 
+        StartingWindowViewBuilder(@NonNull Context context, @NonNull ActivityInfo aInfo) {
+            mContext = context;
+            mActivityInfo = aInfo;
+        }
+
         StartingWindowViewBuilder setWindowBGColor(@ColorInt int background) {
             mThemeColor = background;
-            mBuildComplete = false;
             return this;
         }
 
         StartingWindowViewBuilder makeEmptyView(boolean empty) {
             mEmptyView = empty;
-            mBuildComplete = false;
-            return this;
-        }
-
-        StartingWindowViewBuilder setActivityInfo(ActivityInfo ai) {
-            mActivityInfo = ai;
-            mBuildComplete = false;
-            return this;
-        }
-
-        StartingWindowViewBuilder setContext(Context context) {
-            mContext = context;
-            mBuildComplete = false;
             return this;
         }
 
         SplashScreenView build() {
-            if (mBuildComplete) {
-                return mCachedResult;
-            }
-            if (mContext == null || mActivityInfo == null) {
-                Slog.e(TAG, "Unable to create StartingWindowView, lack of materials!");
-                return null;
-            }
-
             Drawable iconDrawable;
             final int animationDuration;
             if (mEmptyView) {
@@ -292,7 +285,9 @@ public class SplashscreenContentDrawer {
                 final int densityDpi = mContext.getResources().getDisplayMetrics().densityDpi;
                 final int scaledIconDpi =
                         (int) (0.5f + iconScale * densityDpi * NO_BACKGROUND_SCALE);
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "getIcon");
                 iconDrawable = mIconProvider.getIcon(mActivityInfo, scaledIconDpi);
+                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
                 if (iconDrawable == null) {
                     iconDrawable = mContext.getPackageManager().getDefaultActivityIcon();
                 }
@@ -306,9 +301,7 @@ public class SplashscreenContentDrawer {
                 animationDuration = 0;
             }
 
-            mCachedResult = fillViewWithIcon(mFinalIconSize, mFinalIconDrawable, animationDuration);
-            mBuildComplete = true;
-            return mCachedResult;
+            return fillViewWithIcon(mFinalIconSize, mFinalIconDrawable, animationDuration);
         }
 
         private void createIconDrawable(Drawable iconDrawable, boolean legacy) {
@@ -331,28 +324,20 @@ public class SplashscreenContentDrawer {
             }
 
             Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "processAdaptiveIcon");
-            final AdaptiveIconDrawable adaptiveIconDrawable =
-                    (AdaptiveIconDrawable) iconDrawable;
-            final DrawableColorTester backIconTester =
-                    new DrawableColorTester(adaptiveIconDrawable.getBackground());
-
+            final AdaptiveIconDrawable adaptiveIconDrawable = (AdaptiveIconDrawable) iconDrawable;
             final Drawable iconForeground = adaptiveIconDrawable.getForeground();
-            final DrawableColorTester foreIconTester =
-                    new DrawableColorTester(iconForeground, true /* filterTransparent */);
-
-            final boolean foreComplex = foreIconTester.isComplexColor();
-            final int foreMainColor = foreIconTester.getDominateColor();
+            final ColorCache.IconColor iconColor = mColorCache.getIconColor(
+                    mActivityInfo.packageName, mActivityInfo.getIconResource(),
+                    mLastPackageContextConfigHash,
+                    () -> new DrawableColorTester(iconForeground, true /* filterTransparent */),
+                    () -> new DrawableColorTester(adaptiveIconDrawable.getBackground()));
 
             if (DEBUG) {
-                Slog.d(TAG, "foreground complex color? " + foreComplex + " main color: "
-                        + Integer.toHexString(foreMainColor));
-            }
-            final boolean backComplex = backIconTester.isComplexColor();
-            final int backMainColor = backIconTester.getDominateColor();
-            if (DEBUG) {
-                Slog.d(TAG, "background complex color? " + backComplex + " main color: "
-                        + Integer.toHexString(backMainColor));
-                Slog.d(TAG, "theme color? " + Integer.toHexString(mThemeColor));
+                Slog.d(TAG, "FgMainColor=" + Integer.toHexString(iconColor.mFgColor)
+                        + " BgMainColor=" + Integer.toHexString(iconColor.mBgColor)
+                        + " IsBgComplex=" + iconColor.mIsBgComplex
+                        + " FromCache=" + (iconColor.mReuseCount > 0)
+                        + " ThemeColor=" + Integer.toHexString(mThemeColor));
             }
 
             // Only draw the foreground of AdaptiveIcon to the splash screen if below condition
@@ -363,17 +348,17 @@ public class SplashscreenContentDrawer {
             // C. The background of the adaptive icon is grayscale, and the foreground of the
             // adaptive icon forms a certain contrast with the theme color.
             // D. Didn't specify icon background color.
-            if (!backComplex && mTmpAttrs.mIconBgColor == Color.TRANSPARENT
-                    && (isRgbSimilarInHsv(mThemeColor, backMainColor)
-                    || (backIconTester.isGrayscale()
-                    && !isRgbSimilarInHsv(mThemeColor, foreMainColor)))) {
+            if (!iconColor.mIsBgComplex && mTmpAttrs.mIconBgColor == Color.TRANSPARENT
+                    && (isRgbSimilarInHsv(mThemeColor, iconColor.mBgColor)
+                            || (iconColor.mIsBgGrayscale
+                                    && !isRgbSimilarInHsv(mThemeColor, iconColor.mFgColor)))) {
                 if (DEBUG) {
                     Slog.d(TAG, "makeSplashScreenContentView: choose fg icon");
                 }
                 // Reference AdaptiveIcon description, outer is 108 and inner is 72, so we
                 // scale by 192/160 if we only draw adaptiveIcon's foreground.
                 final float noBgScale =
-                        foreIconTester.nonTransparentRatio() < ENLARGE_FOREGROUND_ICON_THRESHOLD
+                        iconColor.mFgNonTransparentRatio < ENLARGE_FOREGROUND_ICON_THRESHOLD
                                 ? NO_BACKGROUND_SCALE : 1f;
                 // Using AdaptiveIconDrawable here can help keep the shape consistent with the
                 // current settings.
@@ -689,6 +674,150 @@ public class SplashscreenContentDrawer {
         }
     }
 
+    /** Cache the result of {@link DrawableColorTester} to reduce expensive calculation. */
+    @VisibleForTesting
+    static class ColorCache extends BroadcastReceiver {
+        /**
+         * The color may be different according to resource id and configuration (e.g. night mode),
+         * so this allows to cache more than one color per package.
+         */
+        private static final int CACHE_SIZE = 2;
+
+        /** The computed colors of packages. */
+        private final ArrayMap<String, Colors> mColorMap = new ArrayMap<>();
+
+        private static class Colors {
+            final WindowColor[] mWindowColors = new WindowColor[CACHE_SIZE];
+            final IconColor[] mIconColors = new IconColor[CACHE_SIZE];
+        }
+
+        private static class Cache {
+            /** The hash used to check whether this cache is hit. */
+            final int mHash;
+
+            /** The number of times this cache has been reused. */
+            int mReuseCount;
+
+            Cache(int hash) {
+                mHash = hash;
+            }
+        }
+
+        static class WindowColor extends Cache {
+            final int mBgColor;
+
+            WindowColor(int hash, int bgColor) {
+                super(hash);
+                mBgColor = bgColor;
+            }
+        }
+
+        static class IconColor extends Cache {
+            final int mFgColor;
+            final int mBgColor;
+            final boolean mIsBgComplex;
+            final boolean mIsBgGrayscale;
+            final float mFgNonTransparentRatio;
+
+            IconColor(int hash, int fgColor, int bgColor, boolean isBgComplex,
+                    boolean isBgGrayscale, float fgNonTransparentRatio) {
+                super(hash);
+                mFgColor = fgColor;
+                mBgColor = bgColor;
+                mIsBgComplex = isBgComplex;
+                mIsBgGrayscale = isBgGrayscale;
+                mFgNonTransparentRatio = fgNonTransparentRatio;
+            }
+        }
+
+        ColorCache(Context context, Handler handler) {
+            // This includes reinstall and uninstall.
+            final IntentFilter filter = new IntentFilter(Intent.ACTION_PACKAGE_REMOVED);
+            filter.addDataScheme(IntentFilter.SCHEME_PACKAGE);
+            context.registerReceiverAsUser(this, UserHandle.ALL, filter,
+                    null /* broadcastPermission */, handler);
+        }
+
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            final Uri packageUri = intent.getData();
+            if (packageUri != null) {
+                mColorMap.remove(packageUri.getEncodedSchemeSpecificPart());
+            }
+        }
+
+        /**
+         * Gets the existing cache if the hash matches. If null is returned, the caller can use
+         * outLeastUsedIndex to put the new cache.
+         */
+        private static <T extends Cache> T getCache(T[] caches, int hash, int[] outLeastUsedIndex) {
+            int minReuseCount = Integer.MAX_VALUE;
+            for (int i = 0; i < CACHE_SIZE; i++) {
+                final T cache = caches[i];
+                if (cache == null) {
+                    // Empty slot has the highest priority to put new cache.
+                    minReuseCount = -1;
+                    outLeastUsedIndex[0] = i;
+                    continue;
+                }
+                if (cache.mHash == hash) {
+                    cache.mReuseCount++;
+                    return cache;
+                }
+                if (cache.mReuseCount < minReuseCount) {
+                    minReuseCount = cache.mReuseCount;
+                    outLeastUsedIndex[0] = i;
+                }
+            }
+            return null;
+        }
+
+        @NonNull WindowColor getWindowColor(String packageName, int configHash, int windowBgColor,
+                int windowBgResId, IntSupplier windowBgColorSupplier) {
+            Colors colors = mColorMap.get(packageName);
+            int hash = 31 * configHash + windowBgColor;
+            hash = 31 * hash + windowBgResId;
+            final int[] leastUsedIndex = { 0 };
+            if (colors != null) {
+                final WindowColor windowColor = getCache(colors.mWindowColors, hash,
+                        leastUsedIndex);
+                if (windowColor != null) {
+                    return windowColor;
+                }
+            } else {
+                colors = new Colors();
+                mColorMap.put(packageName, colors);
+            }
+            final WindowColor windowColor = new WindowColor(hash, windowBgColorSupplier.getAsInt());
+            colors.mWindowColors[leastUsedIndex[0]] = windowColor;
+            return windowColor;
+        }
+
+        @NonNull IconColor getIconColor(String packageName, int configHash, int iconResId,
+                Supplier<DrawableColorTester> fgColorTesterSupplier,
+                Supplier<DrawableColorTester> bgColorTesterSupplier) {
+            Colors colors = mColorMap.get(packageName);
+            final int hash = configHash * 31 + iconResId;
+            final int[] leastUsedIndex = { 0 };
+            if (colors != null) {
+                final IconColor iconColor = getCache(colors.mIconColors, hash, leastUsedIndex);
+                if (iconColor != null) {
+                    return iconColor;
+                }
+            } else {
+                colors = new Colors();
+                mColorMap.put(packageName, colors);
+            }
+            final DrawableColorTester fgTester = fgColorTesterSupplier.get();
+            final DrawableColorTester bgTester = bgColorTesterSupplier.get();
+            final IconColor iconColor = new IconColor(hash, fgTester.getDominateColor(),
+                    bgTester.getDominateColor(), bgTester.isComplexColor(), bgTester.isGrayscale(),
+                    fgTester.nonTransparentRatio());
+            colors.mIconColors[leastUsedIndex[0]] = iconColor;
+            return iconColor;
+        }
+    }
+
     /**
      * Create and play the default exit animation for splash screen view.
      */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
index ff91d827b280..8463da6a0ecb 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
@@ -17,6 +17,7 @@
 package com.android.wm.shell.startingsurface;
 
 import static android.content.Context.CONTEXT_RESTRICTED;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.view.Choreographer.CALLBACK_INSETS_ANIMATION;
 import static android.view.Display.DEFAULT_DISPLAY;
 
@@ -34,6 +35,7 @@ import android.graphics.Rect;
 import android.hardware.display.DisplayManager;
 import android.os.IBinder;
 import android.os.SystemProperties;
+import android.os.Trace;
 import android.os.UserHandle;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -49,8 +51,10 @@ import android.window.StartingWindowInfo;
 import android.window.TaskSnapshot;
 
 import com.android.internal.R;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.common.annotations.ShellSplashscreenThread;
 
 import java.util.function.Supplier;
 
@@ -90,6 +94,7 @@ import java.util.function.Supplier;
  * => WM#addView -> .. waiting for Choreographer#doFrame -> relayout -> draw -> (draw the Paint
  * directly).
  */
+@ShellSplashscreenThread
 public class StartingSurfaceDrawer {
     static final String TAG = StartingSurfaceDrawer.class.getSimpleName();
     static final boolean DEBUG_SPLASH_SCREEN = StartingWindowController.DEBUG_SPLASH_SCREEN;
@@ -98,7 +103,8 @@ public class StartingSurfaceDrawer {
     private final Context mContext;
     private final DisplayManager mDisplayManager;
     private final ShellExecutor mSplashScreenExecutor;
-    private final SplashscreenContentDrawer mSplashscreenContentDrawer;
+    @VisibleForTesting
+    final SplashscreenContentDrawer mSplashscreenContentDrawer;
     private Choreographer mChoreographer;
 
     private static final boolean DEBUG_ENABLE_REVEAL_ANIMATION =
@@ -276,6 +282,7 @@ public class StartingSurfaceDrawer {
         final SplashScreenViewSupplier viewSupplier = new SplashScreenViewSupplier();
         final FrameLayout rootLayout = new FrameLayout(context);
         final Runnable setViewSynchronized = () -> {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "addSplashScreenView");
             // waiting for setContentView before relayoutWindow
             SplashScreenView contentView = viewSupplier.get();
             final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
@@ -294,13 +301,14 @@ public class StartingSurfaceDrawer {
                 }
                 record.setSplashScreenView(contentView);
             }
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         };
         mSplashscreenContentDrawer.createContentView(context, emptyView, activityInfo, taskId,
                 viewSupplier::setView);
 
         try {
             final WindowManager wm = context.getSystemService(WindowManager.class);
-            if (postAddWindow(taskId, appToken, rootLayout, wm, params)) {
+            if (addWindow(taskId, appToken, rootLayout, wm, params)) {
                 // We use the splash screen worker thread to create SplashScreenView while adding
                 // the window, as otherwise Choreographer#doFrame might be delayed on this thread.
                 // And since Choreographer#doFrame won't happen immediately after adding the window,
@@ -393,10 +401,11 @@ public class StartingSurfaceDrawer {
         ActivityTaskManager.getInstance().onSplashScreenViewCopyFinished(taskId, parcelable);
     }
 
-    protected boolean postAddWindow(int taskId, IBinder appToken, View view, WindowManager wm,
+    protected boolean addWindow(int taskId, IBinder appToken, View view, WindowManager wm,
             WindowManager.LayoutParams params) {
         boolean shouldSaveView = true;
         try {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "addRootView");
             wm.addView(view, params);
         } catch (WindowManager.BadTokenException e) {
             // ignore
@@ -404,6 +413,7 @@ public class StartingSurfaceDrawer {
                     + e.getMessage());
             shouldSaveView = false;
         } finally {
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
             if (view != null && view.getParent() == null) {
                 Slog.w(TAG, "view not successfully added to wm, removing view");
                 wm.removeViewImmediate(view);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
index 4e3e133f85b0..903e63ad6554 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
@@ -18,25 +18,27 @@ package com.android.wm.shell.startingsurface;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 
 import android.app.ActivityManager;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.graphics.Rect;
+import android.net.Uri;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
+import android.os.UserHandle;
 import android.testing.TestableContext;
 import android.view.SurfaceControl;
 import android.view.View;
@@ -58,6 +60,8 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.function.IntSupplier;
+
 /**
  * Tests for the starting surface drawer.
  */
@@ -71,6 +75,9 @@ public class StartingSurfaceDrawerTests {
     @Mock
     private TransactionPool mTransactionPool;
 
+    private final Handler mTestHandler = new Handler(Looper.getMainLooper());
+    private final TestableContext mTestContext = new TestContext(
+            InstrumentationRegistry.getInstrumentation().getTargetContext());
     TestStartingSurfaceDrawer mStartingSurfaceDrawer;
 
     static final class TestStartingSurfaceDrawer extends StartingSurfaceDrawer{
@@ -83,7 +90,7 @@ public class StartingSurfaceDrawerTests {
         }
 
         @Override
-        protected boolean postAddWindow(int taskId, IBinder appToken,
+        protected boolean addWindow(int taskId, IBinder appToken,
                 View view, WindowManager wm, WindowManager.LayoutParams params) {
             // listen for addView
             mAddWindowForTask = taskId;
@@ -101,45 +108,50 @@ public class StartingSurfaceDrawerTests {
         }
     }
 
+    private static class TestContext extends TestableContext {
+        TestContext(Context context) {
+            super(context);
+        }
+
+        @Override
+        public Context createPackageContextAsUser(String packageName, int flags, UserHandle user)
+                throws PackageManager.NameNotFoundException {
+            return this;
+        }
+
+        @Override
+        public Intent registerReceiverAsUser(BroadcastReceiver receiver, UserHandle user,
+                IntentFilter filter, String broadcastPermission, Handler scheduler) {
+            return null;
+        }
+    }
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        final TestableContext context = new TestableContext(
-                InstrumentationRegistry.getInstrumentation().getTargetContext(), null);
-        final WindowManager realWindowManager = context.getSystemService(WindowManager.class);
+        final WindowManager realWindowManager = mTestContext.getSystemService(WindowManager.class);
         final WindowMetrics metrics = realWindowManager.getMaximumWindowMetrics();
-        context.addMockSystemService(WindowManager.class, mMockWindowManager);
-
-        spyOn(context);
-        spyOn(realWindowManager);
-        try {
-            doReturn(context).when(context)
-                    .createPackageContextAsUser(anyString(), anyInt(), any());
-        } catch (PackageManager.NameNotFoundException e) {
-            //
-        }
+        mTestContext.addMockSystemService(WindowManager.class, mMockWindowManager);
+
         doReturn(metrics).when(mMockWindowManager).getMaximumWindowMetrics();
         doNothing().when(mMockWindowManager).addView(any(), any());
 
-        final HandlerExecutor testExecutor =
-                new HandlerExecutor(new Handler(Looper.getMainLooper()));
-        mStartingSurfaceDrawer = spy(new TestStartingSurfaceDrawer(context, testExecutor,
-                mTransactionPool));
+        mStartingSurfaceDrawer = spy(new TestStartingSurfaceDrawer(mTestContext,
+                new HandlerExecutor(mTestHandler), mTransactionPool));
     }
 
     @Test
     public void testAddSplashScreenSurface() {
         final int taskId = 1;
-        final Handler mainLoop = new Handler(Looper.getMainLooper());
         final StartingWindowInfo windowInfo =
                 createWindowInfo(taskId, android.R.style.Theme);
         mStartingSurfaceDrawer.addSplashScreenStartingWindow(windowInfo, mBinder, false);
-        waitHandlerIdle(mainLoop);
-        verify(mStartingSurfaceDrawer).postAddWindow(eq(taskId), eq(mBinder), any(), any(), any());
+        waitHandlerIdle(mTestHandler);
+        verify(mStartingSurfaceDrawer).addWindow(eq(taskId), eq(mBinder), any(), any(), any());
         assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, taskId);
 
         mStartingSurfaceDrawer.removeStartingWindow(windowInfo.taskInfo.taskId, null, null, false);
-        waitHandlerIdle(mainLoop);
+        waitHandlerIdle(mTestHandler);
         verify(mStartingSurfaceDrawer).removeWindowSynced(eq(taskId), any(), any(), eq(false));
         assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, 0);
     }
@@ -147,15 +159,45 @@ public class StartingSurfaceDrawerTests {
     @Test
     public void testFallbackDefaultTheme() {
         final int taskId = 1;
-        final Handler mainLoop = new Handler(Looper.getMainLooper());
         final StartingWindowInfo windowInfo =
                 createWindowInfo(taskId, 0);
         mStartingSurfaceDrawer.addSplashScreenStartingWindow(windowInfo, mBinder, false);
-        waitHandlerIdle(mainLoop);
-        verify(mStartingSurfaceDrawer).postAddWindow(eq(taskId), eq(mBinder), any(), any(), any());
+        waitHandlerIdle(mTestHandler);
+        verify(mStartingSurfaceDrawer).addWindow(eq(taskId), eq(mBinder), any(), any(), any());
         assertNotEquals(mStartingSurfaceDrawer.mViewThemeResId, 0);
     }
 
+    @Test
+    public void testColorCache() {
+        final String packageName = mTestContext.getPackageName();
+        final int configHash = 1;
+        final int windowBgColor = 0xff000000;
+        final int windowBgResId = 1;
+        final IntSupplier windowBgColorSupplier = () -> windowBgColor;
+        final SplashscreenContentDrawer.ColorCache colorCache =
+                mStartingSurfaceDrawer.mSplashscreenContentDrawer.mColorCache;
+        final SplashscreenContentDrawer.ColorCache.WindowColor windowColor1 =
+                colorCache.getWindowColor(packageName, configHash, windowBgColor, windowBgResId,
+                        windowBgColorSupplier);
+        assertEquals(windowBgColor, windowColor1.mBgColor);
+        assertEquals(0, windowColor1.mReuseCount);
+
+        final SplashscreenContentDrawer.ColorCache.WindowColor windowColor2 =
+                colorCache.getWindowColor(packageName, configHash, windowBgColor, windowBgResId,
+                        windowBgColorSupplier);
+        assertEquals(windowColor1, windowColor2);
+        assertEquals(1, windowColor1.mReuseCount);
+
+        final Intent packageRemoved = new Intent(Intent.ACTION_PACKAGE_REMOVED);
+        packageRemoved.setData(Uri.parse("package:" + packageName));
+        colorCache.onReceive(mTestContext, packageRemoved);
+
+        final SplashscreenContentDrawer.ColorCache.WindowColor windowColor3 =
+                colorCache.getWindowColor(packageName, configHash, windowBgColor, windowBgResId,
+                        windowBgColorSupplier);
+        assertEquals(0, windowColor3.mReuseCount);
+    }
+
     private StartingWindowInfo createWindowInfo(int taskId, int themeResId) {
         StartingWindowInfo windowInfo = new StartingWindowInfo();
         final ActivityInfo info = new ActivityInfo();