diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index d9f7804..9777b4d 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -7846,14 +7846,13 @@
 
 namespace {
 
-ui::Dataspace pickBestDataspace(ui::Dataspace requestedDataspace, const DisplayDevice* display,
+ui::Dataspace pickBestDataspace(ui::Dataspace requestedDataspace,
+                                const compositionengine::impl::OutputCompositionState& state,
                                 bool capturingHdrLayers, bool hintForSeamlessTransition) {
-    if (requestedDataspace != ui::Dataspace::UNKNOWN || display == nullptr) {
+    if (requestedDataspace != ui::Dataspace::UNKNOWN) {
         return requestedDataspace;
     }
 
-    const auto& state = display->getCompositionDisplay()->getState();
-
     const auto dataspaceForColorMode = ui::pickDataspaceFor(state.colorMode);
 
     // TODO: Enable once HDR screenshots are ready.
@@ -8358,8 +8357,7 @@
 
     if (display != nullptr) {
         captureResults.capturedDataspace =
-                pickBestDataspace(requestedDataspace, display.get(),
-                                  captureResults.capturedHdrLayers,
+                pickBestDataspace(requestedDataspace, state, captureResults.capturedHdrLayers,
                                   renderArea->getHintForSeamlessTransition());
         sdrWhitePointNits = state.sdrWhitePointNits;