Listen to folding state change for SplitController

This CL enables to update the current top SplitContainer if
there's a folding state change.

Test: manual - following CLs use this CL
Bug: 207494880
Bug: 241043111
Change-Id: Ia4d08d7f71d935101f3aeda94384295db2fa2214
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 8d8e2d5..138eee4 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -19,6 +19,7 @@
 import static android.app.ActivityManager.START_SUCCESS;
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
+import static android.view.Display.DEFAULT_DISPLAY;
 import static android.window.TaskFragmentOrganizer.KEY_ERROR_CALLBACK_OP_TYPE;
 import static android.window.TaskFragmentOrganizer.KEY_ERROR_CALLBACK_TASK_FRAGMENT_INFO;
 import static android.window.TaskFragmentOrganizer.KEY_ERROR_CALLBACK_THROWABLE;
@@ -46,6 +47,7 @@
 import android.app.ActivityClient;
 import android.app.ActivityOptions;
 import android.app.ActivityThread;
+import android.app.Application;
 import android.app.Instrumentation;
 import android.content.ComponentName;
 import android.content.Context;
@@ -70,12 +72,16 @@
 import androidx.annotation.GuardedBy;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.window.common.CommonFoldingFeature;
 import androidx.window.common.EmptyLifecycleCallbacksAdapter;
+import androidx.window.extensions.WindowExtensionsProvider;
+import androidx.window.extensions.layout.WindowLayoutComponentImpl;
 
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
@@ -106,26 +112,65 @@
     @GuardedBy("mLock")
     final SparseArray<TaskContainer> mTaskContainers = new SparseArray<>();
 
-    // Callback to Jetpack to notify about changes to split states.
-    @NonNull
+    /** Callback to Jetpack to notify about changes to split states. */
+    @Nullable
     private Consumer<List<SplitInfo>> mEmbeddingCallback;
     private final List<SplitInfo> mLastReportedSplitStates = new ArrayList<>();
     private final Handler mHandler;
     final Object mLock = new Object();
     private final ActivityStartMonitor mActivityStartMonitor;
+    @NonNull
+    final WindowLayoutComponentImpl mWindowLayoutComponent;
 
     public SplitController() {
+        this((WindowLayoutComponentImpl) Objects.requireNonNull(WindowExtensionsProvider
+                .getWindowExtensions().getWindowLayoutComponent()));
+    }
+
+    @VisibleForTesting
+    SplitController(@NonNull WindowLayoutComponentImpl windowLayoutComponent) {
         final MainThreadExecutor executor = new MainThreadExecutor();
         mHandler = executor.mHandler;
         mPresenter = new SplitPresenter(executor, this);
-        ActivityThread activityThread = ActivityThread.currentActivityThread();
+        final ActivityThread activityThread = ActivityThread.currentActivityThread();
+        final Application application = activityThread.getApplication();
         // Register a callback to be notified about activities being created.
-        activityThread.getApplication().registerActivityLifecycleCallbacks(
-                new LifecycleCallbacks());
+        application.registerActivityLifecycleCallbacks(new LifecycleCallbacks());
         // Intercept activity starts to route activities to new containers if necessary.
         Instrumentation instrumentation = activityThread.getInstrumentation();
+
         mActivityStartMonitor = new ActivityStartMonitor();
         instrumentation.addMonitor(mActivityStartMonitor);
+        mWindowLayoutComponent = windowLayoutComponent;
+        mWindowLayoutComponent.addFoldingStateChangedCallback(new FoldingFeatureListener());
+    }
+
+    private class FoldingFeatureListener implements Consumer<List<CommonFoldingFeature>> {
+        @Override
+        public void accept(List<CommonFoldingFeature> foldingFeatures) {
+            synchronized (mLock) {
+                final WindowContainerTransaction wct = new WindowContainerTransaction();
+                for (int i = 0; i < mTaskContainers.size(); i++) {
+                    final TaskContainer taskContainer = mTaskContainers.valueAt(i);
+                    if (!taskContainer.isVisible()) {
+                        continue;
+                    }
+                    if (taskContainer.getDisplayId() != DEFAULT_DISPLAY) {
+                        continue;
+                    }
+                    // TODO(b/238948678): Support reporting display features in all windowing modes.
+                    if (taskContainer.isInMultiWindow()) {
+                        continue;
+                    }
+                    if (taskContainer.isEmpty()) {
+                        continue;
+                    }
+                    updateContainersInTask(wct, taskContainer);
+                    updateAnimationOverride(taskContainer);
+                }
+                mPresenter.applyTransaction(wct);
+            }
+        }
     }
 
     /** Updates the embedding rules applied to future activity launches. */
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
index f24401f..d783317 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
@@ -47,6 +47,7 @@
 import androidx.window.util.DataProducer;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -80,6 +81,11 @@
         mFoldingFeatureProducer.addDataChangedCallback(this::onDisplayFeaturesChanged);
     }
 
