From cb5bb989e3e8b2d0b43f8a97aeefdfc1389eeaae Mon Sep 17 00:00:00 2001
From: y <sfufa@google.com>
Date: Wed, 25 Aug 2021 11:58:46 -0700
Subject: [PATCH] [Work] Fix duplicate work button issue

Bug: 195623679
Test: manual
Change-Id: I882c5f7a196e07fa0dab3a2f27437abd233ec782
---
 .../allapps/AllAppsContainerView.java         | 20 ++++++++++---------
 .../launcher3/allapps/WorkModeSwitch.java     | 12 +++++------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index d2c71b2afe..9c784f7a84 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -454,7 +454,6 @@ public class AllAppsContainerView extends SpringRelativeLayout implements DragSo
         mAllAppsStore.unregisterIconContainer(mAH[AdapterHolder.WORK].recyclerView);
 
         if (mUsingTabs) {
-            setupWorkToggle();
             mAH[AdapterHolder.MAIN].setup(mViewPager.getChildAt(0), mPersonalMatcher);
             mAH[AdapterHolder.WORK].setup(mViewPager.getChildAt(1), mWorkMatcher);
             mAH[AdapterHolder.WORK].recyclerView.setId(R.id.apps_list_view_work);
@@ -485,6 +484,7 @@ public class AllAppsContainerView extends SpringRelativeLayout implements DragSo
     }
 
     private void setupWorkToggle() {
+        removeWorkToggle();
         if (Utilities.ATLEAST_P) {
             mWorkModeSwitch = (WorkModeSwitch) mLauncher.getLayoutInflater().inflate(
                     R.layout.work_mode_fab, this, false);
@@ -497,6 +497,14 @@ public class AllAppsContainerView extends SpringRelativeLayout implements DragSo
         }
     }
 
+    private void removeWorkToggle() {
+        if (mWorkModeSwitch == null) return;
+        if (mWorkModeSwitch.getParent() == this) {
+            this.removeView(mWorkModeSwitch);
+        }
+        mWorkModeSwitch = null;
+    }
+
     private void replaceRVContainer(boolean showTabs) {
         for (int i = 0; i < mAH.length; i++) {
             if (mAH[i].recyclerView != null) {
@@ -517,8 +525,10 @@ public class AllAppsContainerView extends SpringRelativeLayout implements DragSo
             mViewPager = (AllAppsPagedView) newView;
             mViewPager.initParentViews(this);
             mViewPager.getPageIndicator().setOnActivePageChangedListener(this);
+            setupWorkToggle();
         } else {
             mViewPager = null;
+            removeWorkToggle();
         }
     }
 
@@ -537,14 +547,6 @@ public class AllAppsContainerView extends SpringRelativeLayout implements DragSo
             mWorkModeSwitch.setWorkTabVisible(currentActivePage == AdapterHolder.WORK
                     && mAllAppsStore.hasModelFlag(
                     FLAG_HAS_SHORTCUT_PERMISSION | FLAG_QUIET_MODE_CHANGE_PERMISSION));
-
-            if (currentActivePage == AdapterHolder.WORK) {
-                if (mWorkModeSwitch.getParent() == null) {
-                    addView(mWorkModeSwitch);
-                }
-            } else {
-                removeView(mWorkModeSwitch);
-            }
         }
     }
 
diff --git a/src/com/android/launcher3/allapps/WorkModeSwitch.java b/src/com/android/launcher3/allapps/WorkModeSwitch.java
index a800d34758..5d3af08b46 100644
--- a/src/com/android/launcher3/allapps/WorkModeSwitch.java
+++ b/src/com/android/launcher3/allapps/WorkModeSwitch.java
@@ -51,6 +51,7 @@ public class WorkModeSwitch extends Button implements Insettable, View.OnClickLi
 
     @Nullable
     private KeyboardInsetAnimationCallback mKeyboardInsetAnimationCallback;
+    private boolean mWorkTabVisible;
 
     public WorkModeSwitch(Context context) {
         this(context, null, 0);
@@ -91,11 +92,10 @@ public class WorkModeSwitch extends Button implements Insettable, View.OnClickLi
      */
     public void setWorkTabVisible(boolean workTabVisible) {
         clearAnimation();
-        if (workTabVisible) {
+        mWorkTabVisible = workTabVisible;
+        if (workTabVisible && mWorkEnabled) {
             setEnabled(true);
-            if (mWorkEnabled) {
-                setVisibility(VISIBLE);
-            }
+            setVisibility(VISIBLE);
             setAlpha(0);
             animate().alpha(1).start();
         } else {
@@ -105,7 +105,7 @@ public class WorkModeSwitch extends Button implements Insettable, View.OnClickLi
 
     @Override
     public void onClick(View view) {
-        if (Utilities.ATLEAST_P) {
+        if (Utilities.ATLEAST_P && mWorkTabVisible) {
             setEnabled(false);
             Launcher.fromContext(getContext()).getStatsLogManager().logger().log(
                     LAUNCHER_TURN_OFF_WORK_APPS_TAP);
@@ -137,7 +137,7 @@ public class WorkModeSwitch extends Button implements Insettable, View.OnClickLi
 
     @Override
     public WindowInsets onApplyWindowInsets(WindowInsets insets) {
-        if (Utilities.ATLEAST_R) {
+        if (Utilities.ATLEAST_R && mWorkTabVisible) {
             setTranslationY(0);
             if (insets.isVisible(WindowInsets.Type.ime())) {
                 Insets keyboardInsets = insets.getInsets(WindowInsets.Type.ime());