Merge "Allow multiple adjacent TFs (1/n)" into main

5 files changed, 317 insertions, 39 deletions

diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 46312aff1fb6..68f812ff25b1 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2124,7 +2124,7 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
                     if (!tf.isOrganizedTaskFragment()) {
                         return;
                     }
-                    tf.resetAdjacentTaskFragment();
+                    tf.clearAdjacentTaskFragments();
                     tf.setCompanionTaskFragment(null /* companionTaskFragment */);
                     tf.setAnimationParams(TaskFragmentAnimationParams.DEFAULT);
                     if (tf.getTopNonFinishingActivity() != null) {
diff --git a/services/core/java/com/android/server/wm/TaskFragment.java b/services/core/java/com/android/server/wm/TaskFragment.java
index 51b8bd1f0091..9fb5bea132ef 100644
--- a/services/core/java/com/android/server/wm/TaskFragment.java
+++ b/services/core/java/com/android/server/wm/TaskFragment.java
@@ -94,6 +94,7 @@ import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.IBinder;
 import android.os.UserHandle;
+import android.util.ArraySet;
 import android.util.DisplayMetrics;
 import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
@@ -239,12 +240,20 @@ class TaskFragment extends WindowContainer<WindowContainer> {
     /** This task fragment will be removed when the cleanup of its children are done. */
     private boolean mIsRemovalRequested;
 
-    /** The TaskFragment that is adjacent to this one. */
+    /** @deprecated b/373709676 replace with {@link #mAdjacentTaskFragments} */
+    @Deprecated
     @Nullable
     private TaskFragment mAdjacentTaskFragment;
 
     /**
-     * Unlike the {@link mAdjacentTaskFragment}, the companion TaskFragment is not always visually
+     * The TaskFragments that are adjacent to each other, including this TaskFragment.
+     * All TaskFragments in this set share the same set instance.
+     */
+    @Nullable
+    private AdjacentSet mAdjacentTaskFragments;
+
+    /**
+     * Unlike the {@link #mAdjacentTaskFragments}, the companion TaskFragment is not always visually
      * adjacent to this one, but this TaskFragment will be removed by the organizer if the
      * companion TaskFragment is removed.
      */
@@ -442,15 +451,24 @@ class TaskFragment extends WindowContainer<WindowContainer> {
         return service.mWindowOrganizerController.getTaskFragment(token);
     }
 
-    void setAdjacentTaskFragment(@Nullable TaskFragment taskFragment) {
-        if (mAdjacentTaskFragment == taskFragment) {
-            return;
-        }
-        resetAdjacentTaskFragment();
-        if (taskFragment != null) {
+    /** @deprecated b/373709676 replace with {@link #setAdjacentTaskFragments}. */
+    @Deprecated
+    void setAdjacentTaskFragment(@NonNull TaskFragment taskFragment) {
+        if (!Flags.allowMultipleAdjacentTaskFragments()) {
+            if (mAdjacentTaskFragment == taskFragment) {
+                return;
+            }
+            resetAdjacentTaskFragment();
             mAdjacentTaskFragment = taskFragment;
             taskFragment.setAdjacentTaskFragment(this);
+            return;
         }
+
+        setAdjacentTaskFragments(new AdjacentSet(this, taskFragment));
+    }
+
+    void setAdjacentTaskFragments(@NonNull AdjacentSet adjacentTaskFragments) {
+        adjacentTaskFragments.setAsAdjacent();
     }
 
     void setCompanionTaskFragment(@Nullable TaskFragment companionTaskFragment) {
@@ -461,7 +479,14 @@ class TaskFragment extends WindowContainer<WindowContainer> {
         return mCompanionTaskFragment;
     }
 
-    void resetAdjacentTaskFragment() {
+    /** @deprecated b/373709676 replace with {@link #clearAdjacentTaskFragments()}. */
+    @Deprecated
+    private void resetAdjacentTaskFragment() {
+        if (Flags.allowMultipleAdjacentTaskFragments()) {
+            throw new IllegalStateException("resetAdjacentTaskFragment shouldn't be called when"
+                    + " allowMultipleAdjacentTaskFragments is enabled. Use either"
+                    + " #clearAdjacentTaskFragments or #removeFromAdjacentTaskFragments");
+        }
         // Reset the adjacent TaskFragment if its adjacent TaskFragment is also this TaskFragment.
         if (mAdjacentTaskFragment != null && mAdjacentTaskFragment.mAdjacentTaskFragment == this) {
             mAdjacentTaskFragment.mAdjacentTaskFragment = null;
@@ -471,6 +496,57 @@ class TaskFragment extends WindowContainer<WindowContainer> {
         mDelayLastActivityRemoval = false;
     }
 
+    void clearAdjacentTaskFragments() {
+        if (!Flags.allowMultipleAdjacentTaskFragments()) {
+            resetAdjacentTaskFragment();
+            return;
+        }
+
+        if (mAdjacentTaskFragments != null) {
+            mAdjacentTaskFragments.clear();
+        }
+    }
+
+    void removeFromAdjacentTaskFragments() {
+        if (!Flags.allowMultipleAdjacentTaskFragments()) {
+            resetAdjacentTaskFragment();
+            return;
+        }
+
+        if (mAdjacentTaskFragments != null) {
+            mAdjacentTaskFragments.remove(this);
+        }
+    }
+
+    // TODO(b/373709676): update usages.
+    /** @deprecated b/373709676 replace with {@link #getAdjacentTaskFragments()}. */
+    @Deprecated
+    @Nullable
+    TaskFragment getAdjacentTaskFragment() {
+        return mAdjacentTaskFragment;
+    }
+
+    @Nullable
+    AdjacentSet getAdjacentTaskFragments() {
+        return mAdjacentTaskFragments;
+    }
+
+    boolean hasAdjacentTaskFragment() {
+        if (!Flags.allowMultipleAdjacentTaskFragments()) {
+            return mAdjacentTaskFragment != null;
+        }
+        return mAdjacentTaskFragments != null;
+    }
+
+    boolean isAdjacentTo(@NonNull TaskFragment other) {
+        if (!Flags.allowMultipleAdjacentTaskFragments()) {
+            return mAdjacentTaskFragment == other;
+        }
+        return other != this
+                && mAdjacentTaskFragments != null
+                && mAdjacentTaskFragments.contains(other);
+    }
+
     void setTaskFragmentOrganizer(@NonNull TaskFragmentOrganizerToken organizer, int uid,
             @NonNull String processName) {
         mTaskFragmentOrganizer = ITaskFragmentOrganizer.Stub.asInterface(organizer.asBinder());
@@ -566,10 +642,6 @@ class TaskFragment extends WindowContainer<WindowContainer> {
         return isEmbedded() && mPinned;
     }
 
-    TaskFragment getAdjacentTaskFragment() {
-        return mAdjacentTaskFragment;
-    }
-
     /** Returns the currently topmost resumed activity. */
     @Nullable
     ActivityRecord getTopResumedActivity() {
@@ -616,7 +688,7 @@ class TaskFragment extends WindowContainer<WindowContainer> {
         mResumedActivity = r;
         final ActivityRecord topResumed = mTaskSupervisor.updateTopResumedActivityIfNeeded(reason);
         if (mResumedActivity != null && topResumed != null && topResumed.isEmbedded()
-                && topResumed.getTaskFragment().getAdjacentTaskFragment() == this) {
+                && topResumed.getTaskFragment().isAdjacentTo(this)) {
             // Explicitly updates the last resumed Activity if the resumed activity is
             // adjacent to the top-resumed embedded activity.
             mAtmService.setLastResumedActivityUncheckLocked(mResumedActivity, reason);
@@ -2036,7 +2108,7 @@ class TaskFragment extends WindowContainer<WindowContainer> {
     private boolean shouldReportOrientationUnspecified() {
         // Apps and their containers are not allowed to specify orientation from adjacent
         // TaskFragment.
-        return getAdjacentTaskFragment() != null && isVisibleRequested();
+        return hasAdjacentTaskFragment() && isVisibleRequested();
     }
 
     @Override
@@ -3086,7 +3158,7 @@ class TaskFragment extends WindowContainer<WindowContainer> {
             EventLogTags.writeWmTfRemoved(System.identityHashCode(this), getTaskId());
         }
         mIsRemovalRequested = false;
-        resetAdjacentTaskFragment();
+        removeFromAdjacentTaskFragments();
         cleanUpEmbeddedTaskFragment();
         final boolean shouldExecuteAppTransition =
                 mClearedTaskFragmentForPip && isTaskVisibleRequested();
@@ -3267,9 +3339,16 @@ class TaskFragment extends WindowContainer<WindowContainer> {
             sb.append(" organizerProc=");
             sb.append(mTaskFragmentOrganizerProcessName);
         }
-        if (mAdjacentTaskFragment != null) {
-            sb.append(" adjacent=");
-            sb.append(mAdjacentTaskFragment);
+        if (Flags.allowMultipleAdjacentTaskFragments()) {
+            if (mAdjacentTaskFragments != null) {
+                sb.append(" adjacent=");
+                sb.append(mAdjacentTaskFragments);
+            }
+        } else {
+            if (mAdjacentTaskFragment != null) {
+                sb.append(" adjacent=");
+                sb.append(mAdjacentTaskFragment);
+            }
         }
         sb.append('}');
         return sb.toString();
@@ -3385,4 +3464,94 @@ class TaskFragment extends WindowContainer<WindowContainer> {
 
         proto.end(token);
     }
+
+    /** Set of {@link TaskFragment}s that are adjacent to each other. */
+    static class AdjacentSet {
+        private final ArraySet<TaskFragment> mAdjacentSet;
+
+        AdjacentSet(@NonNull TaskFragment... taskFragments) {
+            this(new ArraySet<>(taskFragments));
+        }
+
+        AdjacentSet(@NonNull ArraySet<TaskFragment> taskFragments) {
+            if (!Flags.allowMultipleAdjacentTaskFragments()) {
+                throw new IllegalStateException("allowMultipleAdjacentTaskFragments must be"
+                        + " enabled to set more than two TaskFragments adjacent to each other.");
+            }
+            if (taskFragments.size() < 2) {
+                throw new IllegalArgumentException("Adjacent TaskFragments must contain at least"
+                        + " two TaskFragments, but only " + taskFragments.size()
+                        + " were provided.");
+            }
+            mAdjacentSet = taskFragments;
+        }
+
+        /** Updates each {@link TaskFragment} in the set to be adjacent to each other. */
+        private void setAsAdjacent() {
+            if (mAdjacentSet.isEmpty()
+                    || equals(mAdjacentSet.valueAt(0).mAdjacentTaskFragments)) {
+                // No need to update if any TaskFragment in the set has already been updated to the
+                // same set.
+                return;
+            }
+            for (int i = mAdjacentSet.size() - 1; i >= 0; i--) {
+                final TaskFragment taskFragment = mAdjacentSet.valueAt(i);
+                taskFragment.removeFromAdjacentTaskFragments();
+                taskFragment.mAdjacentTaskFragments = this;
+            }
+        }
+
+        /** Removes the {@link TaskFragment} from the adjacent set. */
+        private void remove(@NonNull TaskFragment taskFragment) {
+            taskFragment.mAdjacentTaskFragments = null;
+            taskFragment.mDelayLastActivityRemoval = false;
+            mAdjacentSet.remove(taskFragment);
+            if (mAdjacentSet.size() < 2) {
+                // To be considered as adjacent, there must be at least 2 TaskFragments in the set.
+                clear();
+            }
+        }
+
+        /** Clears the adjacent relationship. */
+        private void clear() {
+            for (int i = mAdjacentSet.size() - 1; i >= 0; i--) {
+                final TaskFragment taskFragment = mAdjacentSet.valueAt(i);
+                // Clear all reference.
+                taskFragment.mAdjacentTaskFragments = null;
+                taskFragment.mDelayLastActivityRemoval = false;
+            }
+            mAdjacentSet.clear();
+        }
+
+        /** Whether the {@link TaskFragment} is in this adjacent set. */
+        boolean contains(@NonNull TaskFragment taskFragment) {
+            return mAdjacentSet.contains(taskFragment);
+        }
+
+        @Override
+        public boolean equals(@Nullable Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (!(o instanceof AdjacentSet other)) {
+                return false;
+            }
+            return mAdjacentSet.equals(other.mAdjacentSet);
+        }
+
+        @Override
+        public String toString() {
+            final StringBuilder sb = new StringBuilder();
+            sb.append("AdjacentSet{");
+            final int size = mAdjacentSet.size();
+            for (int i = 0; i < size; i++) {
+                if (i != 0) {
+                    sb.append(", ");
+                }
+                sb.append(mAdjacentSet.valueAt(i));
+            }
+            sb.append("}");
+            return sb.toString();
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/wm/WindowOrganizerController.java b/services/core/java/com/android/server/wm/WindowOrganizerController.java
index 66921ff3adeb..1eb84650d591 100644
--- a/services/core/java/com/android/server/wm/WindowOrganizerController.java
+++ b/services/core/java/com/android/server/wm/WindowOrganizerController.java
@@ -1155,7 +1155,7 @@ class WindowOrganizerController extends IWindowOrganizerController.Stub
                 } else if (!task.mCreatedByOrganizer) {
                     throw new UnsupportedOperationException(
                             "Cannot set non-organized task as adjacent flag root: " + wc);
-                } else if (task.getAdjacentTaskFragment() == null && !clearRoot) {
+                } else if (!task.hasAdjacentTaskFragment() && !clearRoot) {
                     throw new UnsupportedOperationException(
                             "Cannot set non-adjacent task as adjacent flag root: " + wc);
                 }
@@ -1645,9 +1645,15 @@ class WindowOrganizerController extends IWindowOrganizerController.Stub
                             opType, exception);
                     break;
                 }
-                if (taskFragment.getAdjacentTaskFragment() != secondaryTaskFragment) {
+                if (!taskFragment.isAdjacentTo(secondaryTaskFragment)) {
                     // Only have lifecycle effect if the adjacent changed.
-                    taskFragment.setAdjacentTaskFragment(secondaryTaskFragment);
+                    if (Flags.allowMultipleAdjacentTaskFragments()) {
+                        // Activity Embedding only set two TFs adjacent.
+                        taskFragment.setAdjacentTaskFragments(
+                                new TaskFragment.AdjacentSet(taskFragment, secondaryTaskFragment));
+                    } else {
+                        taskFragment.setAdjacentTaskFragment(secondaryTaskFragment);
+                    }
                     effects |= TRANSACT_EFFECTS_LIFECYCLE;
                 }
 
@@ -1663,21 +1669,25 @@ class WindowOrganizerController extends IWindowOrganizerController.Stub
                 break;
             }
             case OP_TYPE_CLEAR_ADJACENT_TASK_FRAGMENTS: {
-                final TaskFragment adjacentTaskFragment = taskFragment.getAdjacentTaskFragment();
-                if (adjacentTaskFragment == null) {
+                if (!taskFragment.hasAdjacentTaskFragment()) {
                     break;
                 }
-                taskFragment.resetAdjacentTaskFragment();
-                effects |= TRANSACT_EFFECTS_LIFECYCLE;
 
-                // Clear the focused app if the focused app is no longer visible after reset the
-                // adjacent TaskFragments.
+                // Check if the focused app is in the adjacent set that will be cleared.
                 final ActivityRecord focusedApp = taskFragment.getDisplayContent().mFocusedApp;
                 final TaskFragment focusedTaskFragment = focusedApp != null
                         ? focusedApp.getTaskFragment()
                         : null;
-                if ((focusedTaskFragment == taskFragment
-                        || focusedTaskFragment == adjacentTaskFragment)
+                final boolean wasFocusedInAdjacent = focusedTaskFragment == taskFragment
+                        || (focusedTaskFragment != null
+                        && taskFragment.isAdjacentTo(focusedTaskFragment));
+
+                taskFragment.removeFromAdjacentTaskFragments();
+                effects |= TRANSACT_EFFECTS_LIFECYCLE;
+
+                // Clear the focused app if the focused app is no longer visible after reset the
+                // adjacent TaskFragments.
+                if (wasFocusedInAdjacent
                         && !focusedTaskFragment.shouldBeVisible(null /* starting */)) {
                     focusedTaskFragment.getDisplayContent().setFocusedApp(null /* newFocus */);
                 }
@@ -2207,10 +2217,15 @@ class WindowOrganizerController extends IWindowOrganizerController.Stub
             throw new IllegalArgumentException("setAdjacentRootsHierarchyOp: Not created by"
                     + " organizer root1=" + root1 + " root2=" + root2);
         }
-        if (root1.getAdjacentTaskFragment() == root2) {
+        if (root1.isAdjacentTo(root2)) {
             return TRANSACT_EFFECTS_NONE;
         }
-        root1.setAdjacentTaskFragment(root2);
+        if (Flags.allowMultipleAdjacentTaskFragments()) {
+            // TODO(b/373709676): allow three roots.
+            root1.setAdjacentTaskFragments(new TaskFragment.AdjacentSet(root1, root2));
+        } else {
+            root1.setAdjacentTaskFragment(root2);
+        }
         return TRANSACT_EFFECTS_LIFECYCLE;
     }
 
@@ -2225,10 +2240,10 @@ class WindowOrganizerController extends IWindowOrganizerController.Stub
             throw new IllegalArgumentException("clearAdjacentRootsHierarchyOp: Not created by"
                     + " organizer root=" + root);
         }
-        if (root.getAdjacentTaskFragment() == null) {
+        if (!root.hasAdjacentTaskFragment()) {
             return TRANSACT_EFFECTS_NONE;
         }
-        root.resetAdjacentTaskFragment();
+        root.removeFromAdjacentTaskFragments();
         return TRANSACT_EFFECTS_LIFECYCLE;
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
index 40da9ea2d718..579ed6659976 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
@@ -365,8 +365,8 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         assertTrue(outPrevActivities.isEmpty());
         assertTrue(predictable);
         // reset
-        tf1.setAdjacentTaskFragment(null);
-        tf2.setAdjacentTaskFragment(null);
+        tf1.clearAdjacentTaskFragments();
+        tf2.clearAdjacentTaskFragments();
         tf1.setCompanionTaskFragment(null);
         tf2.setCompanionTaskFragment(null);
 
@@ -398,8 +398,8 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         assertTrue(predictable);
         // reset
         outPrevActivities.clear();
-        tf2.setAdjacentTaskFragment(null);
-        tf3.setAdjacentTaskFragment(null);
+        tf2.clearAdjacentTaskFragments();
+        tf3.clearAdjacentTaskFragments();
 
         final TaskFragment tf4 = createTaskFragmentWithActivity(task);
         // Stacked + next companion to top => predict for previous activity below companion.
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
index 65a6a69fc45e..dafa96f91812 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
@@ -51,6 +51,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.clearInvocations;
@@ -62,6 +63,7 @@ import android.content.res.Configuration;
 import android.graphics.Color;
 import android.graphics.Rect;
 import android.os.Binder;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.view.SurfaceControl;
 import android.view.View;
@@ -1066,6 +1068,98 @@ public class TaskFragmentTest extends WindowTestsBase {
                 Math.min(outConfig.screenWidthDp, outConfig.screenHeightDp));
     }
 
+    @EnableFlags(Flags.FLAG_ALLOW_MULTIPLE_ADJACENT_TASK_FRAGMENTS)
+    @Test
+    public void testSetAdjacentTaskFragments() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragment tf0 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf1 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf2 = createTaskFragmentWithActivity(task);
+        final TaskFragment.AdjacentSet adjacentTfs = new TaskFragment.AdjacentSet(tf0, tf1, tf2);
+        assertFalse(tf0.hasAdjacentTaskFragment());
+
+        tf0.setAdjacentTaskFragments(adjacentTfs);
+
+        assertSame(adjacentTfs, tf0.getAdjacentTaskFragments());
+        assertSame(adjacentTfs, tf1.getAdjacentTaskFragments());
+        assertSame(adjacentTfs, tf2.getAdjacentTaskFragments());
+        assertTrue(tf0.hasAdjacentTaskFragment());
+        assertTrue(tf1.hasAdjacentTaskFragment());
+        assertTrue(tf2.hasAdjacentTaskFragment());
+
+        final TaskFragment.AdjacentSet adjacentTfs2 = new TaskFragment.AdjacentSet(tf0, tf1);
+        tf0.setAdjacentTaskFragments(adjacentTfs2);
+
+        assertSame(adjacentTfs2, tf0.getAdjacentTaskFragments());
+        assertSame(adjacentTfs2, tf1.getAdjacentTaskFragments());
+        assertNull(tf2.getAdjacentTaskFragments());
+        assertTrue(tf0.hasAdjacentTaskFragment());
+        assertTrue(tf1.hasAdjacentTaskFragment());
+        assertFalse(tf2.hasAdjacentTaskFragment());
+    }
+
+    @EnableFlags(Flags.FLAG_ALLOW_MULTIPLE_ADJACENT_TASK_FRAGMENTS)
+    @Test
+    public void testClearAdjacentTaskFragments() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragment tf0 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf1 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf2 = createTaskFragmentWithActivity(task);
+        final TaskFragment.AdjacentSet adjacentTfs = new TaskFragment.AdjacentSet(tf0, tf1, tf2);
+        tf0.setAdjacentTaskFragments(adjacentTfs);
+
+        tf0.clearAdjacentTaskFragments();
+
+        assertNull(tf0.getAdjacentTaskFragments());
+        assertNull(tf1.getAdjacentTaskFragments());
+        assertNull(tf2.getAdjacentTaskFragments());
+        assertFalse(tf0.hasAdjacentTaskFragment());
+        assertFalse(tf1.hasAdjacentTaskFragment());
+        assertFalse(tf2.hasAdjacentTaskFragment());
+    }
+
+    @EnableFlags(Flags.FLAG_ALLOW_MULTIPLE_ADJACENT_TASK_FRAGMENTS)
+    @Test
+    public void testRemoveFromAdjacentTaskFragments() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragment tf0 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf1 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf2 = createTaskFragmentWithActivity(task);
+        final TaskFragment.AdjacentSet adjacentTfs = new TaskFragment.AdjacentSet(tf0, tf1, tf2);
+        tf0.setAdjacentTaskFragments(adjacentTfs);
+
+        tf0.removeFromAdjacentTaskFragments();
+
+        assertNull(tf0.getAdjacentTaskFragments());
+        assertSame(adjacentTfs, tf1.getAdjacentTaskFragments());
+        assertSame(adjacentTfs, tf2.getAdjacentTaskFragments());
+        assertFalse(adjacentTfs.contains(tf0));
+        assertTrue(tf1.isAdjacentTo(tf2));
+        assertTrue(tf2.isAdjacentTo(tf1));
+        assertFalse(tf1.isAdjacentTo(tf0));
+        assertFalse(tf0.isAdjacentTo(tf1));
+        assertFalse(tf0.isAdjacentTo(tf0));
+        assertFalse(tf1.isAdjacentTo(tf1));
+    }
+
+    @EnableFlags(Flags.FLAG_ALLOW_MULTIPLE_ADJACENT_TASK_FRAGMENTS)
+    @Test
+    public void testRemoveFromAdjacentTaskFragmentsWhenRemove() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragment tf0 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf1 = createTaskFragmentWithActivity(task);
+        final TaskFragment tf2 = createTaskFragmentWithActivity(task);
+        final TaskFragment.AdjacentSet adjacentTfs = new TaskFragment.AdjacentSet(tf0, tf1, tf2);
+        tf0.setAdjacentTaskFragments(adjacentTfs);
+
+        tf0.removeImmediately();
+
+        assertNull(tf0.getAdjacentTaskFragments());
+        assertSame(adjacentTfs, tf1.getAdjacentTaskFragments());
+        assertSame(adjacentTfs, tf2.getAdjacentTaskFragments());
+        assertFalse(adjacentTfs.contains(tf0));
+    }
+
     private WindowState createAppWindow(ActivityRecord app, String name) {
         final WindowState win = createWindow(null, TYPE_BASE_APPLICATION, app, name,
                 0 /* ownerId */, false /* ownerCanAddInternalSystemWindow */, new TestIWindow());