+    /** Registers to listen to {@link CommonFoldingFeature} changes */
+    public void addFoldingStateChangedCallback(Consumer<List<CommonFoldingFeature>> consumer) {
+        mFoldingFeatureProducer.addDataChangedCallback(consumer);
+    }
+
     /**
      * Adds a listener interested in receiving updates to {@link WindowLayoutInfo}
      *
@@ -225,12 +231,23 @@
      */
     private List<DisplayFeature> getDisplayFeatures(
             @NonNull @UiContext Context context, List<CommonFoldingFeature> storedFeatures) {
-        List<DisplayFeature> features = new ArrayList<>();
         if (!shouldReportDisplayFeatures(context)) {
+            return Collections.emptyList();
+        }
+        return getDisplayFeatures(context.getDisplayId(),
+                context.getResources().getConfiguration().windowConfiguration,
+                storedFeatures);
+    }
+
+    /** @see #getDisplayFeatures(Context, List) */
+    private List<DisplayFeature> getDisplayFeatures(int displayId,
+            @NonNull WindowConfiguration windowConfiguration,
+            List<CommonFoldingFeature> storedFeatures) {
+        List<DisplayFeature> features = new ArrayList<>();
+        if (displayId != DEFAULT_DISPLAY) {
             return features;
         }
 
-        int displayId = context.getDisplay().getDisplayId();
         for (CommonFoldingFeature baseFeature : storedFeatures) {
             Integer state = convertToExtensionState(baseFeature.getState());
             if (state == null) {
@@ -238,7 +255,7 @@
             }
             Rect featureRect = baseFeature.getRect();
             rotateRectToDisplayRotation(displayId, featureRect);
-            transformToWindowSpaceRect(context, featureRect);
+            transformToWindowSpaceRect(windowConfiguration, featureRect);
 
             if (!isZero(featureRect)) {
                 // TODO(b/228641877): Remove guarding when fixed.
@@ -263,6 +280,8 @@
             windowingMode = ActivityClient.getInstance().getTaskWindowingMode(
                     context.getActivityToken());
         } else {
+            // TODO(b/242674941): use task windowing mode for window context that associates with
+            //  activity.
             windowingMode = context.getResources().getConfiguration().windowConfiguration
                     .getWindowingMode();
         }
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/util/ExtensionHelper.java b/libs/WindowManager/Jetpack/src/androidx/window/util/ExtensionHelper.java
index 31bf963..9e2611f 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/util/ExtensionHelper.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/util/ExtensionHelper.java
@@ -21,6 +21,7 @@
 import static android.view.Surface.ROTATION_270;
 import static android.view.Surface.ROTATION_90;
 
+import android.app.WindowConfiguration;
 import android.content.Context;
 import android.graphics.Rect;
 import android.hardware.display.DisplayManagerGlobal;
@@ -89,13 +90,21 @@
     /** Transforms rectangle from absolute coordinate space to the window coordinate space. */
     public static void transformToWindowSpaceRect(@NonNull @UiContext Context context,
             Rect inOutRect) {
-        Rect windowRect = getWindowBounds(context);
-        if (!Rect.intersects(inOutRect, windowRect)) {
+        transformToWindowSpaceRect(getWindowBounds(context), inOutRect);
+    }
+
+    /** @see ExtensionHelper#transformToWindowSpaceRect(Context, Rect) */
+    public static void transformToWindowSpaceRect(@NonNull WindowConfiguration windowConfiguration,
+            Rect inOutRect) {
+        transformToWindowSpaceRect(windowConfiguration.getBounds(), inOutRect);
+    }
+
+    private static void transformToWindowSpaceRect(@NonNull Rect bounds, @NonNull Rect inOutRect) {
+        if (!inOutRect.intersect(bounds)) {
             inOutRect.setEmpty();
             return;
         }
-        inOutRect.intersect(windowRect);
-        inOutRect.offset(-windowRect.left, -windowRect.top);
+        inOutRect.offset(-bounds.left, -bounds.top);
     }
 
     /**