Hide/show ime switch icon properly

Bug: 5481999
Change-Id: I4b3dde7285a9ac49da8d5d1133fddd0807b03f4d
diff --git a/services/java/com/android/server/InputMethodManagerService.java b/services/java/com/android/server/InputMethodManagerService.java
index ddac35c..f5c4ed4 100644
--- a/services/java/com/android/server/InputMethodManagerService.java
+++ b/services/java/com/android/server/InputMethodManagerService.java
@@ -382,6 +382,7 @@
         public void onReceive(Context context, Intent intent) {
             if (intent.getAction().equals(Intent.ACTION_SCREEN_ON)) {
                 mScreenOn = true;
+                refreshImeWindowVisibilityLocked();
             } else if (intent.getAction().equals(Intent.ACTION_SCREEN_OFF)) {
                 mScreenOn = false;
                 setImeWindowVisibilityStatusHiddenLocked();
@@ -640,6 +641,21 @@
         updateImeWindowStatusLocked();
     }
 
+    private void refreshImeWindowVisibilityLocked() {
+        final Configuration conf = mRes.getConfiguration();
+        final boolean haveHardKeyboard = conf.keyboard
+                != Configuration.KEYBOARD_NOKEYS;
+        final boolean hardKeyShown = haveHardKeyboard
+                && conf.hardKeyboardHidden
+                        != Configuration.HARDKEYBOARDHIDDEN_YES;
+        final boolean isScreenLocked = mKeyguardManager != null
+                && mKeyguardManager.isKeyguardLocked()
+                && mKeyguardManager.isKeyguardSecure();
+        mImeWindowVis = (!isScreenLocked && (mInputShown || hardKeyShown)) ?
+                (InputMethodService.IME_ACTIVE | InputMethodService.IME_VISIBLE) : 0;
+        updateImeWindowStatusLocked();
+    }
+
     private void updateImeWindowStatusLocked() {
         setImeWindowStatus(mCurToken, mImeWindowVis, mBackDisposition);
     }
@@ -1285,16 +1301,7 @@
                     }
                     if (mCurMethod != null) {
                         try {
-                            final Configuration conf = mRes.getConfiguration();
-                            final boolean haveHardKeyboard = conf.keyboard
-                                    != Configuration.KEYBOARD_NOKEYS;
-                            final boolean hardKeyShown = haveHardKeyboard
-                                    && conf.hardKeyboardHidden
-                                            != Configuration.HARDKEYBOARDHIDDEN_YES;
-                            mImeWindowVis = (mInputShown || hardKeyShown) ? (
-                                    InputMethodService.IME_ACTIVE | InputMethodService.IME_VISIBLE)
-                                    : 0;
-                            updateImeWindowStatusLocked();
+                            refreshImeWindowVisibilityLocked();
                             // If subtype is null, try to find the most applicable one from
                             // getCurrentInputMethodSubtype.
                             if (subtype == null) {