Introduce #createWindowContext with display

Test: atest ContextIsUiContextTest ContextGetDisplayTest
Test: atest WindowContextPolicyTests
Bug: 174640742
Change-Id: I13bd07fa3a4e79fe44bce34157ee93622cbb431d
diff --git a/core/api/current.txt b/core/api/current.txt
index 14c1dc87..222be7f0 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -10116,6 +10116,7 @@
     method public abstract android.content.Context createDisplayContext(@NonNull android.view.Display);
     method public abstract android.content.Context createPackageContext(String, int) throws android.content.pm.PackageManager.NameNotFoundException;
     method @NonNull public android.content.Context createWindowContext(int, @Nullable android.os.Bundle);
+    method @NonNull public android.content.Context createWindowContext(@NonNull android.view.Display, int, @Nullable android.os.Bundle);
     method public abstract String[] databaseList();
     method public abstract boolean deleteDatabase(String);
     method public abstract boolean deleteFile(String);
diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index 2fec9f7..15e2c0d 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -2470,8 +2470,9 @@
         return context;
     }
 
+    @NonNull
     @Override
-    public @NonNull WindowContext createWindowContext(int type, Bundle options) {
+    public WindowContext createWindowContext(int type, @NonNull Bundle options) {
         if (getDisplay() == null) {
             throw new UnsupportedOperationException("WindowContext can only be created from "
                     + "other visual contexts, such as Activity or one created with "
@@ -2480,13 +2481,26 @@
         return new WindowContext(this, type, options);
     }
 
-    ContextImpl createBaseWindowContext(IBinder token) {
+    @NonNull
+    @Override
+    public WindowContext createWindowContext(@NonNull Display display, int type,
+            @NonNull Bundle options) {
+        if (display == null) {
+            throw new IllegalArgumentException("Display must not be null");
+        }
+        return new WindowContext(this, display, type, options);
+    }
+
+    ContextImpl createBaseWindowContext(IBinder token, Display display) {
         ContextImpl context = new ContextImpl(this, mMainThread, mPackageInfo, mAttributionTag,
                 mSplitName, token, mUser, mFlags, mClassLoader, null);
         // Window contexts receive configurations directly from the server and as such do not
         // need to override their display in ResourcesManager.
         context.mForceDisplayOverrideInResources = false;
         context.mContextType = CONTEXT_TYPE_WINDOW_CONTEXT;
+        if (display != null) {
+            context.mDisplay = display;
+        }
         return context;
     }
 
diff --git a/core/java/android/app/WindowContext.java b/core/java/android/app/WindowContext.java
index 5f72bac..14ed414 100644
--- a/core/java/android/app/WindowContext.java
+++ b/core/java/android/app/WindowContext.java
@@ -26,6 +26,7 @@
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.RemoteException;
+import android.view.Display;
 import android.view.IWindowManager;
 import android.view.WindowManagerGlobal;
 import android.view.WindowManagerImpl;
@@ -59,13 +60,27 @@
      * @hide
      */
     public WindowContext(@NonNull Context base, int type, @Nullable Bundle options) {
+        this(base, null /* display */, type, options);
+    }
+
+    /**
+     * Default constructor. Will generate a {@link WindowTokenClient} and attach this context to
+     * the token.
+     *
+     * @param base Base {@link Context} for this new instance.
+     * @param display the {@link Display} to override.
+     * @param type Window type to be used with this context.
+     * @hide
+     */
+    public WindowContext(@NonNull Context base, @Nullable Display display, int type,
+            @Nullable Bundle options) {
         // Correct base context will be built once the token is resolved, so passing 'null' here.
         super(null /* base */);
 
         mWms = WindowManagerGlobal.getWindowManagerService();
         mToken = new WindowTokenClient();
 
-        final ContextImpl contextImpl = createBaseWindowContext(base, mToken);
+        final ContextImpl contextImpl = createBaseWindowContext(base, mToken, display);
         attachBaseContext(contextImpl);
         contextImpl.setOuterContext(this);
 
@@ -93,9 +108,10 @@
         Reference.reachabilityFence(this);
     }
 
-    private static ContextImpl createBaseWindowContext(Context outer, IBinder token) {
+    private static ContextImpl createBaseWindowContext(Context outer, IBinder token,
+            Display display) {
         final ContextImpl contextImpl = ContextImpl.getImpl(outer);
-        return contextImpl.createBaseWindowContext(token);
+        return contextImpl.createBaseWindowContext(token, display);
     }
 
     @Override
diff --git a/core/java/android/content/Context.java b/core/java/android/content/Context.java
index d920fb3..616194b 100644
--- a/core/java/android/content/Context.java
+++ b/core/java/android/content/Context.java
@@ -5968,22 +5968,22 @@
      * Creating a window context is an expensive operation. Misuse of this API may lead to a huge
      * performance drop. The best practice is to use the same window context when possible.
      * An approach is to create one window context with specific window type and display and
-     * use it everywhere it's needed..
+     * use it everywhere it's needed.
      * </p>
      *
      * @param type Window type in {@link WindowManager.LayoutParams}
-     * @param options Bundle used to pass window-related options.
-     * @return A {@link Context} that can be used to create windows.
-     * @throws UnsupportedOperationException if this is called on a non-UI context, such as
-     *         {@link android.app.Application Application} or {@link android.app.Service Service}.
+     * @param options A bundle used to pass window-related options
+     * @return A {@link Context} that can be used to create
+     *         non-{@link android.app.Activity activity} windows.
      *
      * @see #getSystemService(String)
      * @see #getSystemService(Class)
      * @see #WINDOW_SERVICE
      * @see #LAYOUT_INFLATER_SERVICE
      * @see #WALLPAPER_SERVICE
-     * @throws UnsupportedOperationException if this {@link Context} does not attach to a display or
-     * the current number of window contexts without adding any view by
+     * @throws UnsupportedOperationException if this {@link Context} does not attach to a display,
+     * such as {@link android.app.Application Application} or {@link android.app.Service Service},
+     * or the current number of window contexts without adding any view by
      * {@link WindowManager#addView} <b>exceeds five</b>.
      */
     @UiContext
@@ -5993,6 +5993,32 @@
     }
 
     /**
+     * A special version of {@link #createWindowContext(int, Bundle)} which also takes
+     * {@link Display}. The only difference between this API and
+     * {@link #createWindowContext(int, Bundle)} is that this API can create window context from
+     * any context even if the context which is not associated to a {@link Display} instance.
+     *
+     * @param display The {@link Display} to associate with
+     * @param type Window type in {@link WindowManager.LayoutParams}
+     * @param options A bundle used to pass window-related options.
+     * @return A {@link Context} that can be used to create
+     *         non-{@link android.app.Activity activity} windows.
+     * @throws IllegalArgumentException if the {@link Display} is {@code null}.
+     *
+     * @see #getSystemService(String)
+     * @see #getSystemService(Class)
+     * @see #WINDOW_SERVICE
+     * @see #LAYOUT_INFLATER_SERVICE
+     * @see #WALLPAPER_SERVICE
+     */
+    @UiContext
+    @NonNull
+    public Context createWindowContext(@NonNull Display display, @WindowType int type,
+            @Nullable Bundle options) {
+        throw new RuntimeException("Not implemented. Must override in a subclass.");
+    }
+
+    /**
      * Return a new Context object for the current Context but attribute to a different tag.
      * In complex apps attribution tagging can be used to distinguish between separate logical
      * parts.
diff --git a/core/java/android/content/ContextWrapper.java b/core/java/android/content/ContextWrapper.java
index 56da3cb..e450c08 100644
--- a/core/java/android/content/ContextWrapper.java
+++ b/core/java/android/content/ContextWrapper.java
@@ -988,6 +988,13 @@
     }
 
     @Override
+    @NonNull
+    public Context createWindowContext(@NonNull Display display, @WindowType int type,
+            @Nullable Bundle options) {
+        return mBase.createWindowContext(display, type, options);
+    }
+
+    @Override
     public @NonNull Context createAttributionContext(@Nullable String attributionTag) {
         return mBase.createAttributionContext(attributionTag);
     }
diff --git a/test-mock/src/android/test/mock/MockContext.java b/test-mock/src/android/test/mock/MockContext.java
index cf3b03c..f7cebd1 100644
--- a/test-mock/src/android/test/mock/MockContext.java
+++ b/test-mock/src/android/test/mock/MockContext.java
@@ -817,6 +817,11 @@
     }
 
     @Override
+    public @NonNull Context createWindowContext(Display display, int type, Bundle options) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
     public boolean isRestricted() {
         throw new UnsupportedOperationException();
     }