diff --git a/src/com/android/settings/activityembedding/ActivityEmbeddingRulesController.java b/src/com/android/settings/activityembedding/ActivityEmbeddingRulesController.java
index fa61994..8484324 100644
--- a/src/com/android/settings/activityembedding/ActivityEmbeddingRulesController.java
+++ b/src/com/android/settings/activityembedding/ActivityEmbeddingRulesController.java
@@ -23,7 +23,6 @@
 import android.util.LayoutDirection;
 import android.util.Log;
 
-import androidx.annotation.NonNull;
 import androidx.window.embedding.ActivityFilter;
 import androidx.window.embedding.ActivityRule;
 import androidx.window.embedding.SplitController;
@@ -114,7 +113,7 @@
 
         registerTwoPanePairRule(
                 context,
-                getComponentName(context, Settings.class),
+                new ComponentName(context, Settings.class),
                 secondaryComponent,
                 secondaryIntentAction,
                 finishPrimaryWithSecondary ? SplitRule.FINISH_ADJACENT : SplitRule.FINISH_NEVER,
@@ -123,7 +122,7 @@
 
         registerTwoPanePairRule(
                 context,
-                getComponentName(context, SettingsHomepageActivity.class),
+                new ComponentName(context, SettingsHomepageActivity.class),
                 secondaryComponent,
                 secondaryIntentAction,
                 finishPrimaryWithSecondary ? SplitRule.FINISH_ADJACENT : SplitRule.FINISH_NEVER,
@@ -143,7 +142,7 @@
 
         registerTwoPanePairRule(
                 context,
-                getComponentName(context, SliceDeepLinkHomepageActivity.class),
+                new ComponentName(context, SliceDeepLinkHomepageActivity.class),
                 secondaryComponent,
                 secondaryIntentAction,
                 finishPrimaryWithSecondary ? SplitRule.FINISH_ALWAYS : SplitRule.FINISH_NEVER,
@@ -179,7 +178,7 @@
 
         registerTwoPanePairRuleForSettingsHome(
                 context,
-                getComponentName(context, SubSettings.class),
+                new ComponentName(context, SubSettings.class),
                 null /* secondaryIntentAction */,
                 clearTop);
     }
@@ -191,8 +190,7 @@
         addActivityFilter(activityFilters, SliceDeepLinkHomepageActivity.class);
         addActivityFilter(activityFilters, Settings.class);
 
-        final Intent intent = new Intent();
-        intent.setComponent(getComponentName(Settings.NetworkDashboardActivity.class));
+        final Intent intent = new Intent(mContext, Settings.NetworkDashboardActivity.class);
         final SplitPlaceholderRule placeholderRule = new SplitPlaceholderRule(
                 activityFilters,
                 intent,
@@ -215,23 +213,7 @@
 
     private void addActivityFilter(Set<ActivityFilter> activityFilters,
             Class<? extends Activity> activityClass) {
-        activityFilters.add(new ActivityFilter(getComponentName(activityClass),
+        activityFilters.add(new ActivityFilter(new ComponentName(mContext, activityClass),
                 null /* intentAction */));
     }
-
-    private void addActivityFilter(Set<ActivityFilter> activityFilters,
-            ComponentName componentName) {
-        activityFilters.add(new ActivityFilter(componentName, null /* intentAction */));
-    }
-
-    @NonNull
-    private ComponentName getComponentName(Class<? extends Activity> activityClass) {
-        return getComponentName(mContext, activityClass);
-    }
-
-    @NonNull
-    private static ComponentName getComponentName(Context context,
-            Class<? extends Activity> activityClass) {
-        return new ComponentName(context.getPackageName(), activityClass.getName());
-    }
 }
diff --git a/src/com/android/settings/dashboard/DashboardFeatureProviderImpl.java b/src/com/android/settings/dashboard/DashboardFeatureProviderImpl.java
index 3d1381d..69fdc9e 100644
--- a/src/com/android/settings/dashboard/DashboardFeatureProviderImpl.java
+++ b/src/com/android/settings/dashboard/DashboardFeatureProviderImpl.java
@@ -36,7 +36,6 @@
 import android.app.settings.SettingsEnums;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.DialogInterface.OnCancelListener;
 import android.content.IContentProvider;
 import android.content.Intent;
 import android.content.pm.PackageManager;
@@ -63,6 +62,7 @@
 import com.android.settings.activityembedding.ActivityEmbeddingRulesController;
 import com.android.settings.activityembedding.ActivityEmbeddingUtils;
 import com.android.settings.dashboard.profileselector.ProfileSelectDialog;
+import com.android.settings.homepage.TopLevelHighlightMixin;
 import com.android.settings.homepage.TopLevelSettings;
 import com.android.settings.overlay.FeatureFactory;
 import com.android.settings.widget.PrimarySwitchPreference;
@@ -170,27 +170,23 @@
                 if (action != null) {
                     intent.setAction(action);
                 }
+                // Register the rule for injected apps.
+                ActivityEmbeddingRulesController.registerTwoPanePairRuleForSettingsHome(
+                        mContext,
+                        new ComponentName(tile.getPackageName(), tile.getComponentName()),
+                        action,
+                        true /* clearTop */);
                 pref.setOnPreferenceClickListener(preference -> {
-                    OnCancelListener listener = null;
+                    TopLevelHighlightMixin highlightMixin = null;
                     if (fragment instanceof TopLevelSettings
                             && ActivityEmbeddingUtils.isEmbeddingActivityEnabled(mContext)) {
-                        // Register the rule for injected apps.
-                        ActivityEmbeddingRulesController.registerTwoPanePairRuleForSettingsHome(
-                                mContext,
-                                new ComponentName(tile.getPackageName(), tile.getComponentName()),
-                                null /* secondaryIntentAction */,
-                                true /* clearTop */);
-
-                        // Highlight preference ui.
+                        // Highlight the preference whenever it's clicked
                         final TopLevelSettings topLevelSettings = (TopLevelSettings) fragment;
-                        // Highlight the tile immediately whenever it's clicked
                         topLevelSettings.setHighlightPreferenceKey(key);
-                        // If the tile allows users to select profile, the pop-op dialog may be
-                        // cancelled and then the previous highlight entry should be restored.
-                        listener = dialog -> topLevelSettings.restorePreviousHighlight();
+                        highlightMixin = topLevelSettings.getHighlightMixin();
                     }
                     launchIntentOrSelectProfile(activity, tile, intent, sourceMetricsCategory,
-                            listener);
+                            highlightMixin);
                     return true;
                 });
             }
@@ -223,7 +219,7 @@
                         SettingsEnums.DASHBOARD_SUMMARY)
                 .addFlags(Intent.FLAG_ACTIVITY_CLEAR_TASK);
         launchIntentOrSelectProfile(activity, tile, intent, SettingsEnums.DASHBOARD_SUMMARY,
-                /* listener= */ null);
+                /* highlightMixin= */ null);
     }
 
     private DynamicDataObserver createDynamicDataObserver(String method, Uri uri, Preference pref) {
@@ -438,7 +434,7 @@
     }
 
     private void launchIntentOrSelectProfile(FragmentActivity activity, Tile tile, Intent intent,
-            int sourceMetricCategory, OnCancelListener listener) {
+            int sourceMetricCategory, TopLevelHighlightMixin highlightMixin) {
         if (!isIntentResolvable(intent)) {
             Log.w(TAG, "Cannot resolve intent, skipping. " + intent);
             return;
@@ -469,7 +465,9 @@
             }
 
             ProfileSelectDialog.show(activity.getSupportFragmentManager(), tile,
-                    sourceMetricCategory, listener);
+                    sourceMetricCategory, /* onShowListener= */ highlightMixin,
+                    /* onDismissListener= */ highlightMixin,
+                    /* onCancelListener= */ highlightMixin);
         }
     }
 
