2 files changed, 56 insertions, 26 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageTileSet.java b/packages/SystemUI/src/com/android/systemui/screenshot/ImageTileSet.java
index ae3cd9996f04..6743afa3ab59 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageTileSet.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageTileSet.java
@@ -15,6 +15,7 @@
  */
 package com.android.systemui.screenshot;
 
+import android.annotation.AnyThread;
 import android.graphics.Bitmap;
 import android.graphics.HardwareRenderer;
 import android.graphics.RecordingCanvas;
@@ -26,6 +27,9 @@ import android.util.Log;
 
 import androidx.annotation.UiThread;
 
+import com.android.internal.util.CallbackRegistry;
+import com.android.internal.util.CallbackRegistry.NotifierCallback;
+
 import java.util.ArrayList;
 import java.util.List;
 
@@ -34,10 +38,14 @@ import java.util.List;
  * <p>
  * To display on-screen, use {@link #getDrawable()}.
  */
+@UiThread
 class ImageTileSet {
 
     private static final String TAG = "ImageTileSet";
 
+    private CallbackRegistry<OnBoundsChangedListener, ImageTileSet, Rect> mOnBoundsListeners;
+    private CallbackRegistry<OnContentChangedListener, ImageTileSet, Rect> mContentListeners;
+
     ImageTileSet(@UiThread Handler handler) {
         mHandler = handler;
     }
@@ -64,15 +72,43 @@ class ImageTileSet {
     private OnContentChangedListener mOnContentChangedListener;
     private OnBoundsChangedListener mOnBoundsChangedListener;
 
-    void setOnBoundsChangedListener(OnBoundsChangedListener listener) {
-        mOnBoundsChangedListener = listener;
-    }
-
-    void setOnContentChangedListener(OnContentChangedListener listener) {
-        mOnContentChangedListener = listener;
+    void addOnBoundsChangedListener(OnBoundsChangedListener listener) {
+        if (mOnBoundsListeners == null) {
+            mOnBoundsListeners = new CallbackRegistry<>(
+                    new NotifierCallback<OnBoundsChangedListener, ImageTileSet, Rect>() {
+                        @Override
+                        public void onNotifyCallback(OnBoundsChangedListener callback,
+                                ImageTileSet sender,
+                                int arg, Rect newBounds) {
+                            callback.onBoundsChanged(newBounds.left, newBounds.top, newBounds.right,
+                                    newBounds.bottom);
+                        }
+                    });
+        }
+        mOnBoundsListeners.add(listener);
+    }
+
+    void addOnContentChangedListener(OnContentChangedListener listener) {
+        if (mContentListeners == null) {
+            mContentListeners = new CallbackRegistry<>(
+                    new NotifierCallback<OnContentChangedListener, ImageTileSet, Rect>() {
+                        @Override
+                        public void onNotifyCallback(OnContentChangedListener callback,
+                                ImageTileSet sender,
+                                int arg, Rect newBounds) {
+                            callback.onContentChanged();
+                        }
+                    });
+        }
+        mContentListeners.add(listener);
     }
 
+    @AnyThread
     void addTile(ImageTile tile) {
+        if (!mHandler.getLooper().isCurrentThread()) {
+            mHandler.post(() -> addTile(tile));
+            return;
+        }
         final Rect newBounds = new Rect(mBounds);
         final Rect newRect = tile.getLocation();
         mTiles.add(tile);
@@ -84,27 +120,15 @@ class ImageTileSet {
         notifyContentChanged();
     }
 
-    void notifyContentChanged() {
-        if (mOnContentChangedListener == null) {
-            return;
-        }
-        if (mHandler.getLooper().isCurrentThread()) {
-            mOnContentChangedListener.onContentChanged();
-        } else {
-            mHandler.post(() -> mOnContentChangedListener.onContentChanged());
+    private void notifyContentChanged() {
+        if (mContentListeners != null) {
+            mContentListeners.notifyCallbacks(this, 0, null);
         }
     }
 
-    void notifyBoundsChanged(Rect bounds) {
-        if (mOnBoundsChangedListener == null) {
-            return;
-        }
-        if (mHandler.getLooper().isCurrentThread()) {
-            mOnBoundsChangedListener.onBoundsChanged(
-                    bounds.left, bounds.top, bounds.right, bounds.bottom);
-        } else {
-            mHandler.post(() -> mOnBoundsChangedListener.onBoundsChanged(
-                    bounds.left, bounds.top, bounds.right, bounds.bottom));
+    private void notifyBoundsChanged(Rect bounds) {
+        if (mOnBoundsListeners != null) {
+            mOnBoundsListeners.notifyCallbacks(this, 0, bounds);
         }
     }
 
@@ -180,8 +204,13 @@ class ImageTileSet {
         return mBounds.height();
     }
 
+    @AnyThread
     void clear() {
-        if (mBounds.isEmpty()) {
+        if (!mHandler.getLooper().isCurrentThread()) {
+            mHandler.post(this::clear);
+            return;
+        }
+        if (mTiles.isEmpty()) {
             return;
         }
         mBounds.setEmpty();
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TiledImageDrawable.java b/packages/SystemUI/src/com/android/systemui/screenshot/TiledImageDrawable.java
index 4ec8eb22c67a..71df369aa7b8 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TiledImageDrawable.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TiledImageDrawable.java
@@ -38,9 +38,10 @@ public class TiledImageDrawable extends Drawable {
 
     public TiledImageDrawable(ImageTileSet tiles) {
         mTiles = tiles;
-        mTiles.setOnContentChangedListener(this::onContentChanged);
+        mTiles.addOnContentChangedListener(this::onContentChanged);
     }
 
+
     private void onContentChanged() {
         if (mNode != null && mNode.hasDisplayList()) {
             mNode.discardDisplayList();