Merge "Fix snapshots not updating on app => overview." into ub-launcher3-master
diff --git a/go/quickstep/src/com/android/quickstep/TaskListLoader.java b/go/quickstep/src/com/android/quickstep/TaskListLoader.java
index 51b73f1..850c7e6 100644
--- a/go/quickstep/src/com/android/quickstep/TaskListLoader.java
+++ b/go/quickstep/src/com/android/quickstep/TaskListLoader.java
@@ -38,23 +38,9 @@
 
     private ArrayList<Task> mTaskList = new ArrayList<>();
     private int mTaskListChangeId;
-    private RecentsModel.TaskThumbnailChangeListener listener = (taskId, thumbnailData) -> {
-        Task foundTask = null;
-        for (Task task : mTaskList) {
-            if (task.key.id == taskId) {
-                foundTask = task;
-                break;
-            }
-        }
-        if (foundTask != null) {
-            foundTask.thumbnail = thumbnailData;
-        }
-        return foundTask;
-    };
 
     public TaskListLoader(Context context) {
         mRecentsModel = RecentsModel.INSTANCE.get(context);
-        mRecentsModel.addThumbnailChangeListener(listener);
     }
 
     /**
diff --git a/go/quickstep/src/com/android/quickstep/views/IconRecentsView.java b/go/quickstep/src/com/android/quickstep/views/IconRecentsView.java
index 6075b35..8661402 100644
--- a/go/quickstep/src/com/android/quickstep/views/IconRecentsView.java
+++ b/go/quickstep/src/com/android/quickstep/views/IconRecentsView.java
@@ -51,6 +51,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.R;
 import com.android.quickstep.ContentFillItemAnimator;
+import com.android.quickstep.RecentsModel;
 import com.android.quickstep.RecentsToActivityHelper;
 import com.android.quickstep.TaskActionController;
 import com.android.quickstep.TaskAdapter;
@@ -58,6 +59,7 @@
 import com.android.quickstep.TaskLayoutManager;
 import com.android.quickstep.TaskListLoader;
 import com.android.quickstep.TaskSwipeCallback;
+import com.android.systemui.shared.recents.model.Task;
 
 /**
  * Root view for the icon recents view. Acts as the main interface to the rest of the Launcher code
@@ -114,7 +116,20 @@
     private boolean mTransitionedFromApp;
     private AnimatorSet mLayoutAnimation;
     private final ArraySet<View> mLayingOutViews = new ArraySet<>();
-
+    private final RecentsModel.TaskThumbnailChangeListener listener = (taskId, thumbnailData) -> {
+        TaskItemView[] itemViews = getTaskViews();
+        for (TaskItemView taskView : itemViews) {
+            TaskHolder taskHolder = (TaskHolder) mTaskRecyclerView.getChildViewHolder(taskView);
+            Task task = taskHolder.getTask();
+            if (taskHolder.getTask().key.id == taskId) {
+                // Update thumbnail on the task.
+                task.thumbnail = thumbnailData;
+                taskView.setThumbnail(thumbnailData.thumbnail);
+                return task;
+            }
+        }
+        return null;
+    };
 
     public IconRecentsView(Context context, AttributeSet attrs) {
         super(context, attrs);
@@ -125,6 +140,7 @@
         mTaskAdapter = new TaskAdapter(mTaskLoader);
         mTaskActionController = new TaskActionController(mTaskLoader, mTaskAdapter);
         mTaskAdapter.setActionController(mTaskActionController);
+        RecentsModel.INSTANCE.get(context).addThumbnailChangeListener(listener);
     }
 
     @Override