diff --git a/src/com/android/settings/dashboard/profileselector/ProfileSelectDialog.java b/src/com/android/settings/dashboard/profileselector/ProfileSelectDialog.java
index b05f23b..e0ba378 100644
--- a/src/com/android/settings/dashboard/profileselector/ProfileSelectDialog.java
+++ b/src/com/android/settings/dashboard/profileselector/ProfileSelectDialog.java
@@ -21,6 +21,8 @@
 import android.content.DialogInterface;
 import android.content.DialogInterface.OnCancelListener;
 import android.content.DialogInterface.OnClickListener;
+import android.content.DialogInterface.OnDismissListener;
+import android.content.DialogInterface.OnShowListener;
 import android.content.Intent;
 import android.os.Bundle;
 import android.os.UserHandle;
@@ -45,23 +47,30 @@
 
     private int mSourceMetricCategory;
     private Tile mSelectedTile;
+    private OnShowListener mOnShowListener;
     private OnCancelListener mOnCancelListener;
+    private OnDismissListener mOnDismissListener;
 
     /**
      * Display the profile select dialog, adding the fragment to the given FragmentManager.
      * @param manager The FragmentManager this fragment will be added to.
      * @param tile The tile for this fragment.
      * @param sourceMetricCategory The source metric category.
-     * @param listener The listener listens to the dialog cancelling event.
+     * @param onShowListener The listener listens to the dialog showing event.
+     * @param onDismissListener The listener listens to the dialog dismissing event.
+     * @param onCancelListener The listener listens to the dialog cancelling event.
      */
     public static void show(FragmentManager manager, Tile tile, int sourceMetricCategory,
-            OnCancelListener listener) {
+            OnShowListener onShowListener, OnDismissListener onDismissListener,
+            OnCancelListener onCancelListener) {
         final ProfileSelectDialog dialog = new ProfileSelectDialog();
         final Bundle args = new Bundle();
         args.putParcelable(ARG_SELECTED_TILE, tile);
         args.putInt(ARG_SOURCE_METRIC_CATEGORY, sourceMetricCategory);
         dialog.setArguments(args);
-        dialog.mOnCancelListener = listener;
+        dialog.mOnShowListener = onShowListener;
+        dialog.mOnDismissListener = onDismissListener;
+        dialog.mOnCancelListener = onCancelListener;
         dialog.show(manager, "select_profile");
     }
 
