Respect UI Context for getDisplay

Before InputMethodService initializes, InputMethodService#getDisplay
will throw UnsupportedOperationException because it doesn't attach to
any display. However, it may make developers confuse about why
throwing exception for an UI context. This CL workarounds to also
check isUiContext flags to make the result of IMS#getDisplay consistent.
It should be considered to remove after IMS migrate to WindowContext
concept.

fixes: 161964893
Test: atest InputMethodServiceTest#testGetDisplayBeforeInitialization

Change-Id: Ie0bda707f0196b45a8851f645da7ecab6814bed6

core/java/android/app/ContextImpl.java

@@ -1914,10 +1914,8 @@ class ContextImpl extends Context {
     @Override
     public Object getSystemService(String name) {
         if (vmIncorrectContextUseEnabled()) {
-            // We may override this API from outer context.
-            final boolean isUiContext = isUiContext() || isOuterUiContext();
             // Check incorrect Context usage.
-            if (isUiComponent(name) && !isUiContext) {
+            if (isUiComponent(name) && !isSelfOrOuterUiContext()) {
                 final String errorMessage = "Tried to access visual service "
                         + SystemServiceRegistry.getSystemServiceClassName(name)
                         + " from a non-visual Context:" + getOuterContext();
@@ -1934,15 +1932,17 @@ class ContextImpl extends Context {
         return SystemServiceRegistry.getSystemService(this, name);
     }
 
-    private boolean isOuterUiContext() {
-        return getOuterContext() != null && getOuterContext().isUiContext();
-    }
-
     @Override
     public String getSystemServiceName(Class<?> serviceClass) {
         return SystemServiceRegistry.getSystemServiceName(serviceClass);
     }
 
+    // TODO(b/149463653): check if we still need this method after migrating IMS to WindowContext.
+    private boolean isSelfOrOuterUiContext() {
+        // We may override outer context's isUiContext
+        return isUiContext() || getOuterContext() != null && getOuterContext().isUiContext();
+    }
+
     /** @hide */
     @Override
     public boolean isUiContext() {
@@ -2389,7 +2389,7 @@ class ContextImpl extends Context {
         context.setResources(createResources(mToken, mPackageInfo, mSplitName, displayId,
                 overrideConfiguration, getDisplayAdjustments(displayId).getCompatibilityInfo(),
                 mResources.getLoaders()));
-        context.mIsUiContext = isUiContext() || isOuterUiContext();
+        context.mIsUiContext = isSelfOrOuterUiContext();
         return context;
     }
 
@@ -2494,9 +2494,9 @@ class ContextImpl extends Context {
 
     @Override
     public Display getDisplay() {
-        if (!mIsSystemOrSystemUiContext && !mIsAssociatedWithDisplay) {
+        if (!mIsSystemOrSystemUiContext && !mIsAssociatedWithDisplay && !isSelfOrOuterUiContext()) {
             throw new UnsupportedOperationException("Tried to obtain display from a Context not "
-                    + "associated with  one. Only visual Contexts (such as Activity or one created "
+                    + "associated with one. Only visual Contexts (such as Activity or one created "
                     + "with Context#createWindowContext) or ones created with "
                     + "Context#createDisplayContext are associated with displays. Other types of "
                     + "Contexts are typically related to background entities and may return an "