Apply local color extraction during drag-n-drop

Test: Drag a test widget around and observe local extract color
      is applied.

Bug: b/182282587, b/182816217
Change-Id: If63a9d91ceb2102d5d913bca85997b8be07b1adf
diff --git a/src/com/android/launcher3/CellLayout.java b/src/com/android/launcher3/CellLayout.java
index b8833cf..cc4bfe8 100644
--- a/src/com/android/launcher3/CellLayout.java
+++ b/src/com/android/launcher3/CellLayout.java
@@ -60,6 +60,7 @@
 import com.android.launcher3.accessibility.DragAndDropAccessibilityDelegate;
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.dragndrop.AppWidgetHostViewDrawable;
 import com.android.launcher3.dragndrop.DraggableView;
 import com.android.launcher3.folder.PreviewBackground;
 import com.android.launcher3.graphics.DragPreviewProvider;
@@ -180,6 +181,9 @@
     private final ArrayList<View> mIntersectingViews = new ArrayList<>();
     private final Rect mOccupiedRect = new Rect();
     private final int[] mDirectionVector = new int[2];
+    private final Workspace mWorkspace;
+    private final DeviceProfile mDeviceProfile;
+
     final int[] mPreviousReorderDirection = new int[2];
     private static final int INVALID_DIRECTION = -100;
 
@@ -209,15 +213,15 @@
         setWillNotDraw(false);
         setClipToPadding(false);
         mActivity = ActivityContext.lookupContext(context);
+        mWorkspace = Launcher.cast(mActivity).getWorkspace();
+        mDeviceProfile = mActivity.getDeviceProfile();
 
-        DeviceProfile grid = mActivity.getDeviceProfile();
-
-        mBorderSpacing = grid.cellLayoutBorderSpacingPx;
+        mBorderSpacing = mDeviceProfile.cellLayoutBorderSpacingPx;
         mCellWidth = mCellHeight = -1;
         mFixedCellWidth = mFixedCellHeight = -1;
 
-        mCountX = grid.inv.numColumns;
-        mCountY = grid.inv.numRows;
+        mCountX = mDeviceProfile.inv.numColumns;
+        mCountY = mDeviceProfile.inv.numRows;
         mOccupied =  new GridOccupancy(mCountX, mCountY);
         mTmpOccupied = new GridOccupancy(mCountX, mCountY);
 
@@ -234,7 +238,7 @@
         mBackground.setCallback(this);
         mBackground.setAlpha(0);
 
-        mReorderPreviewAnimationMagnitude = (REORDER_PREVIEW_MAGNITUDE * grid.iconSizePx);
+        mReorderPreviewAnimationMagnitude = (REORDER_PREVIEW_MAGNITUDE * mDeviceProfile.iconSizePx);
 
         // Initialize the data structures used for the drag visualization.
         mEaseOutInterpolator = Interpolators.DEACCEL_2_5; // Quint ease out
@@ -961,15 +965,18 @@
         final int oldDragCellX = mDragCell[0];
         final int oldDragCellY = mDragCell[1];
 
-        if (outlineProvider == null || outlineProvider.generatedDragOutline == null) {
-            return;
-        }
-
-        Bitmap dragOutline = outlineProvider.generatedDragOutline;
         if (cellX != oldDragCellX || cellY != oldDragCellY) {
             mDragCell[0] = cellX;
             mDragCell[1] = cellY;
 
+            applyColorExtraction(dragObject, mDragCell, spanX, spanY);
+
+            if (outlineProvider == null || outlineProvider.generatedDragOutline == null) {
+                return;
+            }
+
+            Bitmap dragOutline = outlineProvider.generatedDragOutline;
+
             final int oldIndex = mDragOutlineCurrent;
             mDragOutlineAnims[oldIndex].animateOut();
             mDragOutlineCurrent = (oldIndex + 1) % mDragOutlines.length;
@@ -1011,6 +1018,20 @@
         }
     }
 