@@ -97,12 +106,30 @@
     }
 
     @Override
+    public void onStart() {
+        super.onStart();
+        // The fragment shows the dialog within onStart()
+        if (mOnShowListener != null) {
+            mOnShowListener.onShow(getDialog());
+        }
+    }
+
+    @Override
     public void onCancel(DialogInterface dialog) {
+        super.onCancel(dialog);
         if (mOnCancelListener != null) {
             mOnCancelListener.onCancel(dialog);
         }
     }
 
+    @Override
+    public void onDismiss(DialogInterface dialog) {
+        super.onDismiss(dialog);
+        if (mOnDismissListener != null) {
+            mOnDismissListener.onDismiss(dialog);
+        }
+    }
+
     public static void updateUserHandlesIfNeeded(Context context, Tile tile) {
         final List<UserHandle> userHandles = tile.userHandle;
         if (tile.userHandle == null || tile.userHandle.size() <= 1) {
diff --git a/src/com/android/settings/display/TopLevelWallpaperPreferenceController.java b/src/com/android/settings/display/TopLevelWallpaperPreferenceController.java
index 7640d08..0136eac 100644
--- a/src/com/android/settings/display/TopLevelWallpaperPreferenceController.java
+++ b/src/com/android/settings/display/TopLevelWallpaperPreferenceController.java
@@ -63,6 +63,11 @@
         super.displayPreference(screen);
         Preference preference = screen.findPreference(getPreferenceKey());
         preference.setTitle(getTitle());
+        ActivityEmbeddingRulesController.registerTwoPanePairRuleForSettingsHome(
+                mContext,
+                getComponentName(),
+                null /* secondaryIntentAction */,
+                true /* clearTop */);
     }
 
     public String getTitle() {
@@ -103,11 +108,6 @@
                     mContext)) {
                 intent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK);
             }
-            ActivityEmbeddingRulesController.registerTwoPanePairRuleForSettingsHome(
-                    mContext,
-                    intent.getComponent(),
-                    null /* secondaryIntentAction */,
-                    true /* clearTop */);
             preference.getContext().startActivity(intent);
             return true;
         }
diff --git a/src/com/android/settings/homepage/SettingsHomepageActivity.java b/src/com/android/settings/homepage/SettingsHomepageActivity.java
index 31cc7aa..e25e1f5 100644
--- a/src/com/android/settings/homepage/SettingsHomepageActivity.java
+++ b/src/com/android/settings/homepage/SettingsHomepageActivity.java
@@ -76,7 +76,7 @@
     public static final String EXTRA_SETTINGS_LARGE_SCREEN_DEEP_LINK_INTENT_DATA =
             "settings_large_screen_deep_link_intent_data";
 
-    private static final int DEFAULT_HIGHLIGHT_MENU_KEY = R.string.menu_key_network;
+    static final int DEFAULT_HIGHLIGHT_MENU_KEY = R.string.menu_key_network;
     private static final long HOMEPAGE_LOADING_TIMEOUT_MS = 300;
 
     private TopLevelSettings mMainFragment;
@@ -94,6 +94,10 @@
         void onHomepageLoaded();
     }
 
