Revert "Revert "Use color extraction for arrow popup.""

This reverts commit 17fda29bf35bc4f21ede07319ffe26c88fea1b81.

Reason for revert: Fixing the memory leak

Leak was detected in caused launcher tapl tests to fail.
For instructions on how to debug leaks, see go/launcher-testing.
Using AS memory profiler, the color extractor class is holding
onto a reference to launcher that has already been destroyed.
The fix is to set the listener to null when view is closed.

Bug: 175329686
Test: run profiler in AS, open and close popups, and verify no leaks

Change-Id: I2d4a87f0481cbc88f3587dce91278e85ca113545
diff --git a/src/com/android/launcher3/CellLayout.java b/src/com/android/launcher3/CellLayout.java
index 2f755e1..af625a5 100644
--- a/src/com/android/launcher3/CellLayout.java
+++ b/src/com/android/launcher3/CellLayout.java
@@ -18,6 +18,7 @@
 
 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;
@@ -193,6 +194,7 @@
     private static final int INVALID_DIRECTION = -100;
 
     private final Rect mTempRect = new Rect();
+    private final RectF mTempRectF = new RectF();
 
     private static final Paint sPaint = new Paint();
 
@@ -1070,11 +1072,16 @@
         // Apply local extracted color if the DragView is an AppWidgetHostViewDrawable.
         View view = dragObject.dragView.getContentView();
         if (view instanceof LauncherAppWidgetHostView) {
-            Workspace workspace =
-                    Launcher.getLauncher(dragObject.dragView.getContext()).getWorkspace();
+            Launcher launcher = Launcher.getLauncher(dragObject.dragView.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(), workspace, mTempRect, false,
+                    mTempRectF);
+            Utilities.setRect(mTempRectF, mTempRect);
             ((LauncherAppWidgetHostView) view).handleDrag(mTempRect, pageId);
         }
     }
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index be5463e..bff3cab 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -33,6 +33,7 @@
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
 import android.content.pm.ShortcutInfo;
+import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.database.ContentObserver;
 import android.graphics.Bitmap;
@@ -84,6 +85,7 @@
 import com.android.launcher3.shortcuts.ShortcutRequest;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.PackageManagerHelper;
+import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.PendingAddShortcutInfo;
 
 import java.lang.reflect.Method;
@@ -104,6 +106,8 @@
     private static final Pattern sTrimPattern =
             Pattern.compile("^[\\s|\\p{javaSpaceChar}]*(.*)[\\s|\\p{javaSpaceChar}]*$");
 
+    private static final float[] sTmpFloatArray = new float[4];
+
     private static final int[] sLoc0 = new int[2];
     private static final int[] sLoc1 = new int[2];
     private static final Matrix sMatrix = new Matrix();
@@ -133,6 +137,15 @@
             Build.TYPE.toLowerCase(Locale.ROOT).contains("debug") ||
             Build.TYPE.toLowerCase(Locale.ROOT).equals("eng");
 
+    /**
+     * Returns true if theme is dark.
+     */
+    public static boolean isDarkTheme(Context context) {
+        Configuration configuration = context.getResources().getConfiguration();
+        int nightMode = configuration.uiMode & Configuration.UI_MODE_NIGHT_MASK;
+        return nightMode == Configuration.UI_MODE_NIGHT_YES;
+    }
+
     public static boolean isDevelopersOptionsEnabled(Context context) {
         return Settings.Global.getInt(context.getApplicationContext().getContentResolver(),
                         Settings.Global.DEVELOPMENT_SETTINGS_ENABLED, 0) != 0;
@@ -219,6 +232,33 @@
     }
 
     /**
+     * Returns bounds for a child view of DragLayer, in drag layer coordinates.
+     *
+     * see {@link com.android.launcher3.dragndrop.DragLayer}.
+     *
+     * @param viewBounds Bounds of the view wanted in drag layer coordinates, relative to the view
+     *                   itself. eg. (0, 0, view.getWidth, view.getHeight)
+     * @param ignoreTransform If true, view transform is ignored
+     * @param outRect The out rect where we return the bounds of {@param view} in drag layer coords.
+     */
+    public static void getBoundsForViewInDragLayer(BaseDragLayer dragLayer, View view,
+            Rect viewBounds, boolean ignoreTransform, RectF outRect) {
+        float[] points = sTmpFloatArray;
+        points[0] = viewBounds.left;
+        points[1] = viewBounds.top;
+        points[2] = viewBounds.right;
+        points[3] = viewBounds.bottom;
+
+        Utilities.getDescendantCoordRelativeToAncestor(view, dragLayer, points,
+                false, ignoreTransform);
+        outRect.set(
+                Math.min(points[0], points[2]),
+                Math.min(points[1], points[3]),
+                Math.max(points[0], points[2]),
+                Math.max(points[1], points[3]));
+    }
+
+    /**
      * Inverse of {@link #getDescendantCoordRelativeToAncestor(View, View, float[], boolean)}.
      */
     public static void mapCoordInSelfToDescendant(View descendant, View root, float[] coord) {
@@ -273,6 +313,16 @@
         return new int[] {sLoc1[0] - sLoc0[0], sLoc1[1] - sLoc0[1]};
     }
 