+    /** Applies the local color extraction to a dragging widget object. */
+    private void applyColorExtraction(DropTarget.DragObject dragObject, int[] targetCell, int spanX,
+            int spanY) {
+        // Apply local extracted color if the DragView is an AppWidgetHostViewDrawable.
+        Drawable drawable = dragObject.dragView.getDrawable();
+        if (drawable instanceof AppWidgetHostViewDrawable) {
+            int screenId = mWorkspace.getIdForScreen(this);
+            int pageId = mWorkspace.getPageIndexForScreenId(screenId);
+            AppWidgetHostViewDrawable hostViewDrawable = ((AppWidgetHostViewDrawable) drawable);
+            cellToRect(targetCell[0], targetCell[1], spanX, spanY, mTempRect);
+            hostViewDrawable.getAppWidgetHostView().handleDrag(mTempRect, pageId);
+        }
+    }
+
     @SuppressLint("StringFormatMatches")
     public String getItemMoveDescription(int cellX, int cellY) {
         if (mContainerType == HOTSEAT) {
@@ -2076,7 +2097,7 @@
     private void commitTempPlacement() {
         mTmpOccupied.copyTo(mOccupied);
 
-        int screenId = Launcher.cast(mActivity).getWorkspace().getIdForScreen(this);
+        int screenId = mWorkspace.getIdForScreen(this);
         int container = Favorites.CONTAINER_DESKTOP;
 
         if (mContainerType == HOTSEAT) {
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index d1daac8..478effe 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -111,6 +111,7 @@
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 import com.android.launcher3.widget.PendingAppWidgetHostView;
 import com.android.launcher3.widget.WidgetManagerHelper;
+import com.android.launcher3.widget.dragndrop.AppWidgetHostViewDragListener;
 import com.android.systemui.plugins.shared.LauncherOverlayManager.LauncherOverlay;
 
 import java.util.ArrayList;
@@ -1527,6 +1528,9 @@
             }
         }
 
+        if (drawable instanceof AppWidgetHostViewDrawable) {
+            mDragController.addDragListener(new AppWidgetHostViewDragListener(mLauncher));
+        }
         DragView dv = mDragController.startDrag(
                 drawable,
                 draggableView,
diff --git a/src/com/android/launcher3/dragndrop/AppWidgetHostViewDrawable.java b/src/com/android/launcher3/dragndrop/AppWidgetHostViewDrawable.java
index 477bc6e..92ae670 100644
--- a/src/com/android/launcher3/dragndrop/AppWidgetHostViewDrawable.java
+++ b/src/com/android/launcher3/dragndrop/AppWidgetHostViewDrawable.java
@@ -76,4 +76,9 @@
     public ColorFilter getColorFilter() {
         return mPaint.getColorFilter();
     }
+
+    /** Returns the {@link LauncherAppWidgetHostView}. */
+    public LauncherAppWidgetHostView getAppWidgetHostView() {
+        return mAppWidgetHostView;
+    }
 }
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index df4d811..e2816f4 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -561,6 +561,11 @@
         return mInitialScale;
     }
 
+    /** Returns the current {@link Drawable} that is rendered in this view. */
+    public Drawable getDrawable() {
+        return mDrawable;
+    }
+
     private static class SpringFloatValue {
 
         private static final FloatPropertyCompat<SpringFloatValue> VALUE =
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
index df01295..8df70fb 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
@@ -16,7 +16,6 @@
 
 package com.android.launcher3.widget;
 
-import android.app.WallpaperManager;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.content.res.Configuration;
@@ -50,6 +49,7 @@
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.BaseDragLayer.TouchCompleteListener;
+import com.android.launcher3.widget.dragndrop.AppWidgetHostViewDragListener;
 
 import java.util.List;
 
@@ -74,7 +74,6 @@
     private final CheckLongPressHelper mLongPressHelper;
     protected final Launcher mLauncher;
     private final Workspace mWorkspace;
-    private final WallpaperManager mWallpaperManager;
 
     @ViewDebug.ExportedProperty(category = "launcher")
     private boolean mReinflateOnConfigChange;
@@ -85,10 +84,14 @@
     private boolean mIsScrollable;
     private boolean mIsAttachedToWindow;
     private boolean mIsAutoAdvanceRegistered;
+    private boolean mIsInDragMode = false;
     private Runnable mAutoAdvanceRunnable;
     private RectF mLastLocationRegistered = null;
+    @Nullable private AppWidgetHostViewDragListener mDragListener;
+
     // Used to store the widget size during onLayout.
     private final Rect mCurrentWidgetSize = new Rect();
+    private final Rect mWidgetSizeAtDrag = new Rect();
     private final RectF mTempRectF = new RectF();
     private final boolean mIsRtl;
 
@@ -106,7 +109,6 @@
             setOnLightBackground(true);
         }
         mIsRtl = Utilities.isRtl(context.getResources());
-        mWallpaperManager = WallpaperManager.getInstance(getContext());
         mColorExtractor = LocalColorExtractor.newInstance(getContext());
         mColorExtractor.setListener(this);
     }
@@ -118,12 +120,16 @@
         } else {
             super.setColorResources(colors);
         }
+
+        if (mDragListener != null) {
+            mDragListener.onDragContentChanged();
+        }
     }
 
     @Override
     public boolean onLongClick(View view) {
         if (mIsScrollable) {
-            DragLayer dragLayer = Launcher.getLauncher(getContext()).getDragLayer();
+            DragLayer dragLayer = mLauncher.getDragLayer();
             dragLayer.requestDisallowInterceptTouchEvent(false);
         }
         view.performLongClick();
@@ -172,7 +178,7 @@
 
     public boolean onInterceptTouchEvent(MotionEvent ev) {
         if (ev.getAction() == MotionEvent.ACTION_DOWN) {
-            DragLayer dragLayer = Launcher.getLauncher(getContext()).getDragLayer();
+            DragLayer dragLayer = mLauncher.getDragLayer();
             if (mIsScrollable) {
                 dragLayer.requestDisallowInterceptTouchEvent(true);
             }
@@ -252,70 +258,89 @@
 
         mIsScrollable = checkScrollableRecursively(this);
 
-        mCurrentWidgetSize.left = left;
-        mCurrentWidgetSize.top = top;
-        mCurrentWidgetSize.right = right;
-        mCurrentWidgetSize.bottom = bottom;
-        updateColorExtraction(mCurrentWidgetSize);
+        if (!mIsInDragMode && getTag() instanceof LauncherAppWidgetInfo) {
+            mCurrentWidgetSize.left = left;
+            mCurrentWidgetSize.top = top;
+            mCurrentWidgetSize.right = right;
+            mCurrentWidgetSize.bottom = bottom;
+            LauncherAppWidgetInfo info = (LauncherAppWidgetInfo) getTag();
+            int pageId = mWorkspace.getPageIndexForScreenId(info.screenId);
+            updateColorExtraction(mCurrentWidgetSize, pageId);
+        }
     }
 
-    private void updateColorExtraction(Rect widgetLocation) {
+    /** Starts the drag mode. */
+    public void startDrag(AppWidgetHostViewDragListener dragListener) {
+        mIsInDragMode = true;
+        mDragListener = dragListener;
+    }
+
+    /** Handles a drag event occurred on a workspace page, {@code pageId}. */
+    public void handleDrag(Rect rect, int pageId) {
+        mWidgetSizeAtDrag.set(rect);
+        updateColorExtraction(mWidgetSizeAtDrag, pageId);
+    }
+
+    /** Ends the drag mode. */
+    public void endDrag() {
+        mIsInDragMode = false;
+        mDragListener = null;
+        mWidgetSizeAtDrag.setEmpty();
+        requestLayout();
+    }
+
+    private void updateColorExtraction(Rect widgetLocation, int pageId) {
         // If the widget hasn't been measured and laid out, we cannot do this.
         if (widgetLocation.isEmpty()) {
             return;
         }
-        LauncherAppWidgetInfo info = (LauncherAppWidgetInfo) getTag();
-        if (info != null) {
-            int screenWidth = mLauncher.getDeviceProfile().widthPx;
-            int screenHeight = mLauncher.getDeviceProfile().heightPx;
-            int numScreens = mWorkspace.getNumPagesForWallpaperParallax();
-            int screenId = mIsRtl ? numScreens - info.screenId : info.screenId;
-            float relativeScreenWidth = 1f / numScreens;
-            float absoluteTop = widgetLocation.top;
-            float absoluteBottom = widgetLocation.bottom;
-            for (View v = (View) getParent();
-                    v != null && v.getId() != R.id.launcher;
-                    v = (View) v.getParent()) {
-                absoluteBottom += v.getTop();
-                absoluteTop += v.getTop();
-            }
-            float xOffset = 0;
-            View parentView = (View) getParent();
-            // The layout depends on the orientation.
-            if (getResources().getConfiguration().orientation
-                    == Configuration.ORIENTATION_LANDSCAPE) {
-                xOffset = screenHeight - mWorkspace.getPaddingRight()
-                        - parentView.getWidth();
-            } else {
-                xOffset = mWorkspace.getPaddingLeft() + parentView.getPaddingLeft();
-            }
-            // This is the position of the widget relative to the wallpaper, as expected by the
-            // local color extraction of the WallpaperManager.
-            // The coordinate system is such that, on the horizontal axis, each screen has a
-            // distinct range on the [0,1] segment. So if there are 3 screens, they will have the
-            // ranges [0, 1/3], [1/3, 2/3] and [2/3, 1]. The position on the subrange should be
-            // the position of the widget relative to the screen. For the vertical axis, this is
-            // simply the location of the widget relative to the screen.
-            mTempRectF.left = ((widgetLocation.left + xOffset) / screenWidth + screenId)
-                    * relativeScreenWidth;
-            mTempRectF.right = ((widgetLocation.right + xOffset) / screenWidth + screenId)
-                    * relativeScreenWidth;
-            mTempRectF.top = absoluteTop / screenHeight;
-            mTempRectF.bottom = absoluteBottom / screenHeight;
-            if (mTempRectF.left < 0 || mTempRectF.right > 1 || mTempRectF.top < 0
-                    || mTempRectF.bottom > 1) {
-                Log.e(LOG_TAG, "   Error, invalid relative position");
-                return;
-            }
-            if (!mTempRectF.equals(mLastLocationRegistered)) {
-                if (mLastLocationRegistered != null) {
-                    mColorExtractor.removeLocations();
-                }
-                mLastLocationRegistered = new RectF(mTempRectF);
-                mColorExtractor.addLocation(List.of(mLastLocationRegistered));
-            }
+        int screenWidth = mLauncher.getDeviceProfile().widthPx;
+        int screenHeight = mLauncher.getDeviceProfile().heightPx;
+        int numScreens = mWorkspace.getNumPagesForWallpaperParallax();
+        pageId = mIsRtl ? numScreens - pageId - 1 : pageId;
+        float relativeScreenWidth = 1f / numScreens;
+        float absoluteTop = widgetLocation.top;
+        float absoluteBottom = widgetLocation.bottom;
+        for (View v = (View) getParent();
+                v != null && v.getId() != R.id.launcher;
+                v = (View) v.getParent()) {
+            absoluteBottom += v.getTop();
+            absoluteTop += v.getTop();
+        }
+        float xOffset = 0;
+        View parentView = (View) getParent();
+        // The layout depends on the orientation.
+        if (getResources().getConfiguration().orientation == Configuration.ORIENTATION_LANDSCAPE) {
+            int parentViewWidth = parentView == null ? 0 : parentView.getWidth();
+            xOffset = screenHeight - mWorkspace.getPaddingRight() - parentViewWidth;
         } else {
-            mColorExtractor.removeLocations();
+            int parentViewPaddingLeft = parentView == null ? 0 : parentView.getPaddingLeft();
+            xOffset = mWorkspace.getPaddingLeft() + parentViewPaddingLeft;
+        }
+        // This is the position of the widget relative to the wallpaper, as expected by the
+        // local color extraction of the WallpaperManager.
+        // The coordinate system is such that, on the horizontal axis, each screen has a
+        // distinct range on the [0,1] segment. So if there are 3 screens, they will have the
+        // ranges [0, 1/3], [1/3, 2/3] and [2/3, 1]. The position on the subrange should be
+        // the position of the widget relative to the screen. For the vertical axis, this is
+        // simply the location of the widget relative to the screen.
+        mTempRectF.left = ((widgetLocation.left + xOffset) / screenWidth + pageId)
+                * relativeScreenWidth;
+        mTempRectF.right = ((widgetLocation.right + xOffset) / screenWidth + pageId)
+                * relativeScreenWidth;
+        mTempRectF.top = absoluteTop / screenHeight;
+        mTempRectF.bottom = absoluteBottom / screenHeight;
+        if (mTempRectF.left < 0 || mTempRectF.right > 1 || mTempRectF.top < 0
+                || mTempRectF.bottom > 1) {
+            Log.e(LOG_TAG, "   Error, invalid relative position");
+            return;
+        }
+        if (!mTempRectF.equals(mLastLocationRegistered)) {
+            if (mLastLocationRegistered != null) {
+                mColorExtractor.removeLocations();
+            }
+            mLastLocationRegistered = new RectF(mTempRectF);
+            mColorExtractor.addLocation(List.of(mLastLocationRegistered));
         }
     }
 
diff --git a/src/com/android/launcher3/widget/PendingItemDragHelper.java b/src/com/android/launcher3/widget/PendingItemDragHelper.java
index 8961f36..247a748 100644
--- a/src/com/android/launcher3/widget/PendingItemDragHelper.java
+++ b/src/com/android/launcher3/widget/PendingItemDragHelper.java
@@ -39,6 +39,7 @@
 import com.android.launcher3.dragndrop.DraggableView;
 import com.android.launcher3.graphics.DragPreviewProvider;
 import com.android.launcher3.icons.LauncherIcons;
+import com.android.launcher3.widget.dragndrop.AppWidgetHostViewDragListener;
 
 /**
  * Extension of {@link DragPreviewProvider} with logic specific to pending widgets/shortcuts
@@ -110,6 +111,8 @@
             }
             if (mAppWidgetHostViewPreview != null) {
                 preview = new AppWidgetHostViewDrawable(mAppWidgetHostViewPreview);
+                launcher.getDragController()
+                        .addDragListener(new AppWidgetHostViewDragListener(launcher));
             }
             if (preview == null) {
                 preview = new FastBitmapDrawable(
diff --git a/src/com/android/launcher3/widget/dragndrop/AppWidgetHostViewDragListener.java b/src/com/android/launcher3/widget/dragndrop/AppWidgetHostViewDragListener.java
new file mode 100644
index 0000000..c5e6fbd
--- /dev/null
+++ b/src/com/android/launcher3/widget/dragndrop/AppWidgetHostViewDragListener.java
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2021 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.launcher3.widget.dragndrop;
+
+import com.android.launcher3.DropTarget;
+import com.android.launcher3.Launcher;
+import com.android.launcher3.dragndrop.AppWidgetHostViewDrawable;
+import com.android.launcher3.dragndrop.DragController;
+import com.android.launcher3.dragndrop.DragOptions;
+import com.android.launcher3.widget.LauncherAppWidgetHostView;
+
+/** A drag listener of {@link LauncherAppWidgetHostView}. */
+public final class AppWidgetHostViewDragListener implements DragController.DragListener {
+    private final Launcher mLauncher;
+    private DropTarget.DragObject mDragObject;
+    private AppWidgetHostViewDrawable mAppWidgetHostViewDrawable;
+    private LauncherAppWidgetHostView mAppWidgetHostView;
+
+    public AppWidgetHostViewDragListener(Launcher launcher) {
+        mLauncher = launcher;
+    }
+
+    @Override
+    public void onDragStart(DropTarget.DragObject dragObject, DragOptions unused) {
+        if (dragObject.dragView.getDrawable() instanceof AppWidgetHostViewDrawable) {
+            mDragObject = dragObject;
+            mAppWidgetHostViewDrawable =
+                    (AppWidgetHostViewDrawable) mDragObject.dragView.getDrawable();
+            mAppWidgetHostView = mAppWidgetHostViewDrawable.getAppWidgetHostView();
+            mAppWidgetHostView.startDrag(this);
+        } else {
+            mLauncher.getDragController().removeDragListener(this);
+        }
+    }
+
+    @Override
+    public void onDragEnd() {
+        mAppWidgetHostView.endDrag();
+        mLauncher.getDragController().removeDragListener(this);
+    }
+
+    /** Notifies when there is a content change in the drag view. */
+    public void onDragContentChanged() {
+        mDragObject.dragView.invalidate();
+    }
+}