Remove InternalInsetsType from InsetsSourceConsumer

This CL uses the ID to identify an insets source instead. The ID will be
created at runtime, so this CL removes the consumer from
InsetsController while the corresponding source is gone.

However, ImeInsetsSourceConsumer is a special case that it doesn't only
interact with the source and the control, but also InputMethodManager.
So this CL makes the IME consumer always exist.

Bug: 234093736
Test: atest ImeInsetsSourceConsumerTest InsetsAnimationControlImplTest
      InsetsControllerTest InsetsSourceConsumerTest
      WindowInsetsControllerTests
Change-Id: I4ce2663e79618e6634546b86c59092f46e0daef1

7 files changed, 162 insertions, 110 deletions

diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index e34aef97a8ab..b8cd7b9d177a 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -20,7 +20,6 @@ import static android.view.ImeInsetsSourceConsumerProto.INSETS_SOURCE_CONSUMER;
 import static android.view.ImeInsetsSourceConsumerProto.IS_HIDE_ANIMATION_RUNNING;
 import static android.view.ImeInsetsSourceConsumerProto.IS_REQUESTED_VISIBLE_AWAITING_CONTROL;
 import static android.view.ImeInsetsSourceConsumerProto.IS_SHOW_REQUESTED_DURING_HIDE_ANIMATION;
-import static android.view.InsetsState.ITYPE_IME;
 
 import android.annotation.Nullable;
 import android.os.IBinder;
@@ -55,9 +54,9 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     private boolean mIsShowRequestedDuringHideAnimation;
 
     public ImeInsetsSourceConsumer(
-            InsetsState state, Supplier<Transaction> transactionSupplier,
+            int id, InsetsState state, Supplier<Transaction> transactionSupplier,
             InsetsController controller) {
-        super(ITYPE_IME, state, transactionSupplier, controller);
+        super(id, WindowInsets.Type.ime(), state, transactionSupplier, controller);
     }
 
     @Override
@@ -137,7 +136,7 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
         // If we had a request before to show from IME (tracked with mImeRequestedShow), reaching
         // this code here means that we now got control, so we can start the animation immediately.
         // If client window is trying to control IME and IME is already visible, it is immediate.
