Merge "Render user's actual workspace in ThemePicker preview (Part 6)" into ub-launcher3-master
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 0927b26..02e98e8 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -87,9 +87,11 @@
 import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
+import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -128,7 +130,8 @@
 
         private static final Set<MainThreadInitializedObject> WHITELIST = new HashSet<>(
                 Arrays.asList(UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
-                        LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE));
+                        LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
+                        CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE));
 
         private final InvariantDeviceProfile mIdp;
         private final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
@@ -145,6 +148,14 @@
             return this;
         }
 
+        public void onDestroy() {
+            CustomWidgetManager customWidgetManager = (CustomWidgetManager) mObjectMap.get(
+                    CustomWidgetManager.INSTANCE);
+            if (customWidgetManager != null) {
+                customWidgetManager.onDestroy();
+            }
+        }
+
         /**
          * Find a cached object from mObjectMap if we have already created one. If not, generate
          * an object using the provider.
@@ -349,6 +360,9 @@
         private void inflateAndAddWidgets(LauncherAppWidgetInfo info, WidgetsModel widgetsModel) {
             WidgetItem widgetItem = widgetsModel.getWidgetProviderInfoByProviderName(
                     info.providerName);
+            if (widgetItem == null) {
+                return;
+            }
             AppWidgetHostView view = new AppWidgetHostView(mContext);
             view.setAppWidget(-1, widgetItem.widgetInfo);
             view.updateAppWidget(null);
@@ -385,9 +399,11 @@
                 }
 
                 WorkspaceFetcher fetcher;
+                PreviewContext previewContext = null;
                 if (needsToMigrate && success) {
+                    previewContext = new PreviewContext(mContext, mIdp);
                     LauncherAppState appForPreview = new LauncherAppState(
-                            new PreviewContext(mContext, mIdp), null /* iconCacheFileName */);
+                            previewContext, null /* iconCacheFileName */);
                     fetcher = new WorkspaceItemsInfoFromPreviewFetcher(appForPreview);
                     MODEL_EXECUTOR.execute(fetcher);
                 } else {
@@ -396,6 +412,9 @@
                             (LauncherModel.ModelUpdateTask) fetcher);
                 }
                 WorkspaceResult workspaceResult = fetcher.get();
+                if (previewContext != null) {
+                    previewContext.onDestroy();
+                }
 
                 if (workspaceResult == null) {
                     return;
@@ -557,6 +576,7 @@
         public WorkspaceResult call() throws Exception {
             List<ShortcutInfo> allShortcuts = new ArrayList<>();
             loadWorkspace(allShortcuts, LauncherSettings.Favorites.PREVIEW_CONTENT_URI);
+            mBgDataModel.widgetsModel.update(mApp, null);
             return new WorkspaceResult(mBgDataModel.workspaceItems, mBgDataModel.appWidgets,
                     mBgDataModel.widgetsModel);
         }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index fbf01fc..62904ae 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -104,7 +104,7 @@
 public class LoaderTask implements Runnable {
     private static final String TAG = "LoaderTask";
 
-    private final LauncherAppState mApp;
+    protected final LauncherAppState mApp;
     private final AllAppsList mBgAllAppsList;
     protected final BgDataModel mBgDataModel;
 
diff --git a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
index 2aed936..0ea7d85 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -51,6 +51,7 @@
     public static final MainThreadInitializedObject<CustomWidgetManager> INSTANCE =
             new MainThreadInitializedObject<>(CustomWidgetManager::new);
 
+    private final Context mContext;
     /**
      * auto provider Id is an ever-increasing number that serves as the providerId whenever a new
      * custom widget has been connected.
@@ -62,6 +63,7 @@
     private Consumer<PackageUserKey> mWidgetRefreshCallback;
 
     private CustomWidgetManager(Context context) {
+        mContext = context;
         mPlugins = new SparseArray<>();
         mCustomWidgets = new ArrayList<>();
         mWidgetsIdMap = new SparseArray<>();
@@ -69,6 +71,10 @@
                 .addPluginListener(this, CustomWidgetPlugin.class, true);
     }
 
+    public void onDestroy() {
+        PluginManagerWrapper.INSTANCE.get(mContext).removePluginListener(this);
+    }
+
     @Override
     public void onPluginConnected(CustomWidgetPlugin plugin, Context context) {
         mPlugins.put(mAutoProviderId, plugin);
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index ae32692..0b99e7a 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -249,6 +249,10 @@
             ComponentName providerName) {
         ArrayList<WidgetItem> widgetsList = mWidgetsList.get(
                 new PackageItemInfo(providerName.getPackageName()));
+        if (widgetsList == null) {
+            return null;
+        }
+
         for (WidgetItem item : widgetsList) {
             if (item.componentName.equals(providerName)) {
                 return item;