Merge "Fix led meta state can't be cleared by another keyboard" into tm-dev
diff --git a/cmds/installd/TEST_MAPPING b/cmds/installd/TEST_MAPPING
index 8ccab4c..fc4cfc9 100644
--- a/cmds/installd/TEST_MAPPING
+++ b/cmds/installd/TEST_MAPPING
@@ -32,7 +32,12 @@
       "name": "CtsCompilationTestCases"
     },
     {
-      "name": "SdkSandboxStorageHostTest"
+      "name": "SdkSandboxStorageHostTest",
+      "options": [
+        {
+          "exclude-annotation": "android.platform.test.annotations.LargeTest"
+        }
+      ]
     }
   ]
 }
diff --git a/data/etc/Android.bp b/data/etc/Android.bp
index 5003151..bdd5172 100644
--- a/data/etc/Android.bp
+++ b/data/etc/Android.bp
@@ -227,6 +227,12 @@
 }
 
 prebuilt_etc {
+    name: "android.software.opengles.deqp.level-2022-03-01.prebuilt.xml",
+    src: "android.software.opengles.deqp.level-2022-03-01.xml",
+    defaults: ["frameworks_native_data_etc_defaults"],
+}
+
+prebuilt_etc {
     name: "android.software.sip.voip.prebuilt.xml",
     src: "android.software.sip.voip.xml",
     defaults: ["frameworks_native_data_etc_defaults"],
@@ -245,6 +251,12 @@
 }
 
 prebuilt_etc {
+    name: "android.software.vulkan.deqp.level-2022-03-01.prebuilt.xml",
+    src: "android.software.vulkan.deqp.level-2022-03-01.xml",
+    defaults: ["frameworks_native_data_etc_defaults"],
+}
+
+prebuilt_etc {
     name: "aosp_excluded_hardware.prebuilt.xml",
     src: "aosp_excluded_hardware.xml",
     defaults: ["frameworks_native_data_etc_defaults"],
diff --git a/services/inputflinger/reader/mapper/CursorInputMapper.cpp b/services/inputflinger/reader/mapper/CursorInputMapper.cpp
index dc5fcec..a9a4c71 100644
--- a/services/inputflinger/reader/mapper/CursorInputMapper.cpp
+++ b/services/inputflinger/reader/mapper/CursorInputMapper.cpp
@@ -27,6 +27,9 @@
 
 namespace android {
 
+// The default velocity control parameters that has no effect.
+static const VelocityControlParameters FLAT_VELOCITY_CONTROL_PARAMS{};
+
 // --- CursorMotionAccumulator ---
 
 CursorMotionAccumulator::CursorMotionAccumulator() {
@@ -154,8 +157,9 @@
         mHWheelScale = 1.0f;
     }
 
-    if ((!changes && config->pointerCaptureRequest.enable) ||
-        (changes & InputReaderConfiguration::CHANGE_POINTER_CAPTURE)) {
+    const bool configurePointerCapture = (!changes && config->pointerCaptureRequest.enable) ||
+            (changes & InputReaderConfiguration::CHANGE_POINTER_CAPTURE);
+    if (configurePointerCapture) {
         if (config->pointerCaptureRequest.enable) {
             if (mParameters.mode == Parameters::MODE_POINTER) {
                 mParameters.mode = Parameters::MODE_POINTER_RELATIVE;
@@ -180,10 +184,18 @@
         }
     }
 
-    if (!changes || (changes & InputReaderConfiguration::CHANGE_POINTER_SPEED)) {
-        mPointerVelocityControl.setParameters(config->pointerVelocityControlParameters);
-        mWheelXVelocityControl.setParameters(config->wheelVelocityControlParameters);
-        mWheelYVelocityControl.setParameters(config->wheelVelocityControlParameters);
+    if (!changes || (changes & InputReaderConfiguration::CHANGE_POINTER_SPEED) ||
+        configurePointerCapture) {
+        if (config->pointerCaptureRequest.enable) {
+            // Disable any acceleration or scaling when Pointer Capture is enabled.
+            mPointerVelocityControl.setParameters(FLAT_VELOCITY_CONTROL_PARAMS);
+            mWheelXVelocityControl.setParameters(FLAT_VELOCITY_CONTROL_PARAMS);
+            mWheelYVelocityControl.setParameters(FLAT_VELOCITY_CONTROL_PARAMS);
+        } else {
+            mPointerVelocityControl.setParameters(config->pointerVelocityControlParameters);
+            mWheelXVelocityControl.setParameters(config->wheelVelocityControlParameters);
+            mWheelYVelocityControl.setParameters(config->wheelVelocityControlParameters);
+        }
     }
 
     if (!changes || (changes & InputReaderConfiguration::CHANGE_DISPLAY_INFO)) {
diff --git a/services/inputflinger/tests/InputReader_test.cpp b/services/inputflinger/tests/InputReader_test.cpp
index bda7755..b29a0a2 100644
--- a/services/inputflinger/tests/InputReader_test.cpp
+++ b/services/inputflinger/tests/InputReader_test.cpp
@@ -375,6 +375,11 @@
 
     float getPointerGestureMovementSpeedRatio() { return mConfig.pointerGestureMovementSpeedRatio; }
 
+    void setVelocityControlParams(const VelocityControlParameters& params) {
+        mConfig.pointerVelocityControlParameters = params;
+        mConfig.wheelVelocityControlParameters = params;
+    }
+
 private:
     uint32_t mNextPointerCaptureSequenceNumber = 0;
 
@@ -2949,7 +2954,10 @@
     }
 
     void configureDevice(uint32_t changes) {
-        if (!changes || (changes & InputReaderConfiguration::CHANGE_DISPLAY_INFO)) {
+        if (!changes ||
+            (changes &
+             (InputReaderConfiguration::CHANGE_DISPLAY_INFO |
+              InputReaderConfiguration::CHANGE_POINTER_CAPTURE))) {
             mReader->requestRefreshConfiguration(changes);
             mReader->loopOnce();
         }
@@ -4914,6 +4922,54 @@
     ASSERT_NO_FATAL_FAILURE(assertPosition(*mFakePointerController, 110.0f, 220.0f));
 }
 
+/**
+ * When Pointer Capture is enabled, we expect to report unprocessed relative movements, so any
+ * pointer acceleration or speed processing should not be applied.
+ */
+TEST_F(CursorInputMapperTest, PointerCaptureDisablesVelocityProcessing) {
+    addConfigurationProperty("cursor.mode", "pointer");
+    const VelocityControlParameters testParams(5.f /*scale*/, 0.f /*low threshold*/,
+                                               100.f /*high threshold*/, 10.f /*acceleration*/);
+    mFakePolicy->setVelocityControlParams(testParams);
+    CursorInputMapper& mapper = addMapperAndConfigure<CursorInputMapper>();
+
+    NotifyDeviceResetArgs resetArgs;
+    ASSERT_NO_FATAL_FAILURE(mFakeListener->assertNotifyDeviceResetWasCalled(&resetArgs));
+    ASSERT_EQ(ARBITRARY_TIME, resetArgs.eventTime);
+    ASSERT_EQ(DEVICE_ID, resetArgs.deviceId);
+
+    NotifyMotionArgs args;
+
+    // Move and verify scale is applied.
+    process(mapper, ARBITRARY_TIME, READ_TIME, EV_REL, REL_X, 10);
+    process(mapper, ARBITRARY_TIME, READ_TIME, EV_REL, REL_Y, 20);
+    process(mapper, ARBITRARY_TIME, READ_TIME, EV_SYN, SYN_REPORT, 0);
+    ASSERT_NO_FATAL_FAILURE(mFakeListener->assertNotifyMotionWasCalled(&args));
+    ASSERT_EQ(AINPUT_SOURCE_MOUSE, args.source);
+    ASSERT_EQ(AMOTION_EVENT_ACTION_HOVER_MOVE, args.action);
+    const float relX = args.pointerCoords[0].getAxisValue(AMOTION_EVENT_AXIS_RELATIVE_X);
+    const float relY = args.pointerCoords[0].getAxisValue(AMOTION_EVENT_AXIS_RELATIVE_Y);
+    ASSERT_GT(relX, 10);
+    ASSERT_GT(relY, 20);
+
+    // Enable Pointer Capture
+    mFakePolicy->setPointerCapture(true);
+    configureDevice(InputReaderConfiguration::CHANGE_POINTER_CAPTURE);
+    NotifyPointerCaptureChangedArgs captureArgs;
+    ASSERT_NO_FATAL_FAILURE(mFakeListener->assertNotifyCaptureWasCalled(&captureArgs));
+    ASSERT_TRUE(captureArgs.request.enable);
+
+    // Move and verify scale is not applied.
+    process(mapper, ARBITRARY_TIME, READ_TIME, EV_REL, REL_X, 10);
+    process(mapper, ARBITRARY_TIME, READ_TIME, EV_REL, REL_Y, 20);
+    process(mapper, ARBITRARY_TIME, READ_TIME, EV_SYN, SYN_REPORT, 0);
+    ASSERT_NO_FATAL_FAILURE(mFakeListener->assertNotifyMotionWasCalled(&args));
+    ASSERT_EQ(AINPUT_SOURCE_MOUSE_RELATIVE, args.source);
+    ASSERT_EQ(AMOTION_EVENT_ACTION_MOVE, args.action);
+    ASSERT_EQ(10, args.pointerCoords[0].getX());
+    ASSERT_EQ(20, args.pointerCoords[0].getY());
+}
+
 TEST_F(CursorInputMapperTest, Process_ShouldHandleDisplayId) {
     CursorInputMapper& mapper = addMapperAndConfigure<CursorInputMapper>();
 
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index ee4c41e..3bfc2cc 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -2730,12 +2730,12 @@
         }
 
         const auto displayId = info->id;
-        const auto it = mPhysicalDisplayTokens.find(displayId);
+        const auto token = mPhysicalDisplayTokens.get(displayId);
 
         if (event.connection == hal::Connection::CONNECTED) {
             auto [supportedModes, activeMode] = loadDisplayModes(displayId);
 
-            if (it == mPhysicalDisplayTokens.end()) {
+            if (!token) {
                 ALOGV("Creating display %s", to_string(displayId).c_str());
 
                 DisplayDeviceState state;
@@ -2750,14 +2750,13 @@
 
                 sp<IBinder> token = new BBinder();
                 mCurrentState.displays.add(token, state);
-                mPhysicalDisplayTokens.emplace(displayId, std::move(token));
+                mPhysicalDisplayTokens.try_emplace(displayId, std::move(token));
                 mInterceptor->saveDisplayCreation(state);
             } else {
                 ALOGV("Recreating display %s", to_string(displayId).c_str());
 
-                const auto token = it->second;
-                auto& state = mCurrentState.displays.editValueFor(token);
-                state.sequenceId = DisplayDeviceState{}.sequenceId; // Generate new sequenceId
+                auto& state = mCurrentState.displays.editValueFor(token->get());
+                state.sequenceId = DisplayDeviceState{}.sequenceId; // Generate new sequenceId.
                 state.physical->supportedModes = std::move(supportedModes);
                 state.physical->activeMode = std::move(activeMode);
                 if (getHwComposer().updatesDeviceProductInfoOnHotplugReconnect()) {
@@ -2767,13 +2766,13 @@
         } else {
             ALOGV("Removing display %s", to_string(displayId).c_str());
 
-            const ssize_t index = mCurrentState.displays.indexOfKey(it->second);
-            if (index >= 0) {
+            if (const ssize_t index = mCurrentState.displays.indexOfKey(token->get()); index >= 0) {
                 const DisplayDeviceState& state = mCurrentState.displays.valueAt(index);
                 mInterceptor->saveDisplayDeletion(state.sequenceId);
                 mCurrentState.displays.removeItemsAt(index);
             }
-            mPhysicalDisplayTokens.erase(it);
+
+            mPhysicalDisplayTokens.erase(displayId);
         }
 
         processDisplayChangesLocked();
@@ -2954,15 +2953,16 @@
     }
 
     LOG_FATAL_IF(!displaySurface);
-    const auto display = setupNewDisplayDeviceInternal(displayToken, std::move(compositionDisplay),
-                                                       state, displaySurface, producer);
-    mDisplays.emplace(displayToken, display);
+    auto display = setupNewDisplayDeviceInternal(displayToken, std::move(compositionDisplay), state,
+                                                 displaySurface, producer);
     if (display->isPrimary()) {
         initScheduler(display);
     }
     if (!state.isVirtual()) {
         dispatchDisplayHotplugEvent(display->getPhysicalId(), true);
     }
+
+    mDisplays.try_emplace(displayToken, std::move(display));
 }
 
 void SurfaceFlinger::processDisplayRemoved(const wp<IBinder>& displayToken) {
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 07da731..011aaef 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -27,6 +27,7 @@
 #include <android/gui/DisplayState.h>
 #include <cutils/atomic.h>
 #include <cutils/compiler.h>
+#include <ftl/small_map.h>
 #include <gui/BufferQueue.h>
 #include <gui/FrameTimestamps.h>
 #include <gui/ISurfaceComposer.h>
@@ -905,8 +906,8 @@
     }
 
     sp<DisplayDevice> getDisplayDeviceLocked(const wp<IBinder>& displayToken) REQUIRES(mStateLock) {
-        const auto it = mDisplays.find(displayToken);
-        return it == mDisplays.end() ? nullptr : it->second;
+        const sp<DisplayDevice> nullDisplay;
+        return mDisplays.get(displayToken).value_or(std::cref(nullDisplay));
     }
 
     sp<const DisplayDevice> getDisplayDeviceLocked(PhysicalDisplayId id) const
@@ -1051,8 +1052,8 @@
      */
     sp<IBinder> getPhysicalDisplayTokenLocked(PhysicalDisplayId displayId) const
             REQUIRES(mStateLock) {
-        const auto it = mPhysicalDisplayTokens.find(displayId);
-        return it != mPhysicalDisplayTokens.end() ? it->second : nullptr;
+        const sp<IBinder> nullToken;
+        return mPhysicalDisplayTokens.get(displayId).value_or(std::cref(nullToken));
     }
 
     std::optional<PhysicalDisplayId> getPhysicalDisplayIdLocked(
@@ -1247,10 +1248,14 @@
 
     std::vector<HotplugEvent> mPendingHotplugEvents GUARDED_BY(mStateLock);
 
-    // this may only be written from the main thread with mStateLock held
-    // it may be read from other threads with mStateLock held
-    std::map<wp<IBinder>, sp<DisplayDevice>> mDisplays GUARDED_BY(mStateLock);
-    std::unordered_map<PhysicalDisplayId, sp<IBinder>> mPhysicalDisplayTokens
+    // Displays are composited in `mDisplays` order. Internal displays are inserted at boot and
+    // never removed, so take precedence over external and virtual displays.
+    //
+    // The static capacities were chosen to exceed a typical number of physical/virtual displays.
+    //
+    // May be read from any thread, but must only be written from the main thread.
+    ftl::SmallMap<wp<IBinder>, const sp<DisplayDevice>, 5> mDisplays GUARDED_BY(mStateLock);
+    ftl::SmallMap<PhysicalDisplayId, const sp<IBinder>, 3> mPhysicalDisplayTokens
             GUARDED_BY(mStateLock);
 
     struct {
@@ -1423,10 +1428,8 @@
     std::atomic<ui::Transform::RotationFlags> mActiveDisplayTransformHint;
 
     bool isRefreshRateOverlayEnabled() const REQUIRES(mStateLock) {
-        return std::any_of(mDisplays.begin(), mDisplays.end(),
-                           [](std::pair<wp<IBinder>, sp<DisplayDevice>> display) {
-                               return display.second->isRefreshRateOverlayEnabled();
-                           });
+        return hasDisplay(
+                [](const auto& display) { return display.isRefreshRateOverlayEnabled(); });
     }
 
     wp<IBinder> mActiveDisplayToken GUARDED_BY(mStateLock);
diff --git a/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp b/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
index 2425862..f04221c 100644
--- a/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
@@ -165,36 +165,39 @@
     return displayDevice;
 }
 
-bool DisplayTransactionTest::hasPhysicalHwcDisplay(HWDisplayId hwcDisplayId) {
-    return mFlinger.mutableHwcPhysicalDisplayIdMap().count(hwcDisplayId) == 1;
+bool DisplayTransactionTest::hasPhysicalHwcDisplay(HWDisplayId hwcDisplayId) const {
+    return mFlinger.hwcPhysicalDisplayIdMap().count(hwcDisplayId) == 1;
 }
 
-bool DisplayTransactionTest::hasTransactionFlagSet(int flag) {
-    return mFlinger.mutableTransactionFlags() & flag;
+bool DisplayTransactionTest::hasTransactionFlagSet(int32_t flag) const {
+    return mFlinger.transactionFlags() & flag;
 }
 
-bool DisplayTransactionTest::hasDisplayDevice(sp<IBinder> displayToken) {
-    return mFlinger.mutableDisplays().count(displayToken) == 1;
+bool DisplayTransactionTest::hasDisplayDevice(const sp<IBinder>& displayToken) const {
+    return mFlinger.displays().contains(displayToken);
 }
 
-sp<DisplayDevice> DisplayTransactionTest::getDisplayDevice(sp<IBinder> displayToken) {
-    return mFlinger.mutableDisplays()[displayToken];
+const DisplayDevice& DisplayTransactionTest::getDisplayDevice(
+        const sp<IBinder>& displayToken) const {
+    return *mFlinger.displays().get(displayToken)->get();
 }
 
-bool DisplayTransactionTest::hasCurrentDisplayState(sp<IBinder> displayToken) {
-    return mFlinger.mutableCurrentState().displays.indexOfKey(displayToken) >= 0;
+bool DisplayTransactionTest::hasCurrentDisplayState(const sp<IBinder>& displayToken) const {
+    return mFlinger.currentState().displays.indexOfKey(displayToken) >= 0;
 }
 
-const DisplayDeviceState& DisplayTransactionTest::getCurrentDisplayState(sp<IBinder> displayToken) {
-    return mFlinger.mutableCurrentState().displays.valueFor(displayToken);
+const DisplayDeviceState& DisplayTransactionTest::getCurrentDisplayState(
+        const sp<IBinder>& displayToken) const {
+    return mFlinger.currentState().displays.valueFor(displayToken);
 }
 
-bool DisplayTransactionTest::hasDrawingDisplayState(sp<IBinder> displayToken) {
-    return mFlinger.mutableDrawingState().displays.indexOfKey(displayToken) >= 0;
+bool DisplayTransactionTest::hasDrawingDisplayState(const sp<IBinder>& displayToken) const {
+    return mFlinger.drawingState().displays.indexOfKey(displayToken) >= 0;
 }
 
-const DisplayDeviceState& DisplayTransactionTest::getDrawingDisplayState(sp<IBinder> displayToken) {
-    return mFlinger.mutableDrawingState().displays.valueFor(displayToken);
+const DisplayDeviceState& DisplayTransactionTest::getDrawingDisplayState(
+        const sp<IBinder>& displayToken) const {
+    return mFlinger.drawingState().displays.valueFor(displayToken);
 }
 
 } // namespace android
diff --git a/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h b/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h
index 565c244..f5235ce 100644
--- a/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h
+++ b/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h
@@ -95,14 +95,17 @@
     // --------------------------------------------------------------------
     // Postcondition helpers
 
-    bool hasPhysicalHwcDisplay(hal::HWDisplayId hwcDisplayId);
-    bool hasTransactionFlagSet(int flag);
-    bool hasDisplayDevice(sp<IBinder> displayToken);
-    sp<DisplayDevice> getDisplayDevice(sp<IBinder> displayToken);
-    bool hasCurrentDisplayState(sp<IBinder> displayToken);
-    const DisplayDeviceState& getCurrentDisplayState(sp<IBinder> displayToken);
-    bool hasDrawingDisplayState(sp<IBinder> displayToken);
-    const DisplayDeviceState& getDrawingDisplayState(sp<IBinder> displayToken);
+    bool hasPhysicalHwcDisplay(hal::HWDisplayId) const;
+    bool hasTransactionFlagSet(int32_t flag) const;
+
+    bool hasDisplayDevice(const sp<IBinder>& displayToken) const;
+    const DisplayDevice& getDisplayDevice(const sp<IBinder>& displayToken) const;
+
+    bool hasCurrentDisplayState(const sp<IBinder>& displayToken) const;
+    const DisplayDeviceState& getCurrentDisplayState(const sp<IBinder>& displayToken) const;
+
+    bool hasDrawingDisplayState(const sp<IBinder>& displayToken) const;
+    const DisplayDeviceState& getDrawingDisplayState(const sp<IBinder>& displayToken) const;
 
     // --------------------------------------------------------------------
     // Test instances
diff --git a/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayModeSwitching.cpp b/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayModeSwitching.cpp
index 32d57b5..5872a47 100644
--- a/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayModeSwitching.cpp
+++ b/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayModeSwitching.cpp
@@ -61,7 +61,6 @@
 
 protected:
     void setupScheduler(std::shared_ptr<scheduler::RefreshRateConfigs>);
-    void testChangeRefreshRate(bool isDisplayActive, bool isRefreshRequired);
 
     sp<DisplayDevice> mDisplay;
     mock::EventThread* mAppEventThread;
diff --git a/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayTransactionCommitTest.cpp b/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayTransactionCommitTest.cpp
index 6959ee3..9ac2907 100644
--- a/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayTransactionCommitTest.cpp
+++ b/services/surfaceflinger/tests/unittests/SurfaceFlinger_DisplayTransactionCommitTest.cpp
@@ -107,9 +107,10 @@
 void DisplayTransactionCommitTest::verifyDisplayIsConnected(const sp<IBinder>& displayToken) {
     // The display device should have been set up in the list of displays.
     ASSERT_TRUE(hasDisplayDevice(displayToken));
-    const auto& device = getDisplayDevice(displayToken);
-    EXPECT_EQ(static_cast<bool>(Case::Display::SECURE), device->isSecure());
-    EXPECT_EQ(static_cast<bool>(Case::Display::PRIMARY), device->isPrimary());
+    const auto& display = getDisplayDevice(displayToken);
+
+    EXPECT_EQ(static_cast<bool>(Case::Display::SECURE), display.isSecure());
+    EXPECT_EQ(static_cast<bool>(Case::Display::PRIMARY), display.isPrimary());
 
     std::optional<DisplayDeviceState::Physical> expectedPhysical;
     if (const auto connectionType = Case::Display::CONNECTION_TYPE::value) {
@@ -143,10 +144,11 @@
     // SF should have a display token.
     const auto displayId = Case::Display::DISPLAY_ID::get();
     ASSERT_TRUE(PhysicalDisplayId::tryCast(displayId));
-    ASSERT_EQ(mFlinger.mutablePhysicalDisplayTokens().count(displayId), 1);
-    auto& displayToken = mFlinger.mutablePhysicalDisplayTokens()[displayId];
 
-    verifyDisplayIsConnected<Case>(displayToken);
+    const auto displayTokenOpt = mFlinger.mutablePhysicalDisplayTokens().get(displayId);
+    ASSERT_TRUE(displayTokenOpt);
+
+    verifyDisplayIsConnected<Case>(displayTokenOpt->get());
 }
 
 void DisplayTransactionCommitTest::verifyDisplayIsNotConnected(const sp<IBinder>& displayToken) {
@@ -248,9 +250,9 @@
     // SF should not have a display token.
     const auto displayId = Case::Display::DISPLAY_ID::get();
     ASSERT_TRUE(PhysicalDisplayId::tryCast(displayId));
-    ASSERT_TRUE(mFlinger.mutablePhysicalDisplayTokens().count(displayId) == 0);
+    ASSERT_FALSE(mFlinger.mutablePhysicalDisplayTokens().contains(displayId));
 
-    // The existing token should have been removed
+    // The existing token should have been removed.
     verifyDisplayIsNotConnected(existing.token());
 }
 
@@ -330,7 +332,7 @@
                 // SF should not have a display token.
                 const auto displayId = Case::Display::DISPLAY_ID::get();
                 ASSERT_TRUE(PhysicalDisplayId::tryCast(displayId));
-                ASSERT_TRUE(mFlinger.mutablePhysicalDisplayTokens().count(displayId) == 0);
+                ASSERT_FALSE(mFlinger.mutablePhysicalDisplayTokens().contains(displayId));
             }(),
             testing::KilledBySignal(SIGABRT), "Primary display cannot be disconnected.");
 }
@@ -369,15 +371,16 @@
                 // --------------------------------------------------------------------
                 // Postconditions
 
-                // The existing token should have been removed
+                // The existing token should have been removed.
                 verifyDisplayIsNotConnected(existing.token());
                 const auto displayId = Case::Display::DISPLAY_ID::get();
                 ASSERT_TRUE(PhysicalDisplayId::tryCast(displayId));
-                ASSERT_TRUE(mFlinger.mutablePhysicalDisplayTokens().count(displayId) == 1);
-                EXPECT_NE(existing.token(), mFlinger.mutablePhysicalDisplayTokens()[displayId]);
 
-                // A new display should be connected in its place
+                const auto displayTokenOpt = mFlinger.mutablePhysicalDisplayTokens().get(displayId);
+                ASSERT_TRUE(displayTokenOpt);
+                EXPECT_NE(existing.token(), displayTokenOpt->get());
 
+                // A new display should be connected in its place.
                 verifyPhysicalDisplayIsConnected<Case>();
 
                 // --------------------------------------------------------------------
diff --git a/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp b/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp
index 7948e60..583cf5f 100644
--- a/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp
+++ b/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp
@@ -406,7 +406,7 @@
     display.inject();
 
     // The display is set to PowerMode::ON
-    getDisplayDevice(display.token())->setPowerMode(PowerMode::ON);
+    display.mutableDisplayDevice()->setPowerMode(PowerMode::ON);
 
     // --------------------------------------------------------------------
     // Invocation
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 866d9eb..f1a69fb 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -517,6 +517,12 @@
      * post-conditions.
      */
 
+    const auto& displays() const { return mFlinger->mDisplays; }
+    const auto& currentState() const { return mFlinger->mCurrentState; }
+    const auto& drawingState() const { return mFlinger->mDrawingState; }
+    const auto& transactionFlags() const { return mFlinger->mTransactionFlags; }
+    const auto& hwcPhysicalDisplayIdMap() const { return getHwComposer().mPhysicalDisplayIdMap; }
+
     auto& mutableHasWideColorDisplay() { return SurfaceFlinger::hasWideColorDisplay; }
 
     auto& mutableCurrentState() { return mFlinger->mCurrentState; }
@@ -755,7 +761,9 @@
             return mFlinger.mutableCurrentState().displays.valueFor(mDisplayToken);
         }
 
-        auto& mutableDisplayDevice() { return mFlinger.mutableDisplays()[mDisplayToken]; }
+        const sp<DisplayDevice>& mutableDisplayDevice() {
+            return mFlinger.mutableDisplays().get(mDisplayToken)->get();
+        }
 
         // If `configs` is nullptr, the injector creates RefreshRateConfigs from the `modes`.
         // Otherwise, it uses `configs`, which the caller must create using the same `modes`.
@@ -862,12 +870,14 @@
             if (!display->isVirtual()) {
                 display->setActiveMode(activeModeId);
             }
-            mFlinger.mutableDisplays().emplace(mDisplayToken, display);
+            mFlinger.mutableDisplays().emplace_or_replace(mDisplayToken, display);
+
             mFlinger.mutableCurrentState().displays.add(mDisplayToken, state);
             mFlinger.mutableDrawingState().displays.add(mDisplayToken, state);
 
             if (const auto& physical = state.physical) {
-                mFlinger.mutablePhysicalDisplayTokens()[physical->id] = mDisplayToken;
+                mFlinger.mutablePhysicalDisplayTokens().emplace_or_replace(physical->id,
+                                                                           mDisplayToken);
             }
 
             return display;