Fix Shell::Screenshot for Impeller (#50072)

Fixes flutter/flutter#141571

Shell::Screenshot was impelemnted in a Skia-only way and would crash when invoked. Adds test coverage for the breaking path on iOS that ends up calling `FlutterView drawLayer`.

Author: dnfield

shell/common/rasterizer.cc

@@ -8,6 +8,7 @@
 #include <memory>
 #include <utility>
 
+#include "display_list/dl_builder.h"
 #include "flow/frame_timings.h"
 #include "flutter/common/constants.h"
 #include "flutter/common/graphics/persistent_cache.h"
@@ -17,6 +18,7 @@
 #include "flutter/shell/common/base64.h"
 #include "flutter/shell/common/serialization_callbacks.h"
 #include "fml/make_copyable.h"
+#include "fml/synchronization/waitable_event.h"
 #include "third_party/skia/include/core/SkColorSpace.h"
 #include "third_party/skia/include/core/SkData.h"
 #include "third_party/skia/include/core/SkImage.h"
@@ -34,6 +36,12 @@
 #include "third_party/skia/include/gpu/GrTypes.h"
 #include "third_party/skia/include/gpu/ganesh/SkSurfaceGanesh.h"
 
+#if IMPELLER_SUPPORTS_RENDERING
+#include "impeller/aiks/aiks_context.h"           // nogncheck
+#include "impeller/core/formats.h"                // nogncheck
+#include "impeller/display_list/dl_dispatcher.h"  // nogncheck
+#endif
+
 namespace flutter {
 
 // The rasterizer will tell Skia to purge cached resources that have not been
@@ -802,28 +810,143 @@ static sk_sp<SkData> ScreenshotLayerTreeAsPicture(
   return recorder.finishRecordingAsPicture()->serialize(&procs);
 }
 
-sk_sp<SkData> Rasterizer::ScreenshotLayerTreeAsImage(
-    flutter::LayerTree* tree,
+static void RenderFrameForScreenshot(
     flutter::CompositorContext& compositor_context,
+    DlCanvas* canvas,
+    flutter::LayerTree* tree,
     GrDirectContext* surface_context,
+    const std::shared_ptr<impeller::AiksContext>& aiks_context) {
+  // There is no root surface transformation for the screenshot layer. Reset
+  // the matrix to identity.
+  SkMatrix root_surface_transformation;
+  root_surface_transformation.reset();
+
+  auto frame = compositor_context.AcquireFrame(
+      surface_context,              // skia context
+      canvas,                       // canvas
+      nullptr,                      // view embedder
+      root_surface_transformation,  // root surface transformation
+      false,                        // instrumentation enabled
+      true,                         // render buffer readback supported
+      nullptr,                      // thread merger
+      aiks_context.get()            // aiks context
+  );
+  canvas->Clear(DlColor::kTransparent());
+  frame->Raster(*tree, true, nullptr);
+  canvas->Flush();
+}
+
+#if IMPELLER_SUPPORTS_RENDERING
+Rasterizer::ScreenshotFormat ToScreenshotFormat(impeller::PixelFormat format) {
+  switch (format) {
+    case impeller::PixelFormat::kUnknown:
+    case impeller::PixelFormat::kA8UNormInt:
+    case impeller::PixelFormat::kR8UNormInt:
+    case impeller::PixelFormat::kR8G8UNormInt:
+    case impeller::PixelFormat::kR8G8B8A8UNormIntSRGB:
+    case impeller::PixelFormat::kB8G8R8A8UNormIntSRGB:
+    case impeller::PixelFormat::kB10G10R10XRSRGB:
+    case impeller::PixelFormat::kS8UInt:
+    case impeller::PixelFormat::kD24UnormS8Uint:
+    case impeller::PixelFormat::kD32FloatS8UInt:
+    case impeller::PixelFormat::kR32G32B32A32Float:
+    case impeller::PixelFormat::kB10G10R10XR:
+    case impeller::PixelFormat::kB10G10R10A10XR:
+      FML_DCHECK(false);
+      return Rasterizer::ScreenshotFormat::kUnknown;
+    case impeller::PixelFormat::kR8G8B8A8UNormInt:
+      return Rasterizer::ScreenshotFormat::kR8G8B8A8UNormInt;
+    case impeller::PixelFormat::kB8G8R8A8UNormInt:
+      return Rasterizer::ScreenshotFormat::kB8G8R8A8UNormInt;
+    case impeller::PixelFormat::kR16G16B16A16Float:
+      return Rasterizer::ScreenshotFormat::kR16G16B16A16Float;
+  }
+}
+
+static std::pair<sk_sp<SkData>, Rasterizer::ScreenshotFormat>
+ScreenshotLayerTreeAsImageImpeller(
+    const std::shared_ptr<impeller::AiksContext>& aiks_context,
+    flutter::LayerTree* tree,
+    flutter::CompositorContext& compositor_context,
+    bool compressed) {
+  if (compressed) {
+    FML_LOG(ERROR) << "Compressed screenshots not supported for Impeller";
+    return {nullptr, Rasterizer::ScreenshotFormat::kUnknown};
+  }
+
+  DisplayListBuilder builder(SkRect::MakeSize(
+      SkSize::Make(tree->frame_size().fWidth, tree->frame_size().fHeight)));
+
+  RenderFrameForScreenshot(compositor_context, &builder, tree, nullptr,
+                           aiks_context);
+
+  impeller::DlDispatcher dispatcher;
+  builder.Build()->Dispatch(dispatcher);
+  const auto& picture = dispatcher.EndRecordingAsPicture();
+  const auto& image = picture.ToImage(
+      *aiks_context,
+      impeller::ISize(tree->frame_size().fWidth, tree->frame_size().fHeight));
+  const auto& texture = image->GetTexture();
+  impeller::DeviceBufferDescriptor buffer_desc;
+  buffer_desc.storage_mode = impeller::StorageMode::kHostVisible;
+  buffer_desc.size =
+      texture->GetTextureDescriptor().GetByteSizeOfBaseMipLevel();
+  auto impeller_context = aiks_context->GetContext();
+  auto buffer =
+      impeller_context->GetResourceAllocator()->CreateBuffer(buffer_desc);
+  auto command_buffer = impeller_context->CreateCommandBuffer();
+  command_buffer->SetLabel("BlitTextureToBuffer Command Buffer");
+  auto pass = command_buffer->CreateBlitPass();
+  pass->AddCopy(texture, buffer);
+  pass->EncodeCommands(impeller_context->GetResourceAllocator());
+  fml::AutoResetWaitableEvent latch;
+  sk_sp<SkData> sk_data;
+  auto completion = [buffer, &buffer_desc, &sk_data,
+                     &latch](impeller::CommandBuffer::Status status) {
+    if (status != impeller::CommandBuffer::Status::kCompleted) {
+      FML_LOG(ERROR) << "Failed to complete blit pass.";
+      latch.Signal();
+      return;
+    }
+    sk_data = SkData::MakeWithCopy(buffer->OnGetContents(), buffer_desc.size);
+
+    latch.Signal();
+  };
+
+  if (!command_buffer->SubmitCommands(completion)) {
+    FML_LOG(ERROR) << "Failed to submit commands.";
+  }
+  latch.Wait();
+  return std::make_pair(
+      sk_data, ToScreenshotFormat(texture->GetTextureDescriptor().format));
+}
+#endif
+
+std::pair<sk_sp<SkData>, Rasterizer::ScreenshotFormat>
+Rasterizer::ScreenshotLayerTreeAsImage(
+    flutter::LayerTree* tree,
+    flutter::CompositorContext& compositor_context,
     bool compressed) {
+#if IMPELLER_SUPPORTS_RENDERING
+  if (delegate_.GetSettings().enable_impeller) {
+    return ScreenshotLayerTreeAsImageImpeller(GetAiksContext(), tree,
+                                              compositor_context, compressed);
+  }
+#endif  // IMPELLER_SUPPORTS_RENDERING
+
+  GrDirectContext* surface_context = GetGrContext();
   // Attempt to create a snapshot surface depending on whether we have access
   // to a valid GPU rendering context.
   std::unique_ptr<OffscreenSurface> snapshot_surface =
       std::make_unique<OffscreenSurface>(surface_context, tree->frame_size());
 
   if (!snapshot_surface->IsValid()) {
     FML_LOG(ERROR) << "Screenshot: unable to create snapshot surface";
-    return nullptr;
+    return {nullptr, ScreenshotFormat::kUnknown};
   }
 
   // Draw the current layer tree into the snapshot surface.
-  auto* canvas = snapshot_surface->GetCanvas();
-
-  // There is no root surface transformation for the screenshot layer. Reset
-  // the matrix to identity.
-  SkMatrix root_surface_transformation;
-  root_surface_transformation.reset();
+  DlCanvas* canvas = snapshot_surface->GetCanvas();
 
   // snapshot_surface->makeImageSnapshot needs the GL context to be set if the
   // render context is GL. frame->Raster() pops the gl context in platforms
@@ -832,29 +955,26 @@ sk_sp<SkData> Rasterizer::ScreenshotLayerTreeAsImage(
   auto context_switch = surface_->MakeRenderContextCurrent();
   if (!context_switch->GetResult()) {
     FML_LOG(ERROR) << "Screenshot: unable to make image screenshot";
-    return nullptr;
+    return {nullptr, ScreenshotFormat::kUnknown};
   }
 
-  auto frame = compositor_context.AcquireFrame(
-      surface_context,              // skia context
-      canvas,                       // canvas
-      nullptr,                      // view embedder
-      root_surface_transformation,  // root surface transformation
-      false,                        // instrumentation enabled
-      true,                         // render buffer readback supported
-      nullptr,                      // thread merger
-      nullptr                       // aiks context
-  );
-  canvas->Clear(DlColor::kTransparent());
-  frame->Raster(*tree, true, nullptr);
-  canvas->Flush();
+  RenderFrameForScreenshot(compositor_context, canvas, tree, surface_context,
+                           nullptr);
 
-  return snapshot_surface->GetRasterData(compressed);
+  return std::make_pair(snapshot_surface->GetRasterData(compressed),
+                        ScreenshotFormat::kUnknown);
 }
 
 Rasterizer::Screenshot Rasterizer::ScreenshotLastLayerTree(
     Rasterizer::ScreenshotType type,
     bool base64_encode) {
+  if (delegate_.GetSettings().enable_impeller &&
+      type == ScreenshotType::SkiaPicture) {
+    FML_DCHECK(false);
+    FML_LOG(ERROR) << "Last layer tree cannot be screenshotted as a "
+                      "SkiaPicture when using Impeller.";
+    return {};
+  }
   // TODO(dkwingsmt): Support screenshotting all last layer trees
   // when the shell protocol supports multi-views.
   // https://github.com/flutter/flutter/issues/135534
@@ -865,48 +985,49 @@ Rasterizer::Screenshot Rasterizer::ScreenshotLastLayerTree(
     return {};
   }
 
-  sk_sp<SkData> data = nullptr;
+  std::pair<sk_sp<SkData>, ScreenshotFormat> data{nullptr,
+                                                  ScreenshotFormat::kUnknown};
   std::string format;
 
-  GrDirectContext* surface_context =
-      surface_ ? surface_->GetContext() : nullptr;
-
   switch (type) {
     case ScreenshotType::SkiaPicture:
       format = "ScreenshotType::SkiaPicture";
-      data = ScreenshotLayerTreeAsPicture(layer_tree, *compositor_context_);
+      data.first =
+          ScreenshotLayerTreeAsPicture(layer_tree, *compositor_context_);
       break;
     case ScreenshotType::UncompressedImage:
       format = "ScreenshotType::UncompressedImage";
-      data = ScreenshotLayerTreeAsImage(layer_tree, *compositor_context_,
-                                        surface_context, false);
+      data =
+          ScreenshotLayerTreeAsImage(layer_tree, *compositor_context_, false);
       break;
     case ScreenshotType::CompressedImage:
       format = "ScreenshotType::CompressedImage";
-      data = ScreenshotLayerTreeAsImage(layer_tree, *compositor_context_,
-                                        surface_context, true);
+      data = ScreenshotLayerTreeAsImage(layer_tree, *compositor_context_, true);
       break;
     case ScreenshotType::SurfaceData: {
       Surface::SurfaceData surface_data = surface_->GetSurfaceData();
       format = surface_data.pixel_format;
-      data = surface_data.data;
+      data.first = surface_data.data;
       break;
     }
   }
 
-  if (data == nullptr) {
+  if (data.first == nullptr) {
     FML_LOG(ERROR) << "Screenshot data was null.";
     return {};
   }
 
   if (base64_encode) {
-    size_t b64_size = Base64::EncodedSize(data->size());
+    size_t b64_size = Base64::EncodedSize(data.first->size());
     auto b64_data = SkData::MakeUninitialized(b64_size);
-    Base64::Encode(data->data(), data->size(), b64_data->writable_data());
-    return Rasterizer::Screenshot{b64_data, layer_tree->frame_size(), format};
+    Base64::Encode(data.first->data(), data.first->size(),
+                   b64_data->writable_data());
+    return Rasterizer::Screenshot{b64_data, layer_tree->frame_size(), format,
+                                  data.second};
   }
 
-  return Rasterizer::Screenshot{data, layer_tree->frame_size(), format};
+  return Rasterizer::Screenshot{data.first, layer_tree->frame_size(), format,
+                                data.second};
 }
 
 void Rasterizer::SetNextFrameCallback(const fml::closure& callback) {
@@ -977,8 +1098,12 @@ Rasterizer::Screenshot::Screenshot() {}
 
 Rasterizer::Screenshot::Screenshot(sk_sp<SkData> p_data,
                                    SkISize p_size,
-                                   const std::string& p_format)
-    : data(std::move(p_data)), frame_size(p_size), format(p_format) {}
+                                   const std::string& p_format,
+                                   ScreenshotFormat p_pixel_format)
+    : data(std::move(p_data)),
+      frame_size(p_size),
+      format(p_format),
+      pixel_format(p_pixel_format) {}
 
 Rasterizer::Screenshot::Screenshot(const Screenshot& other) = default;
 

shell/common/rasterizer.h

@@ -25,8 +25,8 @@
 #include "flutter/fml/time/time_delta.h"
 #include "flutter/fml/time/time_point.h"
 #if IMPELLER_SUPPORTS_RENDERING
-// GN is having trouble understanding how this works in the Fuchsia builds.
 #include "impeller/aiks/aiks_context.h"  // nogncheck
+#include "impeller/core/formats.h"       // nogncheck
 #include "impeller/renderer/context.h"   // nogncheck
 #include "impeller/typographer/backends/skia/typographer_context_skia.h"  // nogncheck
 #endif  // IMPELLER_SUPPORTS_RENDERING
@@ -357,9 +357,10 @@ class Rasterizer final : public SnapshotDelegate,
     SkiaPicture,
 
     //--------------------------------------------------------------------------
-    /// A format used to denote uncompressed image data. This format
+    /// A format used to denote uncompressed image data. For Skia, this format
     /// is 32 bits per pixel, 8 bits per component and
-    /// denoted by the `kN32_SkColorType ` Skia color type.
+    /// denoted by the `kN32_SkColorType ` Skia color type. For Impeller, its
+    /// format is specified in Screenshot::pixel_format.
     ///
     UncompressedImage,
 
@@ -377,6 +378,18 @@ class Rasterizer final : public SnapshotDelegate,
     // NOLINTEND(readability-identifier-naming)
   };
 
+  // Specifies the format of pixel data in a Screenshot.
+  enum class ScreenshotFormat {
+    // Unknown format, or Skia default.
+    kUnknown,
+    // RGBA 8 bits per channel.
+    kR8G8B8A8UNormInt,
+    // BGRA 8 bits per channel.
+    kB8G8R8A8UNormInt,
+    // RGBA 16 bit floating point per channel.
+    kR16G16B16A16Float,
+  };
+
   //----------------------------------------------------------------------------
   /// @brief      A POD type used to return the screenshot data along with the
   ///             size of the frame.
@@ -400,6 +413,13 @@ class Rasterizer final : public SnapshotDelegate,
     ///
     std::string format;
 
+    //--------------------------------------------------------------------------
+    /// The pixel format of the data in `data`.
+    ///
+    /// If the impeller backend is not used, this value is always kUnknown and
+    /// the data is in RGBA8888 format.
+    ScreenshotFormat pixel_format = ScreenshotFormat::kUnknown;
+
     //--------------------------------------------------------------------------
     /// @brief      Creates an empty screenshot
     ///
@@ -411,10 +431,12 @@ class Rasterizer final : public SnapshotDelegate,
     /// @param[in]  p_data  The screenshot data
     /// @param[in]  p_size  The screenshot size.
     /// @param[in]  p_format  The screenshot format.
+    /// @param[in]  p_pixel_format  The screenshot format.
     ///
     Screenshot(sk_sp<SkData> p_data,
                SkISize p_size,
-               const std::string& p_format);
+               const std::string& p_format,
+               ScreenshotFormat p_pixel_format);
 
     //--------------------------------------------------------------------------
     /// @brief      The copy constructor for a screenshot.
@@ -663,10 +685,9 @@ class Rasterizer final : public SnapshotDelegate,
     return delegate_.GetIsGpuDisabledSyncSwitch();
   }
 
-  sk_sp<SkData> ScreenshotLayerTreeAsImage(
+  std::pair<sk_sp<SkData>, ScreenshotFormat> ScreenshotLayerTreeAsImage(
       flutter::LayerTree* tree,
       flutter::CompositorContext& compositor_context,
-      GrDirectContext* surface_context,
       bool compressed);
 
   // This method starts with the frame timing recorder at build end. This