Merge "Support ContextWrapper by replacing instanceof checks with token checks" into tm-qpr-dev
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 b70b320..9b16877 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.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 @@
         });
         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 @@
             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 @@
             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'