+    /**
+     * Helper method to set rectOut with rectFSrc.
+     */
+    public static void setRect(RectF rectFSrc, Rect rectOut) {
+        rectOut.left = (int) rectFSrc.left;
+        rectOut.top = (int) rectFSrc.top;
+        rectOut.right = (int) rectFSrc.right;
+        rectOut.bottom = (int) rectFSrc.bottom;
+    }
+
     public static void scaleRectFAboutCenter(RectF r, float scale) {
         if (scale != 1.0f) {
             float cx = r.centerX();
diff --git a/src/com/android/launcher3/graphics/PreloadIconDrawable.java b/src/com/android/launcher3/graphics/PreloadIconDrawable.java
index 13ae866..e45b8f7 100644
--- a/src/com/android/launcher3/graphics/PreloadIconDrawable.java
+++ b/src/com/android/launcher3/graphics/PreloadIconDrawable.java
@@ -21,7 +21,6 @@
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ObjectAnimator;
 import android.content.Context;
-import android.content.res.Configuration;
 import android.graphics.Bitmap;
 import android.graphics.Canvas;
 import android.graphics.Matrix;
@@ -34,6 +33,7 @@
 import android.util.SparseArray;
 import android.view.ContextThemeWrapper;
 
+import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.icons.GraphicsUtils;
@@ -119,9 +119,7 @@
                 info,
                 IconPalette.getPreloadProgressColor(context, info.bitmap.color),
                 getPreloadColors(context),
-            (context.getResources().getConfiguration().uiMode
-                    & Configuration.UI_MODE_NIGHT_MASK
-                    & Configuration.UI_MODE_NIGHT_YES) != 0) /* isDarkMode */;
+                Utilities.isDarkTheme(context));
     }
 
     public PreloadIconDrawable(
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index c19dfe9..5684c5a 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -16,9 +16,9 @@
 
 package com.android.launcher3.popup;
 
-import static com.android.launcher3.AbstractFloatingView.getTopOpenViewWithType;
 import static com.android.launcher3.anim.Interpolators.ACCEL_DEACCEL;
 import static com.android.launcher3.popup.PopupPopulator.MAX_SHORTCUTS;
+import static com.android.launcher3.util.ColorExtractionUtils.getColorExtractionRect;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -31,6 +31,7 @@
 import android.content.res.Resources;
 import android.graphics.Outline;
 import android.graphics.Rect;
+import android.graphics.RectF;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.GradientDrawable;
@@ -41,17 +42,21 @@
 import android.view.View;
 import android.view.ViewGroup;
 import android.view.ViewOutlineProvider;
+import android.view.ViewTreeObserver;
 import android.widget.FrameLayout;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.InsettableFrameLayout;
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAnimUtils;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.Workspace;
 import com.android.launcher3.anim.RevealOutlineAnimation;
 import com.android.launcher3.anim.RoundedRectRevealOutlineProvider;
 import com.android.launcher3.dragndrop.DragLayer;
@@ -59,9 +64,11 @@
 import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.BaseDragLayer;
+import com.android.launcher3.widget.LocalColorExtractor;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 
 /**
  * A container for shortcuts to deep links and notifications associated with an app.
@@ -73,6 +80,9 @@
 
     // +1 for system shortcut view
     private static final int MAX_NUM_CHILDREN = MAX_SHORTCUTS + 1;
+    // Index used to get background color when using local wallpaper color extraction,
+    private static final int LIGHT_COLOR_EXTRACTION_INDEX = android.R.color.system_accent2_50;
+    private static final int DARK_COLOR_EXTRACTION_INDEX = android.R.color.system_accent2_800;
 
     private final Rect mTempRect = new Rect();
 
@@ -104,8 +114,14 @@
 
     private Runnable mOnCloseCallback = () -> { };
 
+    // The rect string of the view that the arrow is attached to, in screen reference frame.
+    private String mArrowColorRectString;
     private int mArrowColor;
     private final int[] mColors;
+    private final HashMap<String, View> mViewForRect = new HashMap<>();
+
+    private final int mColorExtractionIndex;
+    @Nullable private LocalColorExtractor mColorExtractor;
 
     public ArrowPopup(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
@@ -113,7 +129,9 @@
         mOutlineRadius = Themes.getDialogCornerRadius(context);
         mLauncher = BaseDraggingActivity.fromContext(context);
         mIsRtl = Utilities.isRtl(getResources());
-
+        mColorExtractionIndex = Utilities.isDarkTheme(context)
+                ? DARK_COLOR_EXTRACTION_INDEX
+                : LIGHT_COLOR_EXTRACTION_INDEX;
         setClipToOutline(true);
         setOutlineProvider(new ViewOutlineProvider() {
             @Override
@@ -158,6 +176,10 @@
                 mColors[i] =
                         (int) argb.evaluate((i + 1) * step, primaryColor, secondaryColor);
             }
+
+            if (Utilities.ATLEAST_S) {
+                setupColorExtraction();
+            }
         }
     }
 
@@ -342,6 +364,15 @@
             // so we centered it instead. In that case we don't want to showDefaultOptions the arrow.
             mArrow.setVisibility(INVISIBLE);
         } else {
+            updateArrowColor();
+        }
+
+        mArrow.setPivotX(mArrowWidth / 2.0f);
+        mArrow.setPivotY(mIsAboveIcon ? mArrowHeight : 0);
+    }
+
+    private void updateArrowColor() {
+        if (!Gravity.isVertical(mGravity)) {
             mArrow.setBackground(new RoundedArrowDrawable(
                     mArrowWidth, mArrowHeight, mArrowPointRadius,
                     mOutlineRadius, getMeasuredWidth(), getMeasuredHeight(),
@@ -350,9 +381,6 @@
                     mArrowColor));
             mArrow.setElevation(getElevation());
         }
-
-        mArrow.setPivotX(mArrowWidth / 2.0f);
-        mArrow.setPivotY(mIsAboveIcon ? mArrowHeight : 0);
     }
 
     /**
@@ -671,6 +699,12 @@
         getPopupContainer().removeView(this);
         getPopupContainer().removeView(mArrow);
         mOnCloseCallback.run();
+        mArrowColorRectString = null;
+        mViewForRect.clear();
+        if (mColorExtractor != null) {
+            mColorExtractor.removeLocations();
+            mColorExtractor.setListener(null);
+        }
     }
 
     /**
@@ -680,6 +714,68 @@
         mOnCloseCallback = callback;
     }
 
+    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);
+            if (v != null) {
+                int newColor = extractedColors.get(mColorExtractionIndex);
+                setChildColor(v, newColor);
+                if (rectString.equals(mArrowColorRectString)) {
+                    mArrowColor = newColor;
+                    updateArrowColor();
+                }
+            }
+        });
+
+        getViewTreeObserver().addOnPreDrawListener(new ViewTreeObserver.OnPreDrawListener() {
+            @Override
+            public boolean onPreDraw() {
+                getViewTreeObserver().removeOnPreDrawListener(this);
+
+                ArrayList<RectF> locations = new ArrayList<>();
+                Rect r = new Rect();
+
+                int count = getChildCount();
+                int numVisibleChild = 0;
+                for (int i = 0; i < count; i++) {
+                    View view = getChildAt(i);
+                    if (view.getVisibility() == VISIBLE) {
+                        RectF rf = new RectF();
+                        getColorExtractionRect(Launcher.getLauncher(getContext()),
+                                workspace.getCurrentPage(), view, rf);
+                        if (rf.isEmpty()) {
+                            numVisibleChild++;
+                            continue;
+                        }
+
+                        locations.add(rf);
+                        String rectString = rf.toShortString();
+                        mViewForRect.put(rectString, view);
+
+                        // Arrow color matches the first child or the last child.
+                        if (!mIsAboveIcon && numVisibleChild == 0) {
+                            mArrowColorRectString = rectString;
+                        } else if (mIsAboveIcon) {
+                            mArrowColorRectString = rectString;
+                        }
+
+                        numVisibleChild++;
+                    }
+                }
+
+                mColorExtractor.addLocation(locations);
+                return false;
+            }
+        });
+    }
+
     protected BaseDragLayer getPopupContainer() {
         return mLauncher.getDragLayer();
     }
diff --git a/src/com/android/launcher3/util/ColorExtractionUtils.java b/src/com/android/launcher3/util/ColorExtractionUtils.java
new file mode 100644
index 0000000..b377ded
--- /dev/null
+++ b/src/com/android/launcher3/util/ColorExtractionUtils.java
@@ -0,0 +1,110 @@
+/*
+ * 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.util;
+
+import android.content.res.Resources;
+import android.graphics.Rect;
+import android.graphics.RectF;
+import android.view.View;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.Launcher;
+import com.android.launcher3.Utilities;
+
+/**
+ * Utility class used to map launcher views to wallpaper rect.
+ */
+public class ColorExtractionUtils {
+
+    public static final String TAG = "ColorExtractionUtils";
+
+    private static final Rect sTempRect = new Rect();
+    private static final RectF sTempRectF = new RectF();
+
+    /**
+     * 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 static void getColorExtractionRect(Launcher launcher, int pageId, View v,
+            RectF colorExtractionRectOut) {
+        Rect viewRect = sTempRect;
+        viewRect.set(0, 0, v.getWidth(), v.getHeight());
+        Utilities.getBoundsForViewInDragLayer(launcher.getDragLayer(), v, viewRect, false,
+                sTempRectF);
+        Utilities.setRect(sTempRectF, viewRect);
+        getColorExtractionRect(launcher, pageId, viewRect, colorExtractionRectOut);
+    }
+
+    /**
+     * 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 static void getColorExtractionRect(Launcher launcher, int pageId, Rect rectInDragLayer,
+            RectF colorExtractionRectOut) {
+        // If the view hasn't been measured and laid out, we cannot do this.
+        if (rectInDragLayer.isEmpty()) {
+            colorExtractionRectOut.setEmpty();
+            return;
+        }
+
+        Resources res = launcher.getResources();
+        DeviceProfile dp = launcher.getDeviceProfile().inv.getDeviceProfile(launcher);
+        float screenWidth = dp.widthPx;
+        float screenHeight = dp.heightPx;
+        int numScreens = launcher.getWorkspace().getNumPagesForWallpaperParallax();
+        pageId = Utilities.isRtl(res) ? numScreens - pageId - 1 : pageId;
+        float relativeScreenWidth = 1f / numScreens;
+
+        int[] dragLayerBounds = new int[2];
+        launcher.getDragLayer().getLocationOnScreen(dragLayerBounds);
+        // Translate from drag layer coordinates to screen coordinates.
+        int screenLeft = rectInDragLayer.left + dragLayerBounds[0];
+        int screenTop = rectInDragLayer.top + dragLayerBounds[1];
+        int screenRight = rectInDragLayer.right + dragLayerBounds[0];
+        int screenBottom = rectInDragLayer.bottom + dragLayerBounds[1];
+
+        // This is the position of the view 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 view relative to the screen. For the vertical axis, this is
+        // simply the location of the view relative to the screen.
+        // Translate from drag layer coordinates to screen coordinates
+        colorExtractionRectOut.left = (screenLeft / screenWidth + pageId) * relativeScreenWidth;
+        colorExtractionRectOut.right = (screenRight / screenWidth + pageId) * relativeScreenWidth;
+        colorExtractionRectOut.top = screenTop / screenHeight;
+        colorExtractionRectOut.bottom = screenBottom / screenHeight;
+
+        if (colorExtractionRectOut.left < 0
+                || colorExtractionRectOut.right > 1
+                || colorExtractionRectOut.top < 0
+                || colorExtractionRectOut.bottom > 1) {
+            colorExtractionRectOut.setEmpty();
+        }
+    }
+}
diff --git a/src/com/android/launcher3/util/Themes.java b/src/com/android/launcher3/util/Themes.java
index 99942aa..06cac08 100644
--- a/src/com/android/launcher3/util/Themes.java
+++ b/src/com/android/launcher3/util/Themes.java
@@ -22,7 +22,6 @@
 import android.app.WallpaperColors;
 import android.app.WallpaperManager;
 import android.content.Context;
-import android.content.res.Configuration;
 import android.content.res.TypedArray;
 import android.graphics.Color;
 import android.graphics.ColorMatrix;
@@ -54,16 +53,12 @@
     }
 
     public static int getActivityThemeRes(Context context, int wallpaperColorHints) {
-        Configuration configuration = context.getResources().getConfiguration();
-        int nightMode = configuration.uiMode & Configuration.UI_MODE_NIGHT_MASK;
-        boolean darkTheme = nightMode == Configuration.UI_MODE_NIGHT_YES;
-
         boolean supportsDarkText = Utilities.ATLEAST_S
                 && (wallpaperColorHints & HINT_SUPPORTS_DARK_TEXT) != 0;
         boolean isMainColorDark = Utilities.ATLEAST_S
                 && (wallpaperColorHints & HINT_SUPPORTS_DARK_THEME) != 0;
 
-        if (darkTheme) {
+        if (Utilities.isDarkTheme(context)) {
             return supportsDarkText ? R.style.AppTheme_Dark_DarkText
                     : isMainColorDark ? R.style.AppTheme_Dark_DarkMainColor : R.style.AppTheme_Dark;
         } else {
diff --git a/src/com/android/launcher3/views/FloatingIconView.java b/src/com/android/launcher3/views/FloatingIconView.java
index d49320b..1ddd1ba 100644
--- a/src/com/android/launcher3/views/FloatingIconView.java
+++ b/src/com/android/launcher3/views/FloatingIconView.java
@@ -233,15 +233,8 @@
             outViewBounds.set(0, 0, v.getWidth(), v.getHeight());
         }
 
-        float[] points = new float[] {outViewBounds.left, outViewBounds.top, outViewBounds.right,
-                outViewBounds.bottom};
-        Utilities.getDescendantCoordRelativeToAncestor(v, launcher.getDragLayer(), points,
-                false, ignoreTransform);
-        outRect.set(
-                Math.min(points[0], points[2]),
-                Math.min(points[1], points[3]),
-                Math.max(points[0], points[2]),
-                Math.max(points[1], points[3]));
+        Utilities.getBoundsForViewInDragLayer(launcher.getDragLayer(), v, outViewBounds,
+                ignoreTransform, outRect);
     }
 
     /**
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
index 620604a..697c453 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
@@ -16,6 +16,10 @@
 
 package com.android.launcher3.widget;
 
+import static com.android.launcher3.Utilities.getBoundsForViewInDragLayer;
+import static com.android.launcher3.Utilities.setRect;
+import static com.android.launcher3.util.ColorExtractionUtils.getColorExtractionRect;
+
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.content.res.Configuration;
@@ -25,7 +29,6 @@
 import android.graphics.RectF;
 import android.os.Handler;
 import android.os.SystemClock;
-import android.util.Log;
 import android.util.SparseBooleanArray;
 import android.util.SparseIntArray;
 import android.view.LayoutInflater;
@@ -95,11 +98,11 @@
     private RectF mLastLocationRegistered = null;
     @Nullable private AppWidgetHostViewDragListener mDragListener;
 
-    // Used to store the widget size during onLayout.
+    // 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 boolean mIsRtl;
     private final Rect mEnforcedRectangle = new Rect();
     private final float mEnforcedCornerRadius;
     private final ViewOutlineProvider mCornerRadiusEnforcementOutline = new ViewOutlineProvider() {
@@ -129,7 +132,6 @@
         if (Utilities.ATLEAST_Q && Themes.getAttrBoolean(mLauncher, R.attr.isWorkspaceDarkText)) {
             setOnLightBackground(true);
         }
-        mIsRtl = Utilities.isRtl(context.getResources());
         mColorExtractor = LocalColorExtractor.newInstance(getContext());
         mColorExtractor.setListener(this);
 
@@ -319,13 +321,12 @@
         mIsScrollable = checkScrollableRecursively(this);
 
         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);
+            getBoundsForViewInDragLayer(mLauncher.getDragLayer(), this, mCurrentWidgetSize, true,
+                    mTempRectF);
+            setRect(mTempRectF, mCurrentWidgetSize);
+            updateColorExtraction(mCurrentWidgetSize,
+                    mWorkspace.getPageIndexForScreenId(info.screenId));
         }
 
         enforceRoundedCorners();
@@ -338,8 +339,8 @@
     }
 
     /** Handles a drag event occurred on a workspace page, {@code pageId}. */
-    public void handleDrag(Rect rect, int pageId) {
-        mWidgetSizeAtDrag.set(rect);
+    public void handleDrag(Rect rectInDragLayer, int pageId) {
+        mWidgetSizeAtDrag.set(rectInDragLayer);
         updateColorExtraction(mWidgetSizeAtDrag, pageId);
     }
 
@@ -351,53 +352,14 @@
         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;
-        }
-        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;
-        View v = this;
-        while (v.getParent() instanceof View) {
-            v = (View) v.getParent();
-            if (v.getId() != R.id.launcher) {
-                break;
-            }
-            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 = screenWidth - mWorkspace.getPaddingRight() - parentViewWidth;
-        } else {
-            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");
+    /**
+     * @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) {
+        getColorExtractionRect(mLauncher, pageId, rectInDragLayer, mTempRectF);
+
+        if (mTempRectF.isEmpty()) {
             return;
         }
         if (!mTempRectF.equals(mLastLocationRegistered)) {