+    private interface FragmentBuilder<T extends Fragment>  {
+        T build();
+    }
+
     /**
      * Try to add a {@link HomepageLoadedListener}. If homepage is already loaded, the listener
      * will not be notified.
@@ -168,13 +172,15 @@
             initAvatarView();
             showSuggestionFragment();
             if (FeatureFlagUtils.isEnabled(this, FeatureFlags.CONTEXTUAL_HOME)) {
-                showFragment(new ContextualCardsFragment(), R.id.contextual_cards_content);
+                showFragment(() -> new ContextualCardsFragment(), R.id.contextual_cards_content);
             }
         }
-        mMainFragment = new TopLevelSettings();
-        mMainFragment.getArguments().putString(SettingsActivity.EXTRA_FRAGMENT_ARG_KEY,
-                getHighlightMenuKey());
-        showFragment(mMainFragment, R.id.main_content);
+        mMainFragment = showFragment(() -> {
+            final TopLevelSettings fragment = new TopLevelSettings();
+            fragment.getArguments().putString(SettingsActivity.EXTRA_FRAGMENT_ARG_KEY,
+                    getHighlightMenuKey());
+            return fragment;
+        }, R.id.main_content);
 
         ((FrameLayout) findViewById(R.id.main_content))
                 .getLayoutTransition().enableTransitionType(LayoutTransition.CHANGING);
@@ -260,9 +266,9 @@
     }
 
     private void showSuggestionFragment() {
-        final Class<? extends Fragment> fragment = FeatureFactory.getFactory(this)
+        final Class<? extends Fragment> fragmentClass = FeatureFactory.getFactory(this)
                 .getSuggestionFeatureProvider(this).getContextualSuggestionFragment();
-        if (fragment == null) {
+        if (fragmentClass == null) {
             return;
         }
 
@@ -274,28 +280,33 @@
         // Schedule a timer to show the homepage and hide the suggestion on timeout.
         mHomepageView.postDelayed(() -> showHomepageWithSuggestion(false),
                 HOMEPAGE_LOADING_TIMEOUT_MS);
-        try {
-            showFragment(fragment.getConstructor().newInstance(), R.id.suggestion_content);
-            if (mIsEmbeddingActivityEnabled) {
-                showFragment(fragment.getConstructor().newInstance(),
-                        R.id.two_pane_suggestion_content);
+        final FragmentBuilder<?> fragmentBuilder = () -> {
+            try {
+                return fragmentClass.getConstructor().newInstance();
+            } catch (Exception e) {
+                Log.w(TAG, "Cannot show fragment", e);
             }
-        } catch (Exception e) {
-            Log.w(TAG, "Cannot show fragment", e);
+            return null;
+        };
+        showFragment(fragmentBuilder, R.id.suggestion_content);
+        if (mIsEmbeddingActivityEnabled) {
+            showFragment(fragmentBuilder, R.id.two_pane_suggestion_content);
         }
     }
 
-    private void showFragment(Fragment fragment, int id) {
+    private <T extends Fragment> T showFragment(FragmentBuilder<T> fragmentBuilder, int id) {
         final FragmentManager fragmentManager = getSupportFragmentManager();
         final FragmentTransaction fragmentTransaction = fragmentManager.beginTransaction();
-        final Fragment showFragment = fragmentManager.findFragmentById(id);
+        T showFragment = (T) fragmentManager.findFragmentById(id);
 
         if (showFragment == null) {
-            fragmentTransaction.add(id, fragment);
+            showFragment = fragmentBuilder.build();
+            fragmentTransaction.add(id, showFragment);
         } else {
             fragmentTransaction.show(showFragment);
         }
         fragmentTransaction.commit();
+        return showFragment;
     }
 
     private void launchDeepLinkIntentToRight() {
@@ -363,14 +374,14 @@
                 targetIntent.getAction(),
                 SplitRule.FINISH_ALWAYS,
                 SplitRule.FINISH_ALWAYS,
-                true /* clearTop*/);
+                true /* clearTop */);
         ActivityEmbeddingRulesController.registerTwoPanePairRule(this,
-                new ComponentName(Settings.class.getPackageName(), Settings.class.getName()),
+                new ComponentName(getApplicationContext(), Settings.class),
                 targetComponentName,
                 targetIntent.getAction(),
                 SplitRule.FINISH_ALWAYS,
                 SplitRule.FINISH_ALWAYS,
-                true /* clearTop*/);
+                true /* clearTop */);
         startActivity(targetIntent);
     }
 
diff --git a/src/com/android/settings/homepage/TopLevelHighlightMixin.java b/src/com/android/settings/homepage/TopLevelHighlightMixin.java
new file mode 100644
index 0000000..ebfe7f2
--- /dev/null
+++ b/src/com/android/settings/homepage/TopLevelHighlightMixin.java
@@ -0,0 +1,189 @@
+/*
+ * 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.settings.homepage;
+
+import android.content.DialogInterface;
+import android.os.Bundle;
+import android.os.Parcel;
+import android.os.Parcelable;
+import android.text.TextUtils;
+import android.util.Log;
+
+import androidx.fragment.app.FragmentActivity;
+import androidx.preference.PreferenceScreen;
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.settings.SettingsActivity;
+import com.android.settings.widget.HighlightableTopLevelPreferenceAdapter;
+
+/** A highlight mixin for the top level settings fragment. */
+public class TopLevelHighlightMixin implements Parcelable, DialogInterface.OnShowListener,
+        DialogInterface.OnCancelListener, DialogInterface.OnDismissListener {
+
+    private static final String TAG = "TopLevelHighlightMixin";
+
+    private String mCurrentKey;
+    // Stores the previous key for the profile select dialog cancel event
+    private String mPreviousKey;
+    // Stores the key hidden for the search page presence
+    private String mHiddenKey;
+    private DialogInterface mDialog;
+    private HighlightableTopLevelPreferenceAdapter mTopLevelAdapter;
+
+    public TopLevelHighlightMixin() {
+    }
+
+    public TopLevelHighlightMixin(Parcel source) {
+        mCurrentKey = source.readString();
+        mPreviousKey = source.readString();
+        mHiddenKey = source.readString();
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+        dest.writeString(mCurrentKey);
+        dest.writeString(mPreviousKey);
+        dest.writeString(mHiddenKey);
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    public static final Creator<TopLevelHighlightMixin> CREATOR = new Creator<>() {
+        @Override
+        public TopLevelHighlightMixin createFromParcel(Parcel source) {
+            return new TopLevelHighlightMixin(source);
+        }
+
+        @Override
+        public TopLevelHighlightMixin[] newArray(int size) {
+            return new TopLevelHighlightMixin[size];
+        }
+    };
+
+    @Override
+    public void onShow(DialogInterface dialog) {
+        mDialog = dialog;
+    }
+
+    @Override
+    public void onDismiss(DialogInterface dialog) {
+        mDialog = null;
+    }
+
+    @Override
+    public void onCancel(DialogInterface dialog) {
+        if (mTopLevelAdapter != null) {
+            mCurrentKey = mPreviousKey;
+            mPreviousKey = null;
+            mTopLevelAdapter.highlightPreference(mCurrentKey, /* scrollNeeded= */ false);
+        }
+    }
+
+    RecyclerView.Adapter onCreateAdapter(TopLevelSettings topLevelSettings,
+            PreferenceScreen preferenceScreen) {
+        if (TextUtils.isEmpty(mCurrentKey)) {
+            mCurrentKey = getHighlightPrefKeyFromArguments(topLevelSettings.getArguments());
+        }
+
+        Log.d(TAG, "onCreateAdapter, pref key: " + mCurrentKey);
+        mTopLevelAdapter = new HighlightableTopLevelPreferenceAdapter(
+                (SettingsHomepageActivity) topLevelSettings.getActivity(), preferenceScreen,
+                topLevelSettings.getListView(), mCurrentKey);
+        return mTopLevelAdapter;
+    }
+
+    void reloadHighlightMenuKey(Bundle arguments) {
+        if (mTopLevelAdapter == null) {
+            return;
+        }
+        ensureDialogDismissed();
+
+        mCurrentKey = getHighlightPrefKeyFromArguments(arguments);
+        Log.d(TAG, "reloadHighlightMenuKey, pref key: " + mCurrentKey);
+        mTopLevelAdapter.highlightPreference(mCurrentKey, /* scrollNeeded= */ true);
+    }
+
+    void setHighlightPreferenceKey(String prefKey) {
+        if (mTopLevelAdapter != null) {
+            ensureDialogDismissed();
+            mPreviousKey = mCurrentKey;
+            mCurrentKey = prefKey;
+            mTopLevelAdapter.highlightPreference(prefKey, /* scrollNeeded= */ false);
+        }
+    }
+
+    void highlightPreferenceIfNeeded(FragmentActivity activity) {
+        if (mTopLevelAdapter != null) {
+            mTopLevelAdapter.requestHighlight();
+        }
+    }
+
+    void setMenuHighlightShowed(boolean show) {
+        if (mTopLevelAdapter == null) {
+            return;
+        }
+        ensureDialogDismissed();
+
+        if (show) {
+            mCurrentKey = mHiddenKey;
+            mHiddenKey = null;
+        } else {
+            if (mHiddenKey == null) {
+                mHiddenKey = mCurrentKey;
+            }
+            mCurrentKey = null;
+        }
+        mTopLevelAdapter.highlightPreference(mCurrentKey, /* scrollNeeded= */ show);
+    }
+
+    void setHighlightMenuKey(String menuKey, boolean scrollNeeded) {
+        if (mTopLevelAdapter == null) {
+            return;
+        }
+        ensureDialogDismissed();
+
+        final String prefKey = HighlightableMenu.lookupPreferenceKey(menuKey);
+        if (TextUtils.isEmpty(prefKey)) {
+            Log.e(TAG, "Invalid highlight menu key: " + menuKey);
+        } else {
+            Log.d(TAG, "Menu key: " + menuKey);
+            mCurrentKey = prefKey;
+            mTopLevelAdapter.highlightPreference(prefKey, scrollNeeded);
+        }
+    }
+
+    private static String getHighlightPrefKeyFromArguments(Bundle arguments) {
+        final String menuKey = arguments.getString(SettingsActivity.EXTRA_FRAGMENT_ARG_KEY);
+        final String prefKey = HighlightableMenu.lookupPreferenceKey(menuKey);
+        if (TextUtils.isEmpty(prefKey)) {
+            Log.e(TAG, "Invalid highlight menu key: " + menuKey);
+        } else {
+            Log.d(TAG, "Menu key: " + menuKey);
+        }
+        return prefKey;
+    }
+
+    private void ensureDialogDismissed() {
+        if (mDialog != null) {
+            onCancel(mDialog);
+            mDialog.dismiss();
+        }
+    }
+}
diff --git a/src/com/android/settings/homepage/TopLevelSettings.java b/src/com/android/settings/homepage/TopLevelSettings.java
index eb1a066..2bb8e34 100644
--- a/src/com/android/settings/homepage/TopLevelSettings.java
+++ b/src/com/android/settings/homepage/TopLevelSettings.java
@@ -25,7 +25,6 @@
 import android.graphics.drawable.Drawable;
 import android.os.Bundle;
 import android.text.TextUtils;
-import android.util.Log;
 
 import androidx.fragment.app.Fragment;
 import androidx.preference.Preference;
@@ -34,7 +33,6 @@
 import androidx.recyclerview.widget.RecyclerView;
 
 import com.android.settings.R;
-import com.android.settings.SettingsActivity;
 import com.android.settings.Utils;
 import com.android.settings.activityembedding.ActivityEmbeddingRulesController;
 import com.android.settings.activityembedding.ActivityEmbeddingUtils;
@@ -42,7 +40,6 @@
 import com.android.settings.dashboard.DashboardFragment;
 import com.android.settings.search.BaseSearchIndexProvider;
 import com.android.settings.support.SupportPreferenceController;
-import com.android.settings.widget.HighlightableTopLevelPreferenceAdapter;
 import com.android.settings.widget.HomepagePreference;
 import com.android.settingslib.core.instrumentation.Instrumentable;
 import com.android.settingslib.drawer.Tile;
@@ -53,13 +50,10 @@
         PreferenceFragmentCompat.OnPreferenceStartFragmentCallback {
 
     private static final String TAG = "TopLevelSettings";
-    private static final String SAVED_HIGHLIGHTED_PREF = "highlighted_pref";
-    private static final String SAVED_CACHED_PREF = "cached_pref";
+    private static final String SAVED_HIGHLIGHT_MIXIN = "highlight_mixin";
+    private static final String PREF_KEY_SUPPORT = "top_level_support";
 
-    private HighlightableTopLevelPreferenceAdapter mTopLevelAdapter;
-
-    private String mHighlightedPreferenceKey;
-    private String mCachedPreferenceKey;
+    private TopLevelHighlightMixin mHighlightMixin;
 
     public TopLevelSettings() {
         final Bundle args = new Bundle();
@@ -127,17 +121,35 @@
     @Override
     public void onCreate(Bundle icicle) {
         super.onCreate(icicle);
-        if (icicle != null) {
-            mHighlightedPreferenceKey = icicle.getString(SAVED_HIGHLIGHTED_PREF);
-            mCachedPreferenceKey = icicle.getString(SAVED_CACHED_PREF);
+        if (!ActivityEmbeddingUtils.isEmbeddingActivityEnabled(getContext())) {
+            return;
         }
+
+        if (icicle != null) {
+            mHighlightMixin = icicle.getParcelable(SAVED_HIGHLIGHT_MIXIN);
+        }
+        if (mHighlightMixin == null) {
+            mHighlightMixin = new TopLevelHighlightMixin();
+        }
+    }
+
+    @Override
+    public void onStart() {
+        // Set default highlight menu key for 1-pane homepage since it will show the placeholder
+        // page once changing back to 2-pane.
+        if (!ActivityEmbeddingUtils.isTwoPaneResolution(getActivity())) {
+            setHighlightMenuKey(getString(SettingsHomepageActivity.DEFAULT_HIGHLIGHT_MENU_KEY),
+                    /* scrollNeeded= */ false);
+        }
+        super.onStart();
     }
 
     @Override
     public void onSaveInstanceState(Bundle outState) {
         super.onSaveInstanceState(outState);
-        outState.putString(SAVED_HIGHLIGHTED_PREF, mHighlightedPreferenceKey);
-        outState.putString(SAVED_CACHED_PREF, mCachedPreferenceKey);
+        if (mHighlightMixin != null) {
+            outState.putParcelable(SAVED_HIGHLIGHT_MIXIN, mHighlightMixin);
+        }
     }
 
     @Override
@@ -170,58 +182,35 @@
 
     @Override
     public void highlightPreferenceIfNeeded() {
-        if (mTopLevelAdapter != null) {
-            mTopLevelAdapter.requestHighlight();
+        if (mHighlightMixin != null) {
+            mHighlightMixin.highlightPreferenceIfNeeded(getActivity());
         }
     }
 
+    /** Returns a {@link TopLevelHighlightMixin} that performs highlighting */
+    public TopLevelHighlightMixin getHighlightMixin() {
+        return mHighlightMixin;
+    }
+
     /** Highlight a preference with specified preference key */
     public void setHighlightPreferenceKey(String prefKey) {
-        if (mTopLevelAdapter != null) {
-            mCachedPreferenceKey = null;
-            mHighlightedPreferenceKey = prefKey;
-            mTopLevelAdapter.highlightPreference(prefKey, /* scrollNeeded= */ false);
+        // Skip Tips & support since it's full screen
+        if (mHighlightMixin != null && !TextUtils.equals(prefKey, PREF_KEY_SUPPORT)) {
+            mHighlightMixin.setHighlightPreferenceKey(prefKey);
         }
     }
 
-    /** Highlight the previous preference */
-    public void restorePreviousHighlight() {
-        if (mTopLevelAdapter != null) {
-            mTopLevelAdapter.restorePreviousHighlight();
-        }
-    }
-
-    /** Show/hide the highlight on the menu entry */
+    /** Show/hide the highlight on the menu entry for the search page presence */
     public void setMenuHighlightShowed(boolean show) {
-        if (mTopLevelAdapter == null) {
-            return;
+        if (mHighlightMixin != null) {
+            mHighlightMixin.setMenuHighlightShowed(show);
         }
-
-        if (show) {
-            mHighlightedPreferenceKey = mCachedPreferenceKey;
-            mCachedPreferenceKey = null;
-        } else {
-            if (mCachedPreferenceKey == null) {
-                mCachedPreferenceKey = mHighlightedPreferenceKey;
-            }
-            mHighlightedPreferenceKey = null;
-        }
-        mTopLevelAdapter.highlightPreference(mHighlightedPreferenceKey, /* scrollNeeded= */ show);
     }
 
     /** Highlight and scroll to a preference with specified menu key */
-    public void setHighlightMenuKey(String menuKey) {
-        if (mTopLevelAdapter == null) {
-            return;
-        }
-
-        final String prefKey = HighlightableMenu.lookupPreferenceKey(menuKey);
-        if (TextUtils.isEmpty(prefKey)) {
-            Log.e(TAG, "Invalid highlight menu key: " + menuKey);
-        } else {
-            Log.d(TAG, "Menu key: " + menuKey);
-            mHighlightedPreferenceKey = prefKey;
-            mTopLevelAdapter.highlightPreference(prefKey, /* scrollNeeded= */ true);
+    public void setHighlightMenuKey(String menuKey, boolean scrollNeeded) {
+        if (mHighlightMixin != null) {
+            mHighlightMixin.setHighlightMenuKey(menuKey, scrollNeeded);
         }
     }
 
@@ -237,16 +226,7 @@
                 || !(getActivity() instanceof SettingsHomepageActivity)) {
             return super.onCreateAdapter(preferenceScreen);
         }
-
-        if (TextUtils.isEmpty(mHighlightedPreferenceKey)) {
-            mHighlightedPreferenceKey = getHighlightPrefKeyFromArguments();
-        }
-
-        Log.d(TAG, "onCreateAdapter, pref key: " + mHighlightedPreferenceKey);
-        mTopLevelAdapter = new HighlightableTopLevelPreferenceAdapter(
-                (SettingsHomepageActivity) getActivity(), preferenceScreen, getListView(),
-                mHighlightedPreferenceKey);
-        return mTopLevelAdapter;
+        return mHighlightMixin.onCreateAdapter(this, preferenceScreen);
     }
 
     @Override
@@ -255,25 +235,9 @@
     }
 
     void reloadHighlightMenuKey() {
-        if (mTopLevelAdapter == null) {
-            return;
+        if (mHighlightMixin != null) {
+            mHighlightMixin.reloadHighlightMenuKey(getArguments());
         }
-
-        mHighlightedPreferenceKey = getHighlightPrefKeyFromArguments();
-        Log.d(TAG, "reloadHighlightMenuKey, pref key: " + mHighlightedPreferenceKey);
-        mTopLevelAdapter.highlightPreference(mHighlightedPreferenceKey, /* scrollNeeded= */ true);
-    }
-
-    private String getHighlightPrefKeyFromArguments() {
-        final Bundle arguments = getArguments();
-        final String menuKey = arguments.getString(SettingsActivity.EXTRA_FRAGMENT_ARG_KEY);
-        final String prefKey = HighlightableMenu.lookupPreferenceKey(menuKey);
-        if (TextUtils.isEmpty(prefKey)) {
-            Log.e(TAG, "Invalid highlight menu key: " + menuKey);
-        } else {
-            Log.d(TAG, "Menu key: " + menuKey);
-        }
-        return prefKey;
     }
 
     public static final BaseSearchIndexProvider SEARCH_INDEX_DATA_PROVIDER =
diff --git a/src/com/android/settings/search/SearchResultTrampoline.java b/src/com/android/settings/search/SearchResultTrampoline.java
index f9cbc36..8b041b6 100644
--- a/src/com/android/settings/search/SearchResultTrampoline.java
+++ b/src/com/android/settings/search/SearchResultTrampoline.java
@@ -109,7 +109,8 @@
             final SettingsHomepageActivity homeActivity =
                     ((SettingsApplication) getApplicationContext()).getHomeActivity();
             if (homeActivity != null) {
-                homeActivity.getMainFragment().setHighlightMenuKey(highlightMenuKey);
+                homeActivity.getMainFragment().setHighlightMenuKey(highlightMenuKey,
+                        /* scrollNeeded= */ true);
             }
         } else {
             // Two-pane case
diff --git a/src/com/android/settings/widget/HighlightableTopLevelPreferenceAdapter.java b/src/com/android/settings/widget/HighlightableTopLevelPreferenceAdapter.java
index b7f3015..ff8f805 100644
--- a/src/com/android/settings/widget/HighlightableTopLevelPreferenceAdapter.java
+++ b/src/com/android/settings/widget/HighlightableTopLevelPreferenceAdapter.java
@@ -20,6 +20,7 @@
 import android.graphics.drawable.Drawable;
 import android.text.TextUtils;
 import android.util.Log;
+import android.util.SparseArray;
 import android.util.TypedValue;
 import android.view.View;
 import android.widget.ImageView;
@@ -59,17 +60,18 @@
     private final int mNormalBackgroundRes;
     private final int mHighlightBackgroundRes;
     private String mHighlightKey;
-    private String mPreviousHighlightKey;
     private int mHighlightPosition = RecyclerView.NO_POSITION;
     private int mScrollPosition = RecyclerView.NO_POSITION;
     private boolean mHighlightNeeded;
     private boolean mScrolled;
+    private SparseArray<PreferenceViewHolder> mViewHolders;
 
     public HighlightableTopLevelPreferenceAdapter(SettingsHomepageActivity homepageActivity,
             PreferenceGroup preferenceGroup, RecyclerView recyclerView, String key) {
         super(preferenceGroup);
         mRecyclerView = recyclerView;
         mHighlightKey = key;
+        mViewHolders = new SparseArray<>();
         mContext = preferenceGroup.getContext();
         mHomepageActivity = homepageActivity;
         final TypedValue outValue = new TypedValue();
@@ -92,6 +94,7 @@
     @Override
     public void onBindViewHolder(PreferenceViewHolder holder, int position) {
         super.onBindViewHolder(holder, position);
+        mViewHolders.put(position, holder);
         updateBackground(holder, position);
     }
 
@@ -120,9 +123,9 @@
             return;
         }
 
+        final int previousPosition = mHighlightPosition;
         if (TextUtils.isEmpty(mHighlightKey)) {
             // De-highlight previous preference.
-            final int previousPosition = mHighlightPosition;
             mHighlightPosition = RecyclerView.NO_POSITION;
             mScrolled = true;
             if (previousPosition >= 0) {
@@ -145,10 +148,14 @@
 
         // Turn on/off highlight when screen split mode is changed.
         if (highlightNeeded != mHighlightNeeded) {
-            Log.d(TAG, "Highlight change needed: " + highlightNeeded);
+            Log.d(TAG, "Highlight needed change: " + highlightNeeded);
             mHighlightNeeded = highlightNeeded;
             mHighlightPosition = position;
             notifyItemChanged(position);
+            if (!highlightNeeded) {
+                // De-highlight to prevent a flicker
+                removeHighlightAt(previousPosition);
+            }
             return;
         }
 
@@ -156,7 +163,6 @@
             return;
         }
 
-        final int previousPosition = mHighlightPosition;
         mHighlightPosition = position;
         Log.d(TAG, "Request highlight position " + position);
         Log.d(TAG, "Is highlight needed: " + highlightNeeded);
@@ -178,20 +184,11 @@
      * preference is clicked.
      */
     public void highlightPreference(String key, boolean scrollNeeded) {
-        mPreviousHighlightKey = mHighlightKey;
         mHighlightKey = key;
         mScrolled = !scrollNeeded;
         requestHighlight();
     }
 
-    /**
-     * A function that restores the previous highlighted setting.
-     */
-    public void restorePreviousHighlight() {
-        mHighlightKey = mPreviousHighlightKey;
-        requestHighlight();
-    }
-
     @Override
     public void onHomepageLoaded() {
         scroll();
@@ -224,6 +221,17 @@
         }
     }
 
+    private void removeHighlightAt(int position) {
+        if (position >= 0) {
+            // De-highlight the existing preference view holder at an early stage
+            final PreferenceViewHolder holder = mViewHolders.get(position);
+            if (holder != null) {
+                removeHighlightBackground(holder);
+            }
+            notifyItemChanged(position);
+        }
+    }
+
     private void addHighlightBackground(PreferenceViewHolder holder) {
         final View v = holder.itemView;
         v.setBackgroundResource(mHighlightBackgroundRes);