Support ContextWrapper by replacing instanceof checks with token checks

Test: all tests pass

Bug: 258065175

Change-Id: I3dd08929c0ba420d5fa71f0e0fdbc8ca7cc8abb5

1 files changed, 9 insertions, 8 deletions

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 b70b320eee3c..9b16877e39e7 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
@@ -35,7 +35,6 @@ import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.util.ArrayMap;
-import android.window.WindowProvider;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -129,9 +128,10 @@ public class WindowLayoutComponentImpl implements WindowLayoutComponent {
         });
         mWindowLayoutChangeListeners.put(context, consumer);
 
-        // TODO(b/258065175) Further extend this to ContextWrappers.
-        if (context instanceof WindowProvider) {
-            final IBinder windowContextToken = context.getWindowContextToken();
+        final IBinder windowContextToken = context.getWindowContextToken();
+        if (windowContextToken != null) {
+            // We register component callbacks for window contexts. For activity contexts, they will
+            // receive callbacks from NotifyOnConfigurationChanged instead.
             final ConfigurationChangeListener listener =
                     new ConfigurationChangeListener(windowContextToken);
             context.registerComponentCallbacks(listener);
@@ -150,8 +150,8 @@ public class WindowLayoutComponentImpl implements WindowLayoutComponent {
             if (!mWindowLayoutChangeListeners.get(context).equals(consumer)) {
                 continue;
             }
-            if (context instanceof WindowProvider) {
-                final IBinder token = context.getWindowContextToken();
+            final IBinder token = context.getWindowContextToken();
+            if (token != null) {
                 context.unregisterComponentCallbacks(mConfigurationChangeListeners.get(token));
                 mConfigurationChangeListeners.remove(token);
             }
@@ -308,9 +308,10 @@ public class WindowLayoutComponentImpl implements WindowLayoutComponent {
             return false;
         }
         final int windowingMode;
-        if (context instanceof Activity) {
+        IBinder activityToken = context.getActivityToken();
+        if (activityToken != null) {
             final Configuration taskConfig = ActivityClient.getInstance().getTaskConfiguration(
-                    context.getActivityToken());
+                    activityToken);
             if (taskConfig == null) {
                 // If we cannot determine the task configuration for any reason, it is likely that
                 // we won't be able to determine its position correctly as well. DisplayFeatures'