Merge "Reconstruct the containers from the saved state" into main

9 files changed, 240 insertions, 21 deletions

diff --git a/core/java/android/window/TaskFragmentOrganizer.java b/core/java/android/window/TaskFragmentOrganizer.java
index 4cc0d8a77a2b..c316800108bd 100644
--- a/core/java/android/window/TaskFragmentOrganizer.java
+++ b/core/java/android/window/TaskFragmentOrganizer.java
@@ -69,6 +69,23 @@ public class TaskFragmentOrganizer extends WindowOrganizer {
     public static final String KEY_ERROR_CALLBACK_OP_TYPE = "operation_type";
 
     /**
+     * Key to bundle {@link TaskFragmentInfo}s from the system in
+     * {@link #registerOrganizer(boolean, Bundle)}
+     *
+     * @hide
+     */
+    public static final String KEY_RESTORE_TASK_FRAGMENTS_INFO = "key_restore_task_fragments_info";
+
+    /**
+     * Key to bundle {@link TaskFragmentParentInfo} from the system in
+     * {@link #registerOrganizer(boolean, Bundle)}
+     *
+     * @hide
+     */
+    public static final String KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO =
+            "key_restore_task_fragment_parent_info";
+
+    /**
      * No change set.
      */
     @WindowManager.TransitionType
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
index 4ce294213526..bfccb29bc952 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
@@ -16,16 +16,26 @@
 
 package androidx.window.extensions.embedding;
 
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENTS_INFO;
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO;
+
 import android.os.Build;
 import android.os.Bundle;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.MessageQueue;
+import android.util.ArrayMap;
 import android.util.Log;
+import android.util.SparseArray;
+import android.window.TaskFragmentInfo;
+import android.window.TaskFragmentParentInfo;
+import android.window.WindowContainerTransaction;
 
 import androidx.annotation.NonNull;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 
 /**
  * Helper class to back up and restore the TaskFragmentOrganizer state, in order to resume
@@ -40,11 +50,21 @@ class BackupHelper {
     @NonNull
     private final SplitController mController;
     @NonNull
+    private final SplitPresenter mPresenter;
+    @NonNull
     private final BackupIdler mBackupIdler = new BackupIdler();
     private boolean mBackupIdlerScheduled;
 
-    BackupHelper(@NonNull SplitController splitController, @NonNull Bundle savedState) {
+    private final List<ParcelableTaskContainerData> mParcelableTaskContainerDataList =
+            new ArrayList<>();
+    private final ArrayMap<IBinder, TaskFragmentInfo> mTaskFragmentInfos = new ArrayMap<>();
+    private final SparseArray<TaskFragmentParentInfo> mTaskFragmentParentInfos =
+            new SparseArray<>();
+
+    BackupHelper(@NonNull SplitController splitController, @NonNull SplitPresenter splitPresenter,
+            @NonNull Bundle savedState) {
         mController = splitController;
+        mPresenter = splitPresenter;
 
         if (!savedState.isEmpty()) {
             restoreState(savedState);
@@ -67,13 +87,13 @@ class BackupHelper {
         public boolean queueIdle() {
             synchronized (mController.mLock) {
                 mBackupIdlerScheduled = false;
-                startBackup();
+                saveState();
             }
             return false;
         }
     }
 
-    private void startBackup() {
+    private void saveState() {
         final List<TaskContainer> taskContainers = mController.getTaskContainers();
         if (taskContainers.isEmpty()) {
             Log.w(TAG, "No task-container to back up");
@@ -97,13 +117,92 @@ class BackupHelper {
             return;
         }
 
-        final List<ParcelableTaskContainerData> parcelableTaskContainerDataList =
-                savedState.getParcelableArrayList(KEY_TASK_CONTAINERS,
-                        ParcelableTaskContainerData.class);
-        for (ParcelableTaskContainerData data : parcelableTaskContainerDataList) {
-            final TaskContainer taskContainer = new TaskContainer(data, mController);
-            if (DEBUG) Log.d(TAG, "Restoring task " + taskContainer.getTaskId());
-            // TODO(b/289875940): implement the TaskContainer restoration.
+        if (DEBUG) Log.d(TAG, "Start restoring saved-state");
+        mParcelableTaskContainerDataList.addAll(savedState.getParcelableArrayList(
+                KEY_TASK_CONTAINERS, ParcelableTaskContainerData.class));
+        if (DEBUG) Log.d(TAG, "Retrieved tasks : " + mParcelableTaskContainerDataList.size());
+        if (mParcelableTaskContainerDataList.isEmpty()) {
+            return;
+        }
+
+        final List<TaskFragmentInfo> infos = savedState.getParcelableArrayList(
+                KEY_RESTORE_TASK_FRAGMENTS_INFO, TaskFragmentInfo.class);
+        for (TaskFragmentInfo info : infos) {
+            if (DEBUG) Log.d(TAG, "Retrieved: " + info);
+            mTaskFragmentInfos.put(info.getFragmentToken(), info);
+            mPresenter.updateTaskFragmentInfo(info);
+        }
+
+        final List<TaskFragmentParentInfo> parentInfos = savedState.getParcelableArrayList(
+                KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO,
+                TaskFragmentParentInfo.class);
+        for (TaskFragmentParentInfo info : parentInfos) {
+            if (DEBUG) Log.d(TAG, "Retrieved: " + info);
+            mTaskFragmentParentInfos.put(info.getTaskId(), info);
+        }
+    }
+
+    boolean hasPendingStateToRestore() {
+        return !mParcelableTaskContainerDataList.isEmpty();
+    }
+
+    /**
+     * Returns {@code true} if any of the {@link TaskContainer} is restored.
+     * Otherwise, returns {@code false}.
+     */
+    boolean rebuildTaskContainers(@NonNull WindowContainerTransaction wct,
+            @NonNull Set<EmbeddingRule> rules) {
+        if (mParcelableTaskContainerDataList.isEmpty()) {
+            return false;
+        }
+
+        if (DEBUG) Log.d(TAG, "Rebuilding TaskContainers.");
+        final ArrayMap<String, EmbeddingRule> embeddingRuleMap = new ArrayMap<>();
+        for (EmbeddingRule rule : rules) {
+            embeddingRuleMap.put(rule.getTag(), rule);
+        }
+
+        boolean restoredAny = false;
+        for (int i = mParcelableTaskContainerDataList.size() - 1; i >= 0; i--) {
+            final ParcelableTaskContainerData parcelableTaskContainerData =
+                    mParcelableTaskContainerDataList.get(i);
+            final List<String> tags = parcelableTaskContainerData.getSplitRuleTags();
+            if (!embeddingRuleMap.containsAll(tags)) {
+                // has unknown tag, unable to restore.
+                if (DEBUG) {
+                    Log.d(TAG, "Rebuilding TaskContainer abort! Unknown Tag. Task#"
+                            + parcelableTaskContainerData.mTaskId);
+                }
+                continue;
+            }
+
+            mParcelableTaskContainerDataList.remove(parcelableTaskContainerData);
+            final TaskContainer taskContainer = new TaskContainer(parcelableTaskContainerData,
+                    mController, mTaskFragmentInfos);
+            if (DEBUG) Log.d(TAG, "Created TaskContainer " + taskContainer);
+            mController.addTaskContainer(taskContainer.getTaskId(), taskContainer);
+
+            for (ParcelableSplitContainerData splitData :
+                    parcelableTaskContainerData.getParcelableSplitContainerDataList()) {
+                final SplitRule rule = (SplitRule) embeddingRuleMap.get(splitData.mSplitRuleTag);
+                assert rule != null;
+                if (mController.getContainer(splitData.getPrimaryContainerToken()) != null
+                        && mController.getContainer(splitData.getSecondaryContainerToken())
+                        != null) {
+                    taskContainer.addSplitContainer(
+                            new SplitContainer(splitData, mController, rule));
+                }
+            }
+
+            mController.onTaskFragmentParentRestored(wct, taskContainer.getTaskId(),
+                    mTaskFragmentParentInfos.get(taskContainer.getTaskId()));
+            restoredAny = true;
+        }
+
+        if (mParcelableTaskContainerDataList.isEmpty()) {
+            mTaskFragmentParentInfos.clear();
+            mTaskFragmentInfos.clear();
         }
+        return restoredAny;
     }
-}
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
index 817cfce69b2e..cb280c530c1b 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
@@ -89,13 +89,13 @@ class ParcelableSplitContainerData implements Parcelable {
     };
 
     @NonNull