-        if (fromIme || (mState.getSource(getInternalType()).isVisible() && getControl() != null)) {
+        if (fromIme || (mState.getSource(getId()).isVisible() && getControl() != null)) {
             return ShowResult.SHOW_IMMEDIATELY;
         }
 
diff --git a/core/java/android/view/InsetsController.java b/core/java/android/view/InsetsController.java
index e1b27ffe035d..c074e84e5158 100644
--- a/core/java/android/view/InsetsController.java
+++ b/core/java/android/view/InsetsController.java
@@ -21,7 +21,6 @@ import static android.view.InsetsControllerProto.CONTROL;
 import static android.view.InsetsControllerProto.STATE;
 import static android.view.InsetsState.ITYPE_CAPTION_BAR;
 import static android.view.InsetsState.ITYPE_IME;
-import static android.view.InsetsState.toPublicType;
 import static android.view.ViewRootImpl.CAPTION_ON_SHELL;
 import static android.view.WindowInsets.Type.FIRST;
 import static android.view.WindowInsets.Type.LAST;
@@ -569,8 +568,9 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
     private final InsetsState mLastDispatchedState = new InsetsState();
 
     private final Rect mFrame = new Rect();
-    private final BiFunction<InsetsController, Integer, InsetsSourceConsumer> mConsumerCreator;
+    private final BiFunction<InsetsController, InsetsSource, InsetsSourceConsumer> mConsumerCreator;
     private final SparseArray<InsetsSourceConsumer> mSourceConsumers = new SparseArray<>();
+    private final InsetsSourceConsumer mImeSourceConsumer;
     private final Host mHost;
     private final Handler mHandler;
 
@@ -621,19 +621,20 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
             this::invokeControllableInsetsChangedListeners;
 
     public InsetsController(Host host) {
-        this(host, (controller, type) -> {
-            if (type == ITYPE_IME) {
-                return new ImeInsetsSourceConsumer(controller.mState, Transaction::new, controller);
+        this(host, (controller, source) -> {
+            if (source.getType() == ime()) {
+                return new ImeInsetsSourceConsumer(source.getId(), controller.mState,
+                        Transaction::new, controller);
             } else {
-                return new InsetsSourceConsumer(type, controller.mState, Transaction::new,
-                        controller);
+                return new InsetsSourceConsumer(source.getId(), source.getType(), controller.mState,
+                        Transaction::new, controller);
             }
         }, host.getHandler());
     }
 
     @VisibleForTesting
     public InsetsController(Host host,
-            BiFunction<InsetsController, Integer, InsetsSourceConsumer> consumerCreator,
+            BiFunction<InsetsController, InsetsSource, InsetsSourceConsumer> consumerCreator,
             Handler handler) {
         mHost = host;
         mConsumerCreator = consumerCreator;
@@ -683,6 +684,9 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
                 dispatchAnimationEnd(finishedAnimations.get(i));
             }
         };
+
+        // Make mImeSourceConsumer always non-null.
+        mImeSourceConsumer = getSourceConsumer(new InsetsSource(ITYPE_IME, ime()));
     }
 
     @VisibleForTesting
@@ -741,28 +745,37 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
 
     private void updateState(InsetsState newState) {
         mState.set(newState, 0 /* types */);
+        for (int i = mSourceConsumers.size() - 1; i >= 0; i--) {
+            final InsetsSourceConsumer consumer = mSourceConsumers.valueAt(i);
+            final InsetsSource source = newState.peekSource(consumer.getId());
+            if (source == null && consumer != mImeSourceConsumer) {
+                // IME source consumer should always be there since we need to communicate with
+                // InputMethodManager no matter we have the source or not.
+                mSourceConsumers.removeAt(i);
+            }
+        }
         @InsetsType int existingTypes = 0;
         @InsetsType int visibleTypes = 0;
         @InsetsType int disabledUserAnimationTypes = 0;
         @InsetsType int[] cancelledUserAnimationTypes = {0};
-        for (@InternalInsetsType int type = 0; type < InsetsState.SIZE; type++) {
-            InsetsSource source = newState.peekSource(type);
+        for (int i = 0; i < InsetsState.SIZE; i++) {
+            InsetsSource source = newState.peekSource(i);
             if (source == null) continue;
-            @InsetsType int insetsType = toPublicType(type);
-            @AnimationType int animationType = getAnimationType(insetsType);
+            @InsetsType int type = source.getType();
+            @AnimationType int animationType = getAnimationType(type);
             if (!source.isUserControllable()) {
                 // The user animation is not allowed when visible frame is empty.
-                disabledUserAnimationTypes |= insetsType;
+                disabledUserAnimationTypes |= type;
                 if (animationType == ANIMATION_TYPE_USER) {
                     // Existing user animation needs to be cancelled.
                     animationType = ANIMATION_TYPE_NONE;
-                    cancelledUserAnimationTypes[0] |= insetsType;
+                    cancelledUserAnimationTypes[0] |= type;
                 }
             }
-            getSourceConsumer(type).updateSource(source, animationType);
-            existingTypes |= insetsType;
+            getSourceConsumer(source).updateSource(source, animationType);
+            existingTypes |= type;
             if (source.isVisible()) {
-                visibleTypes |= insetsType;
+                visibleTypes |= type;
             }
         }
 
@@ -899,25 +912,34 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
         }
 
         @InsetsType int controllableTypes = 0;
+        int consumedControlCount = 0;
         final int[] showTypes = new int[1];
         final int[] hideTypes = new int[1];
 
         // Ensure to update all existing source consumers
         for (int i = mSourceConsumers.size() - 1; i >= 0; i--) {
             final InsetsSourceConsumer consumer = mSourceConsumers.valueAt(i);
-            final InsetsSourceControl control = mTmpControlArray.get(consumer.getInternalType());
+            final InsetsSourceControl control = mTmpControlArray.get(consumer.getId());
+            if (control != null) {
+                controllableTypes |= control.getType();
+                consumedControlCount++;
+            }
 
             // control may be null, but we still need to update the control to null if it got
             // revoked.
             consumer.setControl(control, showTypes, hideTypes);
         }
 
-        // Ensure to create source consumers if not available yet.
-        for (int i = mTmpControlArray.size() - 1; i >= 0; i--) {
-            final InsetsSourceControl control = mTmpControlArray.valueAt(i);
-            final InsetsSourceConsumer consumer = getSourceConsumer(control.getId());
-            consumer.setControl(control, showTypes, hideTypes);
-            controllableTypes |= control.getType();
+        if (consumedControlCount != mTmpControlArray.size()) {
+            // Whoops! The server sent us some controls without sending corresponding sources.
+            for (int i = mTmpControlArray.size() - 1; i >= 0; i--) {
+                final InsetsSourceControl control = mTmpControlArray.valueAt(i);
+                final InsetsSourceConsumer consumer = mSourceConsumers.get(control.getId());
+                if (consumer == null) {
+                    control.release(SurfaceControl::release);
+                    Log.e(TAG, control + " has no consumer.");
+                }
+            }
         }
 
         if (mTmpControlArray.size() > 0) {
@@ -1144,11 +1166,11 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
             types &= ~mDisabledUserAnimationInsetsTypes;
 
             if (fromIme && (disabledTypes & ime()) != 0
-                    && !mState.getSource(ITYPE_IME).isVisible()) {
+                    && !mState.getSource(mImeSourceConsumer.getId()).isVisible()) {
                 // We've requested IMM to show IME, but the IME is not controllable. We need to
                 // cancel the request.
                 setRequestedVisibleTypes(0 /* visibleTypes */, ime());
-                if (getSourceConsumer(ITYPE_IME).onAnimationStateChanged(false /* running */)) {
+                if (mImeSourceConsumer.onAnimationStateChanged(false /* running */)) {
                     notifyVisibilityChanged();
                 }
             }
@@ -1163,11 +1185,10 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
         if (DEBUG) Log.d(TAG, "controlAnimation types: " + types);
         mLastStartedAnimTypes |= types;
 
-        final ArraySet<Integer> internalTypes = InsetsState.toInternalType(types);
         final SparseArray<InsetsSourceControl> controls = new SparseArray<>();
 
         Pair<Integer, Boolean> typesReadyPair = collectSourceControls(
-                fromIme, internalTypes, controls, animationType);
+                fromIme, types, controls, animationType);
         int typesReady = typesReadyPair.first;
         boolean imeReady = typesReadyPair.second;
         if (DEBUG) Log.d(TAG, String.format(
@@ -1257,13 +1278,15 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
     /**
      * @return Pair of (types ready to animate, IME ready to animate).
      */
-    private Pair<Integer, Boolean> collectSourceControls(boolean fromIme,
-            ArraySet<Integer> internalTypes, SparseArray<InsetsSourceControl> controls,
-            @AnimationType int animationType) {
+    private Pair<Integer, Boolean> collectSourceControls(boolean fromIme, @InsetsType int types,
+            SparseArray<InsetsSourceControl> controls, @AnimationType int animationType) {
         int typesReady = 0;
         boolean imeReady = true;
-        for (int i = internalTypes.size() - 1; i >= 0; i--) {
-            final InsetsSourceConsumer consumer = getSourceConsumer(internalTypes.valueAt(i));
+        for (int i = mSourceConsumers.size() - 1; i >= 0; i--) {
+            final InsetsSourceConsumer consumer = mSourceConsumers.valueAt(i);
+            if ((consumer.getType() & types) == 0) {
+                continue;
+            }
             boolean show = animationType == ANIMATION_TYPE_SHOW
                     || animationType == ANIMATION_TYPE_USER;
             boolean canRun = true;
@@ -1291,7 +1314,7 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
             if (!canRun) {
                 if (WARN) Log.w(TAG, String.format(
                         "collectSourceControls can't continue show for type: %s fromIme: %b",
-                        InsetsState.typeToString(consumer.getInternalType()), fromIme));
+                        WindowInsets.Type.toString(consumer.getType()), fromIme));
                 continue;
             }
             final InsetsSourceControl control = consumer.getControl();
@@ -1437,17 +1460,30 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
     }
 
     @VisibleForTesting
-    public @NonNull InsetsSourceConsumer getSourceConsumer(int id) {
-        InsetsSourceConsumer consumer = mSourceConsumers.get(id);
+    public @NonNull InsetsSourceConsumer getSourceConsumer(InsetsSource source) {
+        final int sourceId = source.getId();
+        InsetsSourceConsumer consumer = mSourceConsumers.get(sourceId);
         if (consumer != null) {
             return consumer;
         }
-        consumer = mConsumerCreator.apply(this, id);
-        mSourceConsumers.put(id, consumer);
+        if (source.getType() == ime() && mImeSourceConsumer != null) {
+            // WindowInsets.Type.ime() should be only provided by one source.
+            mSourceConsumers.remove(mImeSourceConsumer.getId());
+            consumer = mImeSourceConsumer;
+            consumer.setId(sourceId);
+        } else {
+            consumer = mConsumerCreator.apply(this, source);
+        }
+        mSourceConsumers.put(sourceId, consumer);
         return consumer;
     }
 
     @VisibleForTesting
+    public @NonNull InsetsSourceConsumer getImeSourceConsumer() {
+        return mImeSourceConsumer;
+    }
+
+    @VisibleForTesting
     public void notifyVisibilityChanged() {
         mHost.notifyInsetsChanged();
     }
@@ -1467,14 +1503,14 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
      * Called when current window gains focus.
      */
     public void onWindowFocusGained(boolean hasViewFocused) {
-        getSourceConsumer(ITYPE_IME).onWindowFocusGained(hasViewFocused);
+        mImeSourceConsumer.onWindowFocusGained(hasViewFocused);
     }
 
     /**
      * Called when current window loses focus.
      */
     public void onWindowFocusLost() {
-        getSourceConsumer(ITYPE_IME).onWindowFocusLost();
+        mImeSourceConsumer.onWindowFocusLost();
     }
 
     @VisibleForTesting
@@ -1518,13 +1554,12 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
         // TODO(b/166736352): We should only skip the animation of specific types, not all types.
         boolean skipAnim = false;
         if ((types & ime()) != 0) {
-            final InsetsSourceConsumer consumer = mSourceConsumers.get(ITYPE_IME);
-            final InsetsSourceControl imeControl = consumer != null ? consumer.getControl() : null;
+            final InsetsSourceControl imeControl = mImeSourceConsumer.getControl();
             // Skip showing animation once that made by system for some reason.
             // (e.g. starting window with IME snapshot)
             if (imeControl != null) {
                 skipAnim = imeControl.getAndClearSkipAnimationOnce() && show
-                        && consumer.hasViewFocusWhenWindowFocusGain();
+                        && mImeSourceConsumer.hasViewFocusWhenWindowFocusGain();
             }
         }
         applyAnimation(types, show, fromIme, skipAnim, statsToken);
@@ -1683,7 +1718,7 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
         @InsetsType int result = 0;
         for (int i = mSourceConsumers.size() - 1; i >= 0; i--) {
             InsetsSourceConsumer consumer = mSourceConsumers.valueAt(i);
-            InsetsSource source = mState.peekSource(consumer.getInternalType());
+            InsetsSource source = mState.peekSource(consumer.getId());
             if (consumer.getControl() != null && source != null && source.isUserControllable()) {
                 result |= consumer.getType();
             }
diff --git a/core/java/android/view/InsetsSourceConsumer.java b/core/java/android/view/InsetsSourceConsumer.java
index 7aeada8e9def..f46eb34bd397 100644
--- a/core/java/android/view/InsetsSourceConsumer.java
+++ b/core/java/android/view/InsetsSourceConsumer.java
@@ -25,7 +25,6 @@ import static android.view.InsetsSourceConsumerProto.IS_REQUESTED_VISIBLE;
 import static android.view.InsetsSourceConsumerProto.PENDING_FRAME;
 import static android.view.InsetsSourceConsumerProto.PENDING_VISIBLE_FRAME;
 import static android.view.InsetsSourceConsumerProto.SOURCE_CONTROL;
-import static android.view.InsetsState.getDefaultVisibility;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
 
@@ -34,7 +33,6 @@ import android.annotation.Nullable;
 import android.graphics.Rect;
 import android.util.Log;
 import android.util.proto.ProtoOutputStream;
-import android.view.InsetsState.InternalInsetsType;
 import android.view.SurfaceControl.Transaction;
 import android.view.WindowInsets.Type.InsetsType;
 
@@ -72,7 +70,7 @@ public class InsetsSourceConsumer {
 
     protected final InsetsController mController;
     protected final InsetsState mState;
-    private final @InternalInsetsType int mInternalType;
+    private int mId;
     private final @InsetsType int mType;
 
     private static final String TAG = "InsetsSourceConsumer";
@@ -88,16 +86,17 @@ public class InsetsSourceConsumer {
     private Rect mPendingVisibleFrame;
 
     /**
-     * @param type The {@link InternalInsetsType} of the consumed insets.
+     * @param id The ID of the consumed insets.
+     * @param type The {@link InsetsType} of the consumed insets.
      * @param state The current {@link InsetsState} of the consumed insets.
      * @param transactionSupplier The source of new {@link Transaction} instances. The supplier
      *         must provide *new* instances, which will be explicitly closed by this class.
      * @param controller The {@link InsetsController} to use for insets interaction.
      */
-    public InsetsSourceConsumer(@InternalInsetsType int type, InsetsState state,
+    public InsetsSourceConsumer(int id, @InsetsType int type, InsetsState state,
             Supplier<Transaction> transactionSupplier, InsetsController controller) {
-        mType = InsetsState.toPublicType(type);
-        mInternalType = type;
+        mId = id;
+        mType = type;
         mState = state;
         mTransactionSupplier = transactionSupplier;
         mController = controller;
@@ -134,12 +133,11 @@ public class InsetsSourceConsumer {
             mController.notifyControlRevoked(this);
 
             // Check if we need to restore server visibility.
-            final InsetsSource source = mState.getSource(mInternalType);
-            final boolean serverVisibility =
-                    mController.getLastDispatchedState().getSourceOrDefaultVisibility(
-                            mInternalType);
-            if (source.isVisible() != serverVisibility) {
-                source.setVisible(serverVisibility);
+            final InsetsSource localSource = mState.peekSource(mId);
+            final InsetsSource serverSource = mController.getLastDispatchedState().peekSource(mId);
+            if (localSource != null && serverSource != null
+                    && localSource.isVisible() != serverSource.isVisible()) {
+                localSource.setVisible(serverSource.isVisible());
                 mController.notifyVisibilityChanged();
             }
         } else {
@@ -196,12 +194,16 @@ public class InsetsSourceConsumer {
         return (mController.getRequestedVisibleTypes() & mType) != 0;
     }
 
-    @InsetsType int getType() {
-        return mType;
+    int getId() {
+        return mId;
     }
 
-    @InternalInsetsType int getInternalType() {
-        return mInternalType;
+    void setId(int id) {
+        mId = id;
+    }
+
+    @InsetsType int getType() {
+        return mType;
     }
 
     /**
@@ -211,12 +213,14 @@ public class InsetsSourceConsumer {
     public boolean onAnimationStateChanged(boolean running) {
         boolean insetsChanged = false;
         if (!running && mPendingFrame != null) {
-            InsetsSource source = mState.getSource(mInternalType);
-            source.setFrame(mPendingFrame);
-            source.setVisibleFrame(mPendingVisibleFrame);
+            final InsetsSource source = mState.peekSource(mId);
+            if (source != null) {
+                source.setFrame(mPendingFrame);
+                source.setVisibleFrame(mPendingVisibleFrame);
+                insetsChanged = true;
+            }
             mPendingFrame = null;
             mPendingVisibleFrame = null;
-            insetsChanged = true;
         }
 
         // We apply the visibility override after the animation is started. We don't do this before
@@ -248,25 +252,25 @@ public class InsetsSourceConsumer {
 
     @VisibleForTesting(visibility = PACKAGE)
     public boolean applyLocalVisibilityOverride() {
-        final InsetsSource source = mState.peekSource(mInternalType);
-        final boolean isVisible = source != null ? source.isVisible() : getDefaultVisibility(
-                mInternalType);
-        final boolean hasControl = mSourceControl != null;
+        final InsetsSource source = mState.peekSource(mId);
+        if (source == null) {
+            return false;
+        }
         final boolean requestedVisible = (mController.getRequestedVisibleTypes() & mType) != 0;
 
         // If we don't have control, we are not able to change the visibility.
-        if (!hasControl) {
+        if (mSourceControl == null) {
             if (DEBUG) Log.d(TAG, "applyLocalVisibilityOverride: No control in "
                     + mController.getHost().getRootViewTitle()
                     + " requestedVisible=" + requestedVisible);
             return false;
         }
-        if (isVisible == requestedVisible) {
+        if (source.isVisible() == requestedVisible) {
             return false;
         }
         if (DEBUG) Log.d(TAG, String.format("applyLocalVisibilityOverride: %s requestedVisible: %b",
                 mController.getHost().getRootViewTitle(), requestedVisible));
-        mState.getSource(mInternalType).setVisible(requestedVisible);
+        source.setVisible(requestedVisible);
         return true;
     }
 
@@ -301,7 +305,7 @@ public class InsetsSourceConsumer {
 
     @VisibleForTesting(visibility = PACKAGE)
     public void updateSource(InsetsSource newSource, @AnimationType int animationType) {
-        InsetsSource source = mState.peekSource(mInternalType);
+        InsetsSource source = mState.peekSource(mId);
         if (source == null || animationType == ANIMATION_TYPE_NONE
                 || source.getFrame().equals(newSource.getFrame())) {
             mPendingFrame = null;
@@ -345,7 +349,7 @@ public class InsetsSourceConsumer {
 
     void dumpDebug(ProtoOutputStream proto, long fieldId) {
         final long token = proto.start(fieldId);
-        proto.write(INTERNAL_INSETS_TYPE, InsetsState.typeToString(mInternalType));
+        proto.write(INTERNAL_INSETS_TYPE, WindowInsets.Type.toString(mType));
         proto.write(HAS_WINDOW_FOCUS, mHasWindowFocus);
         proto.write(IS_REQUESTED_VISIBLE, (mController.getRequestedVisibleTypes() & mType) != 0);
         if (mSourceControl != null) {
diff --git a/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
index 9b8a0e952beb..958fdc665cb2 100644
--- a/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
@@ -60,7 +60,7 @@ import org.mockito.Spy;
 public class ImeInsetsSourceConsumerTest {
 
     Context mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
-    ImeInsetsSourceConsumer mImeConsumer;
+    InsetsSourceConsumer mImeConsumer;
     @Spy InsetsController mController;
     SurfaceControl mLeash;
 
@@ -86,7 +86,7 @@ public class ImeInsetsSourceConsumerTest {
                     false,
                     TYPE_APPLICATION, WINDOWING_MODE_UNDEFINED,
                     SOFT_INPUT_ADJUST_RESIZE, 0, 0);
-            mImeConsumer = (ImeInsetsSourceConsumer) mController.getSourceConsumer(ITYPE_IME);
+            mImeConsumer = mController.getImeSourceConsumer();
         });
     }
 
diff --git a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
index 597df0ba4949..cc5f7f87636c 100644
--- a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
+++ b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
@@ -89,7 +89,8 @@ public class InsetsAnimationControlImplTest {
         mInsetsState = new InsetsState();
         mInsetsState.getSource(ITYPE_STATUS_BAR).setFrame(new Rect(0, 0, 500, 100));
         mInsetsState.getSource(ITYPE_NAVIGATION_BAR).setFrame(new Rect(400, 0, 500, 500));
-        InsetsSourceConsumer topConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, mInsetsState,
+        InsetsSourceConsumer topConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR,
+                WindowInsets.Type.statusBars(), mInsetsState,
                 () -> mMockTransaction, mMockController);
         topConsumer.setControl(
                 new InsetsSourceControl(ITYPE_STATUS_BAR, WindowInsets.Type.statusBars(),
@@ -97,7 +98,8 @@ public class InsetsAnimationControlImplTest {
                 new int[1], new int[1]);
 
         InsetsSourceConsumer navConsumer = new InsetsSourceConsumer(ITYPE_NAVIGATION_BAR,
-                mInsetsState, () -> mMockTransaction, mMockController);
+                WindowInsets.Type.navigationBars(), mInsetsState,
+                () -> mMockTransaction, mMockController);
         navConsumer.setControl(
                 new InsetsSourceControl(ITYPE_NAVIGATION_BAR, WindowInsets.Type.navigationBars(),
                         mNavLeash, true, new Point(400, 0), Insets.of(0, 0, 100, 0)),
diff --git a/core/tests/coretests/src/android/view/InsetsControllerTest.java b/core/tests/coretests/src/android/view/InsetsControllerTest.java
index 88249ad0db4c..c917302443ae 100644
--- a/core/tests/coretests/src/android/view/InsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsControllerTest.java
@@ -100,6 +100,9 @@ import java.util.concurrent.CountDownLatch;
 @Presubmit
 @RunWith(AndroidJUnit4.class)
 public class InsetsControllerTest {
+    private InsetsSource mStatusSource;
+    private InsetsSource mNavSource;
+    private InsetsSource mImeSource;
     private InsetsController mController;
     private SurfaceSession mSession = new SurfaceSession();
     private SurfaceControl mLeash;
@@ -125,10 +128,10 @@ public class InsetsControllerTest {
             mTestClock = new OffsettableClock();
             mTestHandler = new TestHandler(null, mTestClock);
             mTestHost = spy(new TestHost(mViewRoot));
-            mController = new InsetsController(mTestHost, (controller, type) -> {
-                if (type == ITYPE_IME) {
-                    return new InsetsSourceConsumer(type, controller.getState(),
-                            Transaction::new, controller) {
+            mController = new InsetsController(mTestHost, (controller, source) -> {
+                if (source.getType() == ime()) {
+                    return new InsetsSourceConsumer(source.getId(), source.getType(),
+                            controller.getState(), Transaction::new, controller) {
 
                         private boolean mImeRequestedShow;
 
@@ -143,18 +146,25 @@ public class InsetsControllerTest {
                         }
                     };
                 } else {
-                    return new InsetsSourceConsumer(type, controller.getState(), Transaction::new,
-                            controller);
+                    return new InsetsSourceConsumer(source.getId(), source.getType(),
+                            controller.getState(), Transaction::new, controller);
                 }
             }, mTestHandler);
             final Rect rect = new Rect(5, 5, 5, 5);
-            mController.getState().getSource(ITYPE_STATUS_BAR).setFrame(new Rect(0, 0, 100, 10));
-            mController.getState().getSource(ITYPE_NAVIGATION_BAR).setFrame(
-                    new Rect(0, 90, 100, 100));
-            mController.getState().getSource(ITYPE_IME).setFrame(new Rect(0, 50, 100, 100));
-            mController.getState().setDisplayFrame(new Rect(0, 0, 100, 100));
-            mController.getState().setDisplayCutout(new DisplayCutout(
+            mStatusSource = new InsetsSource(ITYPE_STATUS_BAR, statusBars());
+            mStatusSource.setFrame(new Rect(0, 0, 100, 10));
+            mNavSource = new InsetsSource(ITYPE_NAVIGATION_BAR, navigationBars());
+            mNavSource.setFrame(new Rect(0, 90, 100, 100));
+            mImeSource = new InsetsSource(ITYPE_IME, ime());
+            mImeSource.setFrame(new Rect(0, 0, 100, 10));
+            InsetsState state = new InsetsState();
+            state.addSource(mStatusSource);
+            state.addSource(mNavSource);
+            state.addSource(mImeSource);
+            state.setDisplayFrame(new Rect(0, 0, 100, 100));
+            state.setDisplayCutout(new DisplayCutout(
                     Insets.of(10, 10, 10, 10), rect, rect, rect, rect));
+            mController.onStateChanged(state);
             mController.calculateInsets(
                     false,
                     false,
@@ -168,7 +178,7 @@ public class InsetsControllerTest {
     @Test
     public void testControlsChanged() {
         mController.onControlsChanged(createSingletonControl(ITYPE_STATUS_BAR, statusBars()));
-        assertNotNull(mController.getSourceConsumer(ITYPE_STATUS_BAR).getControl().getLeash());
+        assertNotNull(mController.getSourceConsumer(mStatusSource).getControl().getLeash());
         mController.addOnControllableInsetsChangedListener(
                 ((controller, typeMask) -> assertEquals(statusBars(), typeMask)));
     }
@@ -180,7 +190,7 @@ public class InsetsControllerTest {
         mController.addOnControllableInsetsChangedListener(listener);
         mController.onControlsChanged(createSingletonControl(ITYPE_STATUS_BAR, statusBars()));
         mController.onControlsChanged(new InsetsSourceControl[0]);
-        assertNull(mController.getSourceConsumer(ITYPE_STATUS_BAR).getControl());
+        assertNull(mController.getSourceConsumer(mStatusSource).getControl());
         InOrder inOrder = Mockito.inOrder(listener);
         inOrder.verify(listener).onControllableInsetsChanged(eq(mController), eq(0));
         inOrder.verify(listener).onControllableInsetsChanged(eq(mController), eq(statusBars()));
@@ -240,7 +250,7 @@ public class InsetsControllerTest {
             WindowInsetsAnimationControlListener loggingListener =
                     mock(WindowInsetsAnimationControlListener.class);
             mController.setSystemDrivenInsetsAnimationLoggingListener(loggingListener);
-            mController.getSourceConsumer(ITYPE_IME).onWindowFocusGained(true);
+            mController.getSourceConsumer(mImeSource).onWindowFocusGained(true);
             // since there is no focused view, forcefully make IME visible.
             mController.show(WindowInsets.Type.ime(), true /* fromIme */, null /* statsToken */);
             verify(loggingListener).onReady(notNull(), anyInt());
@@ -252,7 +262,7 @@ public class InsetsControllerTest {
         prepareControls();
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
-            mController.getSourceConsumer(ITYPE_IME).onWindowFocusGained(true);
+            mController.getSourceConsumer(mImeSource).onWindowFocusGained(true);
             // since there is no focused view, forcefully make IME visible.
             mController.show(WindowInsets.Type.ime(), true /* fromIme */, null /* statsToken */);
             mController.show(all());
@@ -265,7 +275,7 @@ public class InsetsControllerTest {
             mController.hide(all());
             mController.cancelExistingAnimations();
             assertEquals(0, mController.getRequestedVisibleTypes() & types);
-            mController.getSourceConsumer(ITYPE_IME).onWindowFocusLost();
+            mController.getSourceConsumer(mImeSource).onWindowFocusLost();
         });
         InstrumentationRegistry.getInstrumentation().waitForIdleSync();
     }
@@ -275,14 +285,14 @@ public class InsetsControllerTest {
         InsetsSourceControl ime = createControl(ITYPE_IME, ime());
         mController.onControlsChanged(new InsetsSourceControl[] { ime });
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
-            mController.getSourceConsumer(ITYPE_IME).onWindowFocusGained(true);
+            mController.getSourceConsumer(mImeSource).onWindowFocusGained(true);
             mController.show(WindowInsets.Type.ime(), true /* fromIme */, null /* statsToken */);
             mController.cancelExistingAnimations();
             assertTrue(isRequestedVisible(mController, ime()));
             mController.hide(ime(), true /* fromIme */, null /* statsToken */);
             mController.cancelExistingAnimations();
             assertFalse(isRequestedVisible(mController, ime()));
-            mController.getSourceConsumer(ITYPE_IME).onWindowFocusLost();
+            mController.getSourceConsumer(mImeSource).onWindowFocusLost();
         });
         InstrumentationRegistry.getInstrumentation().waitForIdleSync();
     }
@@ -901,7 +911,8 @@ public class InsetsControllerTest {
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
             // Simulate IME insets is not controllable
             mController.onControlsChanged(new InsetsSourceControl[0]);
-            final InsetsSourceConsumer imeInsetsConsumer = mController.getSourceConsumer(ITYPE_IME);
+            final InsetsSourceConsumer imeInsetsConsumer =
+                    mController.getSourceConsumer(mImeSource);
             assertNull(imeInsetsConsumer.getControl());
 
             // Verify IME requested visibility should be updated to IME consumer from controller.
diff --git a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
index cec3fc5ba2d8..3a3eeee77d8d 100644
--- a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
@@ -100,7 +100,7 @@ public class InsetsSourceConsumerTest {
             mState.addSource(mSpyInsetsSource);
 
             mController = new InsetsController(new ViewRootInsetsControllerHost(mViewRoot));
-            mConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, mState,
+            mConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, statusBars(), mState,
                     () -> mMockTransaction, mController) {
                 @Override
                 public void removeSurface() {
@@ -146,7 +146,7 @@ public class InsetsSourceConsumerTest {
         InsetsController controller = new InsetsController(new ViewRootInsetsControllerHost(
                 mViewRoot));
         InsetsSourceConsumer consumer = new InsetsSourceConsumer(
-                ITYPE_IME, state, null, controller);
+                ITYPE_IME, ime(), state, null, controller);
 
         InsetsSource source = new InsetsSource(ITYPE_IME, ime());
         source.setFrame(0, 1, 2, 3);
@@ -216,9 +216,9 @@ public class InsetsSourceConsumerTest {
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
             InsetsState state = new InsetsState();
             ViewRootInsetsControllerHost host = new ViewRootInsetsControllerHost(mViewRoot);
-            InsetsController insetsController = new InsetsController(host, (controller, type) -> {
-                if (type == ITYPE_IME) {
-                    return new InsetsSourceConsumer(ITYPE_IME, state,
+            InsetsController insetsController = new InsetsController(host, (controller, source) -> {
+                if (source.getType() == ime()) {
+                    return new InsetsSourceConsumer(ITYPE_IME, ime(), state,
                             () -> mMockTransaction, controller) {
                         @Override
                         public int requestShow(boolean fromController) {
@@ -226,10 +226,11 @@ public class InsetsSourceConsumerTest {
                         }
                     };
                 }
-                return new InsetsSourceConsumer(type, controller.getState(), Transaction::new,
-                        controller);
+                return new InsetsSourceConsumer(source.getId(), source.getType(),
+                        controller.getState(), Transaction::new, controller);
             }, host.getHandler());
-            InsetsSourceConsumer imeConsumer = insetsController.getSourceConsumer(ITYPE_IME);
+            InsetsSource imeSource = new InsetsSource(ITYPE_IME, ime());
+            InsetsSourceConsumer imeConsumer = insetsController.getSourceConsumer(imeSource);
 
             // Initial IME insets source control with its leash.
             imeConsumer.setControl(new InsetsSourceControl(ITYPE_IME, ime(), mLeash,