Fix IME not shown from recents to split-screen task

DislplayImeController only update the ImeSurfaceControl
for a non-null control. So it still keep a stale control
even the control has lost.

The ImeSurfaceControl should be updated to avoid using a
stale control to start animation.

Bug: 258814082
Test: Manual test with bug steps
Test: atest DisplayImeControllerTest
Change-Id: Id7ac6de6a9e8e5f0e620254c39d104272175b3a1
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
index d9b4f47..4edc642 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
@@ -295,8 +295,8 @@
                         mImeSourceControl.release(SurfaceControl::release);
                     }
                 }
-                mImeSourceControl = imeSourceControl;
             }
+            mImeSourceControl = imeSourceControl;
         }
 
         private void applyVisibilityToLeash(InsetsSourceControl imeSourceControl) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
index 40f2e88..8641541 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
@@ -21,7 +21,9 @@
 import static android.view.Surface.ROTATION_0;
 import static android.view.WindowInsets.Type.ime;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
@@ -90,6 +92,22 @@
     }
 
     @Test
+    public void insetsControlChanged_updateExpectedImeSourceControl() {
+        final InsetsSourceControl[] insetsSourceControls = new InsetsSourceControl[]{
+                new InsetsSourceControl(ITYPE_IME, mock(SurfaceControl.class), false,
+                        new Point(0, 0), Insets.NONE)};
+        final InsetsSourceControl imeSourceControl = insetsSourceControls[0];
+
+        mPerDisplay.insetsControlChanged(insetsStateWithIme(false), insetsSourceControls);
+
+        assertEquals(imeSourceControl, mPerDisplay.mImeSourceControl);
+
+        mPerDisplay.insetsControlChanged(insetsStateWithIme(false), null);
+
+        assertNull(mPerDisplay.mImeSourceControl);
+    }
+
+    @Test
     public void insetsChanged_schedulesNoWorkOnExecutor() {
         mPerDisplay.insetsChanged(insetsStateWithIme(false));
         verifyZeroInteractions(mExecutor);