-    private IBinder getPrimaryContainerToken() {
+    IBinder getPrimaryContainerToken() {
         return mSplitContainer != null ? mSplitContainer.getPrimaryContainer().getToken()
                 : mPrimaryContainerToken;
     }
 
     @NonNull
-    private IBinder getSecondaryContainerToken() {
+    IBinder getSecondaryContainerToken() {
         return mSplitContainer != null ? mSplitContainer.getSecondaryContainer().getToken()
                 : mSecondaryContainerToken;
     }
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
index 7377d005cda4..97aa69985907 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
@@ -108,6 +108,15 @@ class ParcelableTaskContainerData implements Parcelable {
                 : mParcelableSplitContainerDataList;
     }
 
+    @NonNull
+    List<String> getSplitRuleTags() {
+        final List<String> tags = new ArrayList<>();
+        for (ParcelableSplitContainerData data : getParcelableSplitContainerDataList()) {
+            tags.add(data.mSplitRuleTag);
+        }
+        return tags;
+    }
+
     @Override
     public int describeContents() {
         return 0;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
index 6d436ec01d98..faf73c24073f 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
@@ -86,6 +86,25 @@ class SplitContainer {
         }
     }
 
+    /** This is only used when restoring it from a {@link ParcelableSplitContainerData}. */
+    SplitContainer(@NonNull ParcelableSplitContainerData parcelableData,
+            @NonNull SplitController splitController, @NonNull SplitRule splitRule) {
+        mParcelableData = parcelableData;
+        mPrimaryContainer = splitController.getContainer(parcelableData.getPrimaryContainerToken());
+        mSecondaryContainer = splitController.getContainer(
+                parcelableData.getSecondaryContainerToken());
+        mSplitRule = splitRule;
+        mDefaultSplitAttributes = splitRule.getDefaultSplitAttributes();
+        mCurrentSplitAttributes = mDefaultSplitAttributes;
+
+        if (shouldFinishPrimaryWithSecondary(splitRule)) {
+            mSecondaryContainer.addContainerToFinishOnExit(mPrimaryContainer);
+        }
+        if (shouldFinishSecondaryWithPrimary(splitRule)) {
+            mPrimaryContainer.addContainerToFinishOnExit(mSecondaryContainer);
+        }
+    }
+
     void setPrimaryContainer(@NonNull TaskFragmentContainer primaryContainer) {
         if (!mParcelableData.mIsPrimaryContainerMutable) {
             throw new IllegalStateException("Cannot update primary TaskFragmentContainer");
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index f2f2b7ea7174..db4bb0e5e75e 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -279,6 +279,26 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
             Log.i(TAG, "Setting embedding rules. Size: " + rules.size());
             mSplitRules.clear();
             mSplitRules.addAll(rules);
+
+            if (!Flags.aeBackStackRestore() || !mPresenter.isRebuildTaskContainersNeeded()) {
+                return;
+            }
+
+            try {
+                final TransactionRecord transactionRecord =
+                        mTransactionManager.startNewTransaction();
+                final WindowContainerTransaction wct = transactionRecord.getTransaction();
+                if (mPresenter.rebuildTaskContainers(wct, rules)) {
+                    transactionRecord.apply(false /* shouldApplyIndependently */);
+                    updateCallbackIfNecessary();
+                } else {
+                    transactionRecord.abort();
+                }
+            } catch (IllegalStateException ex) {
+                Log.e(TAG, "Having an existing transaction while running restoration with"
+                        + "new rules!! It is likely too late to perform the restoration "
+                        + "already!?", ex);
+            }
         }
     }
 
@@ -903,6 +923,12 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
     }
 
     @GuardedBy("mLock")
+    void onTaskFragmentParentRestored(@NonNull WindowContainerTransaction wct, int taskId,
+            @NonNull TaskFragmentParentInfo parentInfo) {
+        onTaskFragmentParentInfoChanged(wct, taskId, parentInfo);
+    }
+
+    @GuardedBy("mLock")
     void updateContainersInTaskIfVisible(@NonNull WindowContainerTransaction wct, int taskId) {
         final TaskContainer taskContainer = getTaskContainer(taskId);
         if (taskContainer == null) {
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
index abc7b291fc32..0c0ded9bad74 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
@@ -24,6 +24,7 @@ import static androidx.window.extensions.embedding.SplitController.TAG;
 import static androidx.window.extensions.embedding.WindowAttributes.DIM_AREA_ON_TASK;
 
 import android.annotation.AnimRes;
+import android.annotation.NonNull;
 import android.app.Activity;
 import android.app.ActivityThread;
 import android.app.WindowConfiguration;
@@ -47,7 +48,6 @@ import android.window.TaskFragmentCreationParams;
 import android.window.WindowContainerTransaction;
 
 import androidx.annotation.IntDef;
-import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.window.extensions.core.util.function.Function;
 import androidx.window.extensions.embedding.SplitAttributes.SplitType;
@@ -67,6 +67,7 @@ import com.android.window.flags.Flags;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.Executor;
 
 /**
@@ -174,7 +175,7 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
         } else {
             registerOrganizer();
         }
-        mBackupHelper = new BackupHelper(controller, outSavedState);
+        mBackupHelper = new BackupHelper(controller, this, outSavedState);
         if (!SplitController.ENABLE_SHELL_TRANSITIONS) {
             // TODO(b/207070762): cleanup with legacy app transition
             // Animation will be handled by WM Shell when Shell transition is enabled.
@@ -186,6 +187,15 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
         mBackupHelper.scheduleBackup();
     }
 
+    boolean isRebuildTaskContainersNeeded() {
+        return mBackupHelper.hasPendingStateToRestore();
+    }
+
+    boolean rebuildTaskContainers(@NonNull WindowContainerTransaction wct,
+            @NonNull Set<EmbeddingRule> rules) {
+        return mBackupHelper.rebuildTaskContainers(wct, rules);
+    }
+
     /**
      * Deletes the specified container and all other associated and dependent containers in the same
      * transaction.
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
index 608a3bee7509..74cce68f270b 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
@@ -31,6 +31,7 @@ import android.app.WindowConfiguration.WindowingMode;
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.os.IBinder;
+import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.DisplayMetrics;
 import android.util.Log;
@@ -147,14 +148,23 @@ class TaskContainer {
 
     /** This is only used when restoring it from a {@link ParcelableTaskContainerData}. */
     TaskContainer(@NonNull ParcelableTaskContainerData data,
-            @NonNull SplitController splitController) {
+            @NonNull SplitController splitController,
+            @NonNull ArrayMap<IBinder, TaskFragmentInfo> taskFragmentInfoMap) {
         mParcelableTaskContainerData = new ParcelableTaskContainerData(data, this);
+        mInfo = new TaskFragmentParentInfo(new Configuration(), 0 /* displayId */, -1 /* taskId */,
+                false /* visible */, false /* hasDirectActivity */, null /* decorSurface */);
         mSplitController = splitController;
         for (ParcelableTaskFragmentContainerData tfData :
                 data.getParcelableTaskFragmentContainerDataList()) {
-            final TaskFragmentContainer container =
-                    new TaskFragmentContainer(tfData, splitController, this);
-            mContainers.add(container);
+            final TaskFragmentInfo info = taskFragmentInfoMap.get(tfData.mToken);
+            if (info != null && !info.isEmpty()) {
+                final TaskFragmentContainer container =
+                        new TaskFragmentContainer(tfData, splitController, this);
+                container.setInfo(new WindowContainerTransaction(), info);
+                mContainers.add(container);
+            } else {
+                Log.d(TAG, "Drop " + tfData + " while restoring Task " + data.mTaskId);
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
index 5aa34d22f00f..92953e5a5041 100644
--- a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
@@ -17,6 +17,8 @@
 package com.android.server.wm;
 
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENTS_INFO;
+import static android.window.TaskFragmentOrganizer.KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO;
 import static android.window.TaskFragmentOrganizer.putErrorInfoInBundle;
 import static android.window.TaskFragmentTransaction.TYPE_ACTIVITY_REPARENTED_TO_TASK;
 import static android.window.TaskFragmentTransaction.TYPE_TASK_FRAGMENT_APPEARED;
@@ -206,7 +208,13 @@ public class TaskFragmentOrganizerController extends ITaskFragmentOrganizerContr
             mOrganizerPid = pid;
             mAppThread = getAppThread(pid, mOrganizerUid);
             for (int i = mOrganizedTaskFragments.size() - 1; i >= 0; i--) {
-                mOrganizedTaskFragments.get(i).onTaskFragmentOrganizerRestarted(organizer);
+                final TaskFragment taskFragment = mOrganizedTaskFragments.get(i);
+                if (taskFragment.isAttached()
+                        && taskFragment.getTopNonFinishingActivity() != null) {
+                    taskFragment.onTaskFragmentOrganizerRestarted(organizer);
+                } else {
+                    mOrganizedTaskFragments.remove(taskFragment);
+                }
             }
             try {
                 mOrganizer.asBinder().linkToDeath(this, 0 /*flags*/);
@@ -575,8 +583,29 @@ public class TaskFragmentOrganizerController extends ITaskFragmentOrganizerContr
         }
 
         mCachedTaskFragmentOrganizerStates.remove(cachedState);
-        outSavedState.putAll(cachedState.mSavedState);
         cachedState.restore(organizer, pid);
+        outSavedState.putAll(cachedState.mSavedState);
+
+        // Collect the organized TfInfo and TfParentInfo in the system.
+        final ArrayList<TaskFragmentInfo> infos = new ArrayList<>();
+        final ArrayMap<Integer, Task> tasks = new ArrayMap<>();
+        final int fragmentCount = cachedState.mOrganizedTaskFragments.size();
+        for (int j = 0; j < fragmentCount; j++) {
+            final TaskFragment tf = cachedState.mOrganizedTaskFragments.get(j);
+            infos.add(tf.getTaskFragmentInfo());
+            if (!tasks.containsKey(tf.getTask().mTaskId)) {
+                tasks.put(tf.getTask().mTaskId, tf.getTask());
+            }
+        }
+        outSavedState.putParcelableArrayList(KEY_RESTORE_TASK_FRAGMENTS_INFO, infos);
+
+        final ArrayList<TaskFragmentParentInfo> parentInfos = new ArrayList<>();
+        for (int j = tasks.size() - 1; j >= 0; j--) {
+            parentInfos.add(tasks.valueAt(j).getTaskFragmentParentInfo());
+        }
+        outSavedState.putParcelableArrayList(KEY_RESTORE_TASK_FRAGMENT_PARENT_INFO,
+                parentInfos);
+
         mTaskFragmentOrganizerState.put(organizer.asBinder(), cachedState);
         mPendingTaskFragmentEvents.put(organizer.asBinder(), new ArrayList<>());
         return true;