diff --git a/src/com/android/launcher3/ButtonDropTarget.java b/src/com/android/launcher3/ButtonDropTarget.java
index 468d490..477dcf8 100644
--- a/src/com/android/launcher3/ButtonDropTarget.java
+++ b/src/com/android/launcher3/ButtonDropTarget.java
@@ -217,7 +217,6 @@
         final DragView dragView = d.dragView;
         final Rect to = getIconRect(d);
         final float scale = (float) to.width() / dragView.getMeasuredWidth();
-        dragView.disableColorExtraction();
         dragView.detachContentView(/* reattachToPreviousParent= */ true);
         mDropTargetBar.deferOnDragEnd();
 
@@ -225,9 +224,6 @@
             completeDrop(d);
             mDropTargetBar.onDragEnd();
             mLauncher.getStateManager().goToState(NORMAL);
-            // Only re-enable updates once the workspace is back to normal, which will be after the
-            // current frame.
-            post(dragView::resumeColorExtraction);
         };
 
         dragLayer.animateView(d.dragView, to, scale, 0.1f, 0.1f,
diff --git a/src/com/android/launcher3/CellLayout.java b/src/com/android/launcher3/CellLayout.java
index af6cce1..28e3b9d 100644
--- a/src/com/android/launcher3/CellLayout.java
+++ b/src/com/android/launcher3/CellLayout.java
@@ -18,7 +18,6 @@
 
 import static android.animation.ValueAnimator.areAnimatorsEnabled;
 
-import static com.android.launcher3.Utilities.getBoundsForViewInDragLayer;
 import static com.android.launcher3.anim.Interpolators.DEACCEL_1_5;
 
 import android.animation.Animator;
@@ -1071,18 +1070,12 @@
         // Apply local extracted color if the DragView is an AppWidgetHostViewDrawable.
         View view = dragObject.dragView.getContentView();
         if (view instanceof LauncherAppWidgetHostView) {
-            Launcher launcher = Launcher.getLauncher(dragObject.dragView.getContext());
+            Launcher launcher = Launcher.getLauncher(getContext());
             Workspace workspace = launcher.getWorkspace();
             int screenId = workspace.getIdForScreen(this);
-            int pageId = workspace.getPageIndexForScreenId(screenId);
             cellToRect(targetCell[0], targetCell[1], spanX, spanY, mTempRect);
 
-            // Now get the rect in drag layer coordinates.
-            getBoundsForViewInDragLayer(launcher.getDragLayer(), this, mTempRect, true,
-                    mTmpFloatArray, mTempRectF);
-            Utilities.setRect(mTempRectF, mTempRect);
-
-            ((LauncherAppWidgetHostView) view).handleDrag(mTempRect, pageId);
+            ((LauncherAppWidgetHostView) view).handleDrag(mTempRect, this, screenId);
         }
     }
 
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index 91eb7f8..740bf5a 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -458,6 +458,10 @@
         cellLayoutBorderSpacingPx = isScalableGrid ? borderSpacing : 0;
     }
 
+    public Info getDisplayInfo() {
+        return mInfo;
+    }
+
     /**
      * We inset the widget padding added by the system and instead rely on the border spacing
      * between cells to create reliable consistency between widgets
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index faf2dd2..5200456 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -198,7 +198,6 @@
     private final int[] mTempXY = new int[2];
     private final float[] mTempFXY = new float[2];
     @Thunk float[] mDragViewVisualCenter = new float[2];
-    private final float[] mTempTouchCoordinates = new float[2];
 
     private SpringLoadedDragController mSpringLoadedDragController;
 
@@ -208,8 +207,6 @@
 
     private boolean mStripScreensOnPageStopMoving = false;
 
-    private DragPreviewProvider mOutlineProvider = null;
-
     private boolean mWorkspaceFadeInAdjacentScreens;
 
     final WallpaperOffsetInterpolator mWallpaperOffset;
@@ -494,7 +491,6 @@
         });
 
         mDragInfo = null;
-        mOutlineProvider = null;
         mDragSourceInternal = null;
     }
 
@@ -1333,11 +1329,7 @@
                 position[0], position[1], 0, null);
     }
 
-    public void prepareDragWithProvider(DragPreviewProvider outlineProvider) {
-        mOutlineProvider = outlineProvider;
-    }
-
-    private void onStartStateTransition(LauncherState state) {
+    private void onStartStateTransition() {
         mIsSwitchingState = true;
         mTransitionProgress = 0;
 
@@ -1358,7 +1350,7 @@
      */
     @Override
     public void setState(LauncherState toState) {
-        onStartStateTransition(toState);
+        onStartStateTransition();
         mStateTransitionAnimation.setState(toState);
         onEndStateTransition();
     }
@@ -1369,7 +1361,7 @@
     @Override
     public void setStateWithAnimation(
             LauncherState toState, StateAnimationConfig config, PendingAnimation animation) {
-        StateTransitionListener listener = new StateTransitionListener(toState);
+        StateTransitionListener listener = new StateTransitionListener();
         mStateTransitionAnimation.setStateWithAnimation(toState, config, animation);
 
         // Invalidate the pages now, so that we have the visible pages before the
@@ -1478,8 +1470,6 @@
             icon.clearPressedBackground();
         }
 
-        mOutlineProvider = previewProvider;
-
         if (draggableView == null && child instanceof DraggableView) {
             draggableView = (DraggableView) child;
         }
@@ -1791,17 +1781,14 @@
 
         int snapScreen = -1;
         boolean resizeOnDrop = false;
+        Runnable onCompleteRunnable = null;
         if (d.dragSource != this || mDragInfo == null) {
             final int[] touchXY = new int[] { (int) mDragViewVisualCenter[0],
                     (int) mDragViewVisualCenter[1] };
             onDropExternal(touchXY, dropTargetLayout, d);
         } else {
             final View cell = mDragInfo.cell;
-            final DragView dragView = d.dragView;
             boolean droppedOnOriginalCellDuringTransition = false;
-            Runnable onCompleteRunnable = dragView::resumeColorExtraction;
-
-            dragView.disableColorExtraction();
 
             if (dropTargetLayout != null && !d.cancelled) {
                 // Move internally
@@ -1912,9 +1899,7 @@
                         AppWidgetProviderInfo pInfo = hostView.getAppWidgetInfo();
                         if (pInfo != null && pInfo.resizeMode != AppWidgetProviderInfo.RESIZE_NONE
                                 && !options.isAccessibleDrag) {
-                            final Runnable previousRunnable = onCompleteRunnable;
                             onCompleteRunnable = () -> {
-                                previousRunnable.run();
                                 if (!isPageInTransition()) {
                                     AppWidgetResizeFrame.showForWidget(hostView, cellLayout);
                                 }
@@ -1983,7 +1968,7 @@
             parent.onDropChild(cell);
 
             mLauncher.getStateManager().goToState(NORMAL, SPRING_LOADED_EXIT_DELAY,
-                    forSuccessCallback(onCompleteRunnable));
+                    onCompleteRunnable == null ? null : forSuccessCallback(onCompleteRunnable));
             mStatsLogManager.logger().withItemInfo(d.dragInfo).withInstanceId(d.logInstanceId)
                     .log(LauncherEvent.LAUNCHER_ITEM_DROP_COMPLETED);
         }
@@ -2301,25 +2286,16 @@
 
         int nextPage = getNextPage();
         if (layout == null && !isPageInTransition()) {
-            // Check if the item is dragged over currentPage - 1 page
-            mTempTouchCoordinates[0] = Math.min(centerX, d.x);
-            mTempTouchCoordinates[1] = d.y;
-            layout = verifyInsidePage(nextPage + (mIsRtl ? 1 : -1), mTempTouchCoordinates);
+            layout = verifyInsidePage(nextPage + (mIsRtl ? 1 : -1), Math.min(centerX, d.x), d.y);
         }
 
         if (layout == null && !isPageInTransition()) {
-            // Check if the item is dragged over currentPage + 1 page
-            mTempTouchCoordinates[0] = Math.max(centerX, d.x);
-            mTempTouchCoordinates[1] = d.y;
-            layout = verifyInsidePage(nextPage + (mIsRtl ? -1 : 1), mTempTouchCoordinates);
+            layout = verifyInsidePage(nextPage + (mIsRtl ? -1 : 1), Math.max(centerX, d.x), d.y);
         }
 
         // If two panel is enabled, users can also drag items to currentPage + 2
         if (isTwoPanelEnabled() && layout == null && !isPageInTransition()) {
-            // Check if the item is dragged over currentPage + 2 page
-            mTempTouchCoordinates[0] = Math.max(centerX, d.x);
-            mTempTouchCoordinates[1] = d.y;
-            layout = verifyInsidePage(nextPage + (mIsRtl ? -2 : 2), mTempTouchCoordinates);
+            layout = verifyInsidePage(nextPage + (mIsRtl ? -2 : 2), Math.max(centerX, d.x), d.y);
         }
 
         // Always pick the current page.
@@ -2337,12 +2313,11 @@
     /**
      * Returns the child CellLayout if the point is inside the page coordinates, null otherwise.
      */
-    private CellLayout verifyInsidePage(int pageNo, float[] touchXy)  {
+    private CellLayout verifyInsidePage(int pageNo, float x, float y) {
         if (pageNo >= 0 && pageNo < getPageCount()) {
             CellLayout cl = (CellLayout) getChildAt(pageNo);
-            mapPointFromSelfToChild(cl, touchXy);
-            if (touchXy[0] >= 0 && touchXy[0] <= cl.getWidth() &&
-                    touchXy[1] >= 0 && touchXy[1] <= cl.getHeight()) {
+            if (x >= cl.getLeft() && x <= cl.getRight()
+                    && y >= cl.getTop() && y <= cl.getBottom()) {
                 // This point is inside the cell layout
                 return cl;
             }
@@ -3341,12 +3316,6 @@
     private class StateTransitionListener extends AnimatorListenerAdapter
             implements AnimatorUpdateListener {
 
-        private final LauncherState mToState;
-
-        StateTransitionListener(LauncherState toState) {
-            mToState = toState;
-        }
-
         @Override
         public void onAnimationUpdate(ValueAnimator anim) {
             mTransitionProgress = anim.getAnimatedFraction();
@@ -3354,7 +3323,7 @@
 
         @Override
         public void onAnimationStart(Animator animation) {
-            onStartStateTransition(mToState);
+            onStartStateTransition();
         }
 
         @Override
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index b9c9b32..57d6cc3 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -65,7 +65,6 @@
 import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
-import com.android.launcher3.widget.LauncherAppWidgetHostView;
 
 /** A custom view for rendering an icon, folder, shortcut or widget during drag-n-drop. */
 public abstract class DragView<T extends Context & ActivityContext> extends FrameLayout {
@@ -299,16 +298,6 @@
         mOnDragStartCallback.executeAllAndDestroy();
     }
 
-    // TODO(b/183609936): This is only for LauncherAppWidgetHostView that is rendered in a drawable.
-    // Once LauncherAppWidgetHostView is directly rendered in this view, removes this method.
-    @Override
-    public void invalidate() {
-        super.invalidate();
-        if (mContent instanceof ImageView) {
-            mContent.invalidate();
-        }
-    }
-
     @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
         super.onMeasure(makeMeasureSpec(mWidth, EXACTLY), makeMeasureSpec(mHeight, EXACTLY));
@@ -488,24 +477,6 @@
     }
 
     /**
-     * If the drag view uses color extraction, block it.
-     */
-    public void disableColorExtraction() {
-        if (mContent instanceof LauncherAppWidgetHostView) {
-            ((LauncherAppWidgetHostView) mContent).disableColorExtraction();
-        }
-    }
-
-    /**
-     * If the drag view uses color extraction, restores it.
-     */
-    public void resumeColorExtraction() {
-        if (mContent instanceof LauncherAppWidgetHostView) {
-            ((LauncherAppWidgetHostView) mContent).enableColorExtraction(/* updateColors= */ false);
-        }
-    }
-
-    /**
      * Removes this view from the {@link DragLayer}.
      *
      * <p>If the drag content is a {@link #mContent}, this call doesn't reattach the
diff --git a/src/com/android/launcher3/folder/LauncherDelegate.java b/src/com/android/launcher3/folder/LauncherDelegate.java
index e599e8c..c5b3913 100644
--- a/src/com/android/launcher3/folder/LauncherDelegate.java
+++ b/src/com/android/launcher3/folder/LauncherDelegate.java
@@ -18,8 +18,6 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_FOLDER_CONVERTED_TO_ICON;
 
 import android.content.Context;
-import android.graphics.Rect;
-import android.graphics.RectF;
 import android.view.MotionEvent;
 import android.view.View;
 
@@ -38,10 +36,7 @@
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
-import com.android.launcher3.views.BaseDragLayer.LayoutParams;
-import com.android.launcher3.widget.LocalColorExtractor;
 
-import java.util.Arrays;
 import java.util.Optional;
 import java.util.function.Consumer;
 
@@ -51,8 +46,6 @@
 public class LauncherDelegate {
 
     private final Launcher mLauncher;
-    private final Rect mTempRect = new Rect();
-    private final RectF mTempRectF = new RectF();
 
     private LauncherDelegate(Launcher launcher) {
         mLauncher = launcher;
@@ -84,15 +77,6 @@
         return mLauncher;
     }
 
-    void addRectForColorExtraction(BaseDragLayer.LayoutParams lp, LocalColorExtractor target) {
-        mTempRect.set(lp.x, lp.y, lp.x + lp.width, lp.y + lp.height);
-        target.getExtractedRectForViewRect(mLauncher,
-                mLauncher.getWorkspace().getCurrentPage(), mTempRect, mTempRectF);
-        if (!mTempRectF.isEmpty()) {
-            target.addLocation(Arrays.asList(mTempRectF));
-        }
-    }
-
     boolean replaceFolderWithFinalItem(Folder folder) {
         // Add the last remaining child to the workspace in place of the folder
         Runnable onCompleteRunnable = new Runnable() {
@@ -215,9 +199,6 @@
             folder.close(true);
             return true;
         }
-
-        @Override
-        void addRectForColorExtraction(LayoutParams lp, LocalColorExtractor target) { }
     }
 
     static LauncherDelegate from(ActivityContext context) {
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index a534ee3..3bb49f5 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -33,7 +33,6 @@
 import android.content.res.Resources;
 import android.graphics.Color;
 import android.graphics.Rect;
-import android.graphics.RectF;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.GradientDrawable;
@@ -50,7 +49,6 @@
 import android.widget.FrameLayout;
 
 import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.BaseDraggingActivity;
@@ -69,7 +67,6 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 
 /**
@@ -126,11 +123,8 @@
     private Runnable mOnCloseCallback = () -> { };
 
     // The rect string of the view that the arrow is attached to, in screen reference frame.
-    protected String mArrowColorRectString;
     private int mArrowColor;
-    protected final HashMap<String, View> mViewForRect = new HashMap<>();
-
-    @Nullable protected LocalColorExtractor mColorExtractor;
+    protected final List<LocalColorExtractor> mColorExtractors;
 
     private final float mElevation;
     private final int mBackgroundColor;
@@ -178,7 +172,9 @@
         boolean isAboveAnotherSurface = getTopOpenViewWithType(mLauncher, TYPE_FOLDER) != null
                 || mLauncher.getStateManager().getState() == LauncherState.ALL_APPS;
         if (!isAboveAnotherSurface && Utilities.ATLEAST_S && ENABLE_LOCAL_COLOR_POPUPS.get()) {
-            setupColorExtraction();
+            mColorExtractors = new ArrayList<>();
+        } else {
+            mColorExtractors = null;
         }
 
         if (isAboveAnotherSurface) {
@@ -323,37 +319,6 @@
         return colors.get(index, mBackgroundColor);
     }
 
-    @TargetApi(Build.VERSION_CODES.S)
-    private void setupColorExtraction() {
-        Workspace workspace = mLauncher.findViewById(R.id.workspace);
-        if (workspace == null) {
-            return;
-        }
-
-        mColorExtractor = LocalColorExtractor.newInstance(mLauncher);
-        mColorExtractor.setListener((rect, extractedColors) -> {
-            String rectString = rect.toShortString();
-            View v = mViewForRect.get(rectString);
-            AnimatorSet colors = new AnimatorSet();
-            if (v != null) {
-                int newColor = getExtractedColor(extractedColors);
-                setChildColor(v, newColor, colors);
-                int numChildren = v instanceof ViewGroup ? ((ViewGroup) v).getChildCount() : 0;
-                for (int i = 0; i < numChildren; ++i) {
-                    View childView = ((ViewGroup) v).getChildAt(i);
-                    setChildColor(childView, newColor, colors);
-
-                }
-                if (rectString.equals(mArrowColorRectString)) {
-                    mArrowColor = newColor;
-                    updateArrowColor();
-                }
-            }
-            colors.setDuration(150);
-            v.post(colors::start);
-        });
-    }
-
     protected void addPreDrawForColorExtraction(Launcher launcher) {
         getViewTreeObserver().addOnPreDrawListener(new ViewTreeObserver.OnPreDrawListener() {
             @Override
@@ -374,40 +339,55 @@
     }
 
     private void initColorExtractionLocations(Launcher launcher) {
-        if (mColorExtractor == null) {
+        if (mColorExtractors == null) {
             return;
         }
-        ArrayList<RectF> locations = new ArrayList<>();
+        Workspace workspace = launcher.getWorkspace();
+        if (workspace == null) {
+            return;
+        }
 
         boolean firstVisibleChild = true;
+        int screenId = workspace.getScreenIdForPageIndex(workspace.getCurrentPage());
+        DragLayer dragLayer = launcher.getDragLayer();
+
+        final View[] viewAlignedWithArrow = new View[1];
+
         // Order matters here, since we need the arrow to match the color of its adjacent view.
-        for (View view : getChildrenForColorExtraction()) {
+        for (final View view : getChildrenForColorExtraction()) {
             if (view != null && view.getVisibility() == VISIBLE) {
-                RectF rf = new RectF();
-                mColorExtractor.getExtractedRectForView(launcher,
-                        launcher.getWorkspace().getCurrentPage(), view, rf);
-                if (!rf.isEmpty()) {
-                    locations.add(rf);
-                    String rectString = rf.toShortString();
-                    mViewForRect.put(rectString, view);
-                    if (mIsAboveIcon) {
-                        mArrowColorRectString = rectString;
-                    } else {
-                        if (firstVisibleChild) {
-                            mArrowColorRectString = rectString;
+                Rect pos = new Rect();
+                dragLayer.getDescendantRectRelativeToSelf(view, pos);
+                if (!pos.isEmpty()) {
+                    LocalColorExtractor extractor = LocalColorExtractor.newInstance(launcher);
+                    extractor.setWorkspaceLocation(pos, dragLayer, screenId);
+                    extractor.setListener(extractedColors -> {
+                        AnimatorSet colors = new AnimatorSet();
+                        int newColor = getExtractedColor(extractedColors);
+                        setChildColor(view, newColor, colors);
+                        int numChildren = view instanceof ViewGroup
+                                ? ((ViewGroup) view).getChildCount() : 0;
+                        for (int i = 0; i < numChildren; ++i) {
+                            View childView = ((ViewGroup) view).getChildAt(i);
+                            setChildColor(childView, newColor, colors);
                         }
-                    }
+                        if (viewAlignedWithArrow[0] == view) {
+                            mArrowColor = newColor;
+                            updateArrowColor();
+                        }
+                        colors.setDuration(150);
+                        view.post(colors::start);
+                    });
+                    mColorExtractors.add(extractor);
 
-                    if (firstVisibleChild) {
-                        firstVisibleChild = false;
+                    if (mIsAboveIcon || firstVisibleChild) {
+                        viewAlignedWithArrow[0] = view;
                     }
-
+                    firstVisibleChild = false;
                 }
             }
         }
-        if (!locations.isEmpty()) {
-            mColorExtractor.addLocation(locations);
-        }
+
     }
 
     /**
@@ -807,11 +787,8 @@
         getPopupContainer().removeView(this);
         getPopupContainer().removeView(mArrow);
         mOnCloseCallback.run();
-        mArrowColorRectString = null;
-        mViewForRect.clear();
-        if (mColorExtractor != null) {
-            mColorExtractor.removeLocations();
-            mColorExtractor.setListener(null);
+        if (mColorExtractors != null) {
+            mColorExtractors.forEach(e -> e.setListener(null));
         }
     }
 
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
index 0313e68..f0b4ba0 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
@@ -19,9 +19,7 @@
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.content.res.Configuration;
-import android.graphics.Canvas;
 import android.graphics.Rect;
-import android.graphics.RectF;
 import android.os.Handler;
 import android.os.SystemClock;
 import android.util.SparseBooleanArray;
@@ -35,23 +33,17 @@
 import android.widget.Advanceable;
 import android.widget.RemoteViews;
 
-import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.CheckLongPressHelper;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
-import com.android.launcher3.Workspace;
 import com.android.launcher3.dragndrop.DragLayer;
-import com.android.launcher3.keyboard.ViewGroupFocusHelper;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.BaseDragLayer.TouchCompleteListener;
-import com.android.launcher3.widget.dragndrop.AppWidgetHostViewDragListener;
-
-import java.util.List;
 
 /**
  * {@inheritDoc}
@@ -60,8 +52,6 @@
         implements TouchCompleteListener, View.OnLongClickListener,
         LocalColorExtractor.Listener {
 
-    private static final String LOG_TAG = "LauncherAppWidgetHostView";
-
     // Related to the auto-advancing of widgets
     private static final long ADVANCE_INTERVAL = 20000;
     private static final long ADVANCE_STAGGER = 250;
@@ -71,9 +61,9 @@
     // Maximum duration for which updates can be deferred.
     private static final long UPDATE_LOCK_TIMEOUT_MILLIS = 1000;
 
+    private final Rect mTempRect = new Rect();
     private final CheckLongPressHelper mLongPressHelper;
     protected final Launcher mLauncher;
-    private final Workspace mWorkspace;
 
     @ViewDebug.ExportedProperty(category = "launcher")
     private boolean mReinflateOnConfigChange;
@@ -85,24 +75,14 @@
     private boolean mIsAttachedToWindow;
     private boolean mIsAutoAdvanceRegistered;
     private Runnable mAutoAdvanceRunnable;
-    private RectF mLastLocationRegistered = null;
 
-    // Used to store the widget sizes in drag layer coordinates.
-    private final Rect mCurrentWidgetSize = new Rect();
-    private final Rect mWidgetSizeAtDrag = new Rect();
-
-    private final RectF mTempRectF = new RectF();
-    private final Object mUpdateLock = new Object();
-    private final ViewGroupFocusHelper mDragLayerRelativeCoordinateHelper;
     private long mDeferUpdatesUntilMillis = 0;
     private RemoteViews mDeferredRemoteViews;
     private boolean mHasDeferredColorChange = false;
     private @Nullable SparseIntArray mDeferredColorChange = null;
-    private boolean mEnableColorExtraction = true;
 
     // The following member variables are only used during drag-n-drop.
     private boolean mIsInDragMode = false;
-    @Nullable private AppWidgetHostViewDragListener mDragListener;
     /** The drag content width which is only set when the drag content scale is not 1f. */
     private int mDragContentWidth = 0;
     /** The drag content height which is only set when the drag content scale is not 1f. */
@@ -111,7 +91,6 @@
     public LauncherAppWidgetHostView(Context context) {
         super(context);
         mLauncher = Launcher.getLauncher(context);
-        mWorkspace = mLauncher.getWorkspace();
         mLongPressHelper = new CheckLongPressHelper(this, this);
         setAccessibilityDelegate(mLauncher.getAccessibilityDelegate());
         setBackgroundResource(R.drawable.widget_internal_focus_bg);
@@ -120,9 +99,6 @@
             setOnLightBackground(true);
         }
         mColorExtractor = LocalColorExtractor.newInstance(getContext());
-        mColorExtractor.setListener(this);
-
-        mDragLayerRelativeCoordinateHelper = new ViewGroupFocusHelper(mLauncher.getDragLayer());
     }
 
     @Override
@@ -135,14 +111,6 @@
     }
 
     @Override
-    protected void onDraw(Canvas canvas) {
-        super.onDraw(canvas);
-        if (mIsInDragMode && mDragListener != null) {
-            mDragListener.onDragContentChanged();
-        }
-    }
-
-    @Override
     public boolean onLongClick(View view) {
         if (mIsScrollable) {
             DragLayer dragLayer = mLauncher.getDragLayer();
@@ -154,13 +122,11 @@
 
     @Override
     public void updateAppWidget(RemoteViews remoteViews) {
-        synchronized (mUpdateLock) {
-            if (isDeferringUpdates()) {
-                mDeferredRemoteViews = remoteViews;
-                return;
-            }
-            mDeferredRemoteViews = null;
+        if (isDeferringUpdates()) {
+            mDeferredRemoteViews = remoteViews;
+            return;
         }
+        mDeferredRemoteViews = null;
 
         super.updateAppWidget(remoteViews);
 
@@ -211,9 +177,7 @@
      * {@link #onColorsChanged} call after {@link #UPDATE_LOCK_TIMEOUT_MILLIS} have elapsed.
      */
     public void beginDeferringUpdates() {
-        synchronized (mUpdateLock) {
-            mDeferUpdatesUntilMillis = SystemClock.uptimeMillis() + UPDATE_LOCK_TIMEOUT_MILLIS;
-        }
+        mDeferUpdatesUntilMillis = SystemClock.uptimeMillis() + UPDATE_LOCK_TIMEOUT_MILLIS;
     }
 
     /**
@@ -225,20 +189,19 @@
         RemoteViews remoteViews;
         SparseIntArray deferredColors;
         boolean hasDeferredColors;
-        synchronized (mUpdateLock) {
-            mDeferUpdatesUntilMillis = 0;
-            remoteViews = mDeferredRemoteViews;
-            mDeferredRemoteViews = null;
-            deferredColors = mDeferredColorChange;
-            hasDeferredColors = mHasDeferredColorChange;
-            mDeferredColorChange = null;
-            mHasDeferredColorChange = false;
-        }
+        mDeferUpdatesUntilMillis = 0;
+        remoteViews = mDeferredRemoteViews;
+        mDeferredRemoteViews = null;
+        deferredColors = mDeferredColorChange;
+        hasDeferredColors = mHasDeferredColorChange;
+        mDeferredColorChange = null;
+        mHasDeferredColorChange = false;
+
         if (remoteViews != null) {
             updateAppWidget(remoteViews);
         }
         if (hasDeferredColors) {
-            onColorsChanged(null /* rectF */, deferredColors);
+            onColorsChanged(deferredColors);
         }
     }
 
@@ -263,13 +226,9 @@
     @Override
     protected void onAttachedToWindow() {
         super.onAttachedToWindow();
-
         mIsAttachedToWindow = true;
         checkIfAutoAdvance();
-
-        if (mLastLocationRegistered != null) {
-            mColorExtractor.addLocation(List.of(mLastLocationRegistered));
-        }
+        mColorExtractor.setListener(this);
     }
 
     @Override
@@ -280,7 +239,7 @@
         // state is updated. So isAttachedToWindow() will return true until next frame.
         mIsAttachedToWindow = false;
         checkIfAutoAdvance();
-        mColorExtractor.removeLocations();
+        mColorExtractor.setListener(null);
     }
 
     @Override
@@ -311,9 +270,13 @@
     @Override
     protected void onLayout(boolean changed, int left, int top, int right, int bottom) {
         super.onLayout(changed, left, top, right, bottom);
-
         mIsScrollable = checkScrollableRecursively(this);
-        updateColorExtraction();
+
+        if (!mIsInDragMode && getTag() instanceof LauncherAppWidgetInfo) {
+            LauncherAppWidgetInfo info = (LauncherAppWidgetInfo) getTag();
+            mTempRect.set(left, top, right, bottom);
+            mColorExtractor.setWorkspaceLocation(mTempRect, (View) getParent(), info.screenId);
+        }
     }
 
     @Override
@@ -327,9 +290,8 @@
     }
 
     /** Starts the drag mode. */
-    public void startDrag(AppWidgetHostViewDragListener dragListener) {
+    public void startDrag() {
         mIsInDragMode = true;
-        mDragListener = dragListener;
         // In the case of dragging a scaled preview from widgets picker, we should reuse the
         // previously measured dimension from WidgetCell#measureAndComputeWidgetPreviewScale, which
         // measures the dimension of a widget preview without its parent's bound before scaling
@@ -340,98 +302,30 @@
         }
     }
 
-    /** Handles a drag event occurred on a workspace page, {@code pageId}. */
-    public void handleDrag(Rect rectInDragLayer, int pageId) {
-        mWidgetSizeAtDrag.set(rectInDragLayer);
-        updateColorExtraction(mWidgetSizeAtDrag, pageId);
+    /** Handles a drag event occurred on a workspace page corresponding to the {@code screenId}. */
+    public void handleDrag(Rect rectInView, View view, int screenId) {
+        if (mIsInDragMode) {
+            mColorExtractor.setWorkspaceLocation(rectInView, view, screenId);
+        }
     }
 
     /** Ends the drag mode. */
     public void endDrag() {
         mIsInDragMode = false;
-        mDragListener = null;
         mDragContentWidth = 0;
         mDragContentHeight = 0;
-        mWidgetSizeAtDrag.setEmpty();
-    }
-
-    /**
-     * @param rectInDragLayer Rect of widget in drag layer coordinates.
-     * @param pageId The workspace page the widget is on.
-     */
-    private void updateColorExtraction(Rect rectInDragLayer, int pageId) {
-        if (!mEnableColorExtraction) return;
-        mColorExtractor.getExtractedRectForViewRect(mLauncher, pageId, rectInDragLayer, mTempRectF);
-
-        if (mTempRectF.isEmpty()) {
-            return;
-        }
-        if (!isSameLocation(mTempRectF, mLastLocationRegistered, /* epsilon= */ 1e-6f)) {
-            if (mLastLocationRegistered != null) {
-                mColorExtractor.removeLocations();
-            }
-            mLastLocationRegistered = new RectF(mTempRectF);
-            mColorExtractor.addLocation(List.of(mLastLocationRegistered));
-        }
-    }
-
-    /**
-     * Update the color extraction, using the current position of the app widget.
-     */
-    private void updateColorExtraction() {
-        if (!mIsInDragMode && getTag() instanceof LauncherAppWidgetInfo) {
-            LauncherAppWidgetInfo info = (LauncherAppWidgetInfo) getTag();
-            mDragLayerRelativeCoordinateHelper.viewToRect(this, mCurrentWidgetSize);
-            updateColorExtraction(mCurrentWidgetSize,
-                    mWorkspace.getPageIndexForScreenId(info.screenId));
-        }
-    }
-
-    /**
-     * Enables the local color extraction.
-     *
-     * @param updateColors If true, this will update the color extraction using the current location
-     *                    of the App Widget.
-     */
-    public void enableColorExtraction(boolean updateColors) {
-        mEnableColorExtraction = true;
-        if (updateColors) {
-            updateColorExtraction();
-        }
-    }
-
-    /**
-     * Disables the local color extraction.
-     */
-    public void disableColorExtraction() {
-        mEnableColorExtraction = false;
-    }
-
-    // Compare two location rectangles. Locations are always in the [0;1] range.
-    private static boolean isSameLocation(@NonNull RectF rect1, @Nullable RectF rect2,
-            float epsilon) {
-        if (rect2 == null) return false;
-        return isSameCoordinate(rect1.left, rect2.left, epsilon)
-                && isSameCoordinate(rect1.right, rect2.right, epsilon)
-                && isSameCoordinate(rect1.top, rect2.top, epsilon)
-                && isSameCoordinate(rect1.bottom, rect2.bottom, epsilon);
-    }
-
-    private static boolean isSameCoordinate(float c1, float c2, float epsilon) {
-        return Math.abs(c1 - c2) < epsilon;
+        requestLayout();
     }
 
     @Override
-    public void onColorsChanged(RectF rectF, SparseIntArray colors) {
-        synchronized (mUpdateLock) {
-            if (isDeferringUpdates()) {
-                mDeferredColorChange = colors;
-                mHasDeferredColorChange = true;
-                return;
-            }
-            mDeferredColorChange = null;
-            mHasDeferredColorChange = false;
+    public void onColorsChanged(SparseIntArray colors) {
+        if (isDeferringUpdates()) {
+            mDeferredColorChange = colors;
+            mHasDeferredColorChange = true;
+            return;
         }
+        mDeferredColorChange = null;
+        mHasDeferredColorChange = false;
 
         // setColorResources will reapply the view, which must happen in the UI thread.
         post(() -> setColorResources(colors));
diff --git a/src/com/android/launcher3/widget/LocalColorExtractor.java b/src/com/android/launcher3/widget/LocalColorExtractor.java
index 23d9e15..96e7531 100644
--- a/src/com/android/launcher3/widget/LocalColorExtractor.java
+++ b/src/com/android/launcher3/widget/LocalColorExtractor.java
@@ -20,18 +20,14 @@
 import android.appwidget.AppWidgetHostView;
 import android.content.Context;
 import android.graphics.Rect;
-import android.graphics.RectF;
 import android.util.SparseIntArray;
 import android.view.View;
 
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
 import com.android.launcher3.util.ResourceBasedOverride;
 
-import java.util.List;
-
 /** Extracts the colors we need from the wallpaper at given locations. */
 public class LocalColorExtractor implements ResourceBasedOverride {
 
@@ -44,7 +40,7 @@
          * their value, in a format that can be passed directly to
          * {@link AppWidgetHostView#setColorResources(SparseIntArray)}.
          */
-        void onColorsChanged(RectF rect, SparseIntArray extractedColors);
+        void onColorsChanged(SparseIntArray extractedColors);
     }
 
     /**
@@ -60,15 +56,13 @@
         // no-op
     }
 
-    /** Adds a list of locations to track with this listener. */
-    public void addLocation(List<RectF> locations) {
-        // no-op
-    }
-
-    /** Stops tracking any locations. */
-    public void removeLocations() {
-        // no-op
-    }
+    /**
+     * Sets the location used for color extraction
+     * @param pos position to use for color extraction
+     * @param child view whose coordinate space is used for {@code pos}
+     * @param screenId the workspace screenId
+     */
+    public void setWorkspaceLocation(Rect pos, View child, int screenId) { }
 
     /**
      * Updates the base context to contain the colors override
@@ -83,32 +77,4 @@
         return null;
     }
 
-    /**
-     * Takes a view and returns its rect that can be used by the wallpaper local color extractor.
-     *
-     * @param launcher Launcher class class.
-     * @param pageId The page the workspace item is on.
-     * @param v The view.
-     * @param colorExtractionRectOut The location rect, but converted to a format expected by the
-     *                               wallpaper local color extractor.
-     */
-    public void getExtractedRectForView(Launcher launcher, int pageId, View v,
-            RectF colorExtractionRectOut) {
-        // no-op
-    }
-
-    /**
-     * Takes a rect in drag layer coordinates and returns the rect that can be used by the wallpaper
-     * local color extractor.
-     *
-     * @param launcher Launcher class.
-     * @param pageId The page the workspace item is on.
-     * @param rectInDragLayer The relevant bounds of the view in drag layer coordinates.
-     * @param colorExtractionRectOut The location rect, but converted to a format expected by the
-     *                               wallpaper local color extractor.
-     */
-    public void getExtractedRectForViewRect(Launcher launcher, int pageId, Rect rectInDragLayer,
-            RectF colorExtractionRectOut) {
-        // no-op
-    }
 }
diff --git a/src/com/android/launcher3/widget/PendingItemDragHelper.java b/src/com/android/launcher3/widget/PendingItemDragHelper.java
index a4003d4..991910d 100644
--- a/src/com/android/launcher3/widget/PendingItemDragHelper.java
+++ b/src/com/android/launcher3/widget/PendingItemDragHelper.java
@@ -205,10 +205,6 @@
             draggableView = DraggableView.ofType(DraggableView.DRAGGABLE_ICON);
         }
 
-        // Since we are not going through the workspace for starting the drag, set drag related
-        // information on the workspace before starting the drag.
-        launcher.getWorkspace().prepareDragWithProvider(this);
-
         int dragLayerX = screenPos.x + previewBounds.left
                 + (int) ((scale * previewWidth - previewWidth) / 2);
         int dragLayerY = screenPos.y + previewBounds.top
diff --git a/src/com/android/launcher3/widget/dragndrop/AppWidgetHostViewDragListener.java b/src/com/android/launcher3/widget/dragndrop/AppWidgetHostViewDragListener.java
index 4a60983..3e54b33 100644
--- a/src/com/android/launcher3/widget/dragndrop/AppWidgetHostViewDragListener.java
+++ b/src/com/android/launcher3/widget/dragndrop/AppWidgetHostViewDragListener.java
@@ -24,7 +24,6 @@
 /** A drag listener of {@link LauncherAppWidgetHostView}. */
 public final class AppWidgetHostViewDragListener implements DragController.DragListener {
     private final Launcher mLauncher;
-    private DropTarget.DragObject mDragObject;
     private LauncherAppWidgetHostView mAppWidgetHostView;
 
     public AppWidgetHostViewDragListener(Launcher launcher) {
@@ -34,9 +33,8 @@
     @Override
     public void onDragStart(DropTarget.DragObject dragObject, DragOptions unused) {
         if (dragObject.dragView.getContentView() instanceof LauncherAppWidgetHostView) {
-            mDragObject = dragObject;
             mAppWidgetHostView = (LauncherAppWidgetHostView) dragObject.dragView.getContentView();
-            mAppWidgetHostView.startDrag(this);
+            mAppWidgetHostView.startDrag();
         } else {
             mLauncher.getDragController().removeDragListener(this);
         }
@@ -47,11 +45,4 @@
         mAppWidgetHostView.endDrag();
         mLauncher.getDragController().removeDragListener(this);
     }
-
-    /** Notifies when there is a content change in the drag view. */
-    public void onDragContentChanged() {
-        if (mDragObject.dragView != null) {
-            mDragObject.dragView.invalidate();
-        }
-    }
 }