[Paint Preview] Support alt compositing mode for Long Screenshots

This addresses two shortcomings of PlayerCompositorDelegate for the
Long Screenshots use case:
1. An in-memory PaintPreviewProto is provided from Java rather than one
   from disk.
2. Compositing should occur to a single SkPicture using the MainFrame
   mode of the compositor rather than the SeparateFrame mode.

This CL also makes the MainFrame mode return the dimensions of the
SkPicture so that Long Screenshot can leverage this information.

Change-Id: If0b2d666f7459b94f4501c23808c504826a0112a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2578076Reviewed-by: Ken Buchanan <kenrb@chromium.org>
Reviewed-by: Mehran Mahmoudi <mahmoudi@chromium.org>
Commit-Queue: Calder Kitagawa <ckitagawa@chromium.org>
Cr-Commit-Position: refs/heads/master@{#835290}

chrome/browser/paint_preview/android/BUILD.gn

@@ -37,6 +37,7 @@ android_library("java") {
     "//chrome/browser/ui/messages/android:java",
     "//components/browser_ui/styles/android:java",
     "//components/paint_preview/browser/android:java",
+    "//components/paint_preview/common/proto:proto_java",
     "//components/paint_preview/player/android:java",
     "//content/public/android:content_java",
     "//third_party/android_deps:androidx_annotation_annotation_java",
@@ -72,6 +73,7 @@ android_library("javatests") {
     "//chrome/browser/ui/messages/android:java",
     "//chrome/test/android:chrome_java_test_support",
     "//components/paint_preview/browser/android:java",
+    "//components/paint_preview/common/proto:proto_java",
     "//components/paint_preview/player/android:java",
     "//content/public/android:content_java",
     "//content/public/test/android:content_java_test_support",

chrome/browser/paint_preview/android/javatests/src/org/chromium/chrome/browser/paint_preview/DemoPaintPreviewTest.java

@@ -60,7 +60,7 @@ public class DemoPaintPreviewTest {
         sMockService = Mockito.mock(PaintPreviewTabService.class);
         TabbedPaintPreview.overridePaintPreviewTabServiceForTesting(sMockService);
         PlayerManager.overrideCompositorDelegateFactoryForTesting(
-                TabbedPaintPreviewTest.TestCompositorDelegate::new);
+                new TabbedPaintPreviewTest.TestCompositorDelegateFactory());
     }
 
     @AfterClass

chrome/browser/paint_preview/android/javatests/src/org/chromium/chrome/browser/paint_preview/StartupPaintPreviewTest.java

@@ -59,7 +59,7 @@ public class StartupPaintPreviewTest {
         Mockito.doReturn(true).when(mockService).hasCaptureForTab(Mockito.anyInt());
         TabbedPaintPreview.overridePaintPreviewTabServiceForTesting(mockService);
         PlayerManager.overrideCompositorDelegateFactoryForTesting(
-                TabbedPaintPreviewTest.TestCompositorDelegate::new);
+                new TabbedPaintPreviewTest.TestCompositorDelegateFactory());
     }
 
     @AfterClass

chrome/browser/paint_preview/android/javatests/src/org/chromium/chrome/browser/paint_preview/TabbedPaintPreviewTest.java

@@ -10,6 +10,7 @@ import android.os.Handler;
 import android.os.Parcel;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.test.filters.MediumTest;
 
 import org.junit.After;
@@ -32,6 +33,7 @@ import org.chromium.chrome.browser.paint_preview.services.PaintPreviewTabService
 import org.chromium.chrome.browser.tab.Tab;
 import org.chromium.chrome.test.ChromeJUnit4ClassRunner;
 import org.chromium.chrome.test.ChromeTabbedActivityTestRule;
+import org.chromium.components.paint_preview.common.proto.PaintPreview.PaintPreviewProto;
 import org.chromium.components.paintpreview.browser.NativePaintPreviewServiceProvider;
 import org.chromium.components.paintpreview.player.PlayerCompositorDelegate;
 import org.chromium.components.paintpreview.player.PlayerManager;
@@ -53,12 +55,37 @@ public class TabbedPaintPreviewTest {
 
     private static final String TEST_URL = "/chrome/test/data/android/about.html";
 
+    /**
+     * Implementation of {@link PlayerCompositorDelegate.Factory} for tests.
+     */
+    public static class TestCompositorDelegateFactory implements PlayerCompositorDelegate.Factory {
+        @Override
+        public PlayerCompositorDelegate create(NativePaintPreviewServiceProvider service, GURL url,
+                String directoryKey, boolean mainFrameMode,
+                @NonNull PlayerCompositorDelegate.CompositorListener compositorListener,
+                Callback<Integer> compositorErrorCallback) {
+            return new TestCompositorDelegate(service, null, url, directoryKey, mainFrameMode,
+                    compositorListener, compositorErrorCallback);
+        }
+
+        @Override
+        public PlayerCompositorDelegate createForProto(NativePaintPreviewServiceProvider service,
+                @Nullable PaintPreviewProto proto, GURL url, String directoryKey,
+                boolean mainFrameMode,
+                @NonNull PlayerCompositorDelegate.CompositorListener compositorListener,
+                Callback<Integer> compositorErrorCallback) {
+            Assert.fail("createForProto shouldn't be called");
+            return null;
+        }
+    }
+
     @Before
     public void setUp() {
         PaintPreviewTabService mockService = Mockito.mock(PaintPreviewTabService.class);
         Mockito.doReturn(true).when(mockService).hasCaptureForTab(Mockito.anyInt());
         TabbedPaintPreview.overridePaintPreviewTabServiceForTesting(mockService);
-        PlayerManager.overrideCompositorDelegateFactoryForTesting(TestCompositorDelegate::new);
+        PlayerManager.overrideCompositorDelegateFactoryForTesting(
+                new TestCompositorDelegateFactory());
         mActivityTestRule.startMainActivityWithURL(
                 mActivityTestRule.getTestServer().getURL(TEST_URL));
     }
@@ -296,9 +323,12 @@ public class TabbedPaintPreviewTest {
     public static class TestCompositorDelegate implements PlayerCompositorDelegate {
         private int mNextRequestId;
 
-        TestCompositorDelegate(NativePaintPreviewServiceProvider service, GURL url,
-                String directoryKey, @NonNull CompositorListener compositorListener,
+        TestCompositorDelegate(NativePaintPreviewServiceProvider service,
+                @Nullable PaintPreviewProto proto, GURL url, String directoryKey,
+                boolean mainFrameMode, @NonNull CompositorListener compositorListener,
                 Callback<Integer> compositorErrorCallback) {
+            Assert.assertNull(proto);
+            Assert.assertFalse(mainFrameMode);
             new Handler().postDelayed(() -> {
                 Parcel parcel = Parcel.obtain();
                 parcel.writeLong(4577L);
@@ -326,6 +356,14 @@ public class TabbedPaintPreviewTest {
             return requestId;
         }
 
+        @Override
+        public int requestBitmap(Rect clipRect, float scaleFactor, Callback<Bitmap> bitmapCallback,
+                Runnable errorCallback) {
+            Assert.fail("The GUIDless version of TestCompositorDelegate#requestBitmap() shouldn't"
+                    + " be called.");
+            return 0;
+        }
+
         @Override
         public boolean cancelBitmapRequest(int requestId) {
             return false;

components/paint_preview/player/android/BUILD.gn

@@ -85,6 +85,7 @@ android_library("java") {
     "//base:base_java",
     "//base:jni_java",
     "//components/paint_preview/browser/android:java",
+    "//components/paint_preview/common/proto:proto_java",
     "//third_party/android_deps:androidx_annotation_annotation_java",
     "//third_party/android_swipe_refresh:android_swipe_refresh_java",
     "//ui/android:ui_java",
@@ -129,6 +130,7 @@ android_library("javatests") {
     ":player_java_test_support",
     "//base:base_java",
     "//base:base_java_test_support",
+    "//components/paint_preview/common/proto:proto_java",
     "//content/public/android:content_java",
     "//content/public/test/android:content_java_test_support",
     "//third_party/android_deps:androidx_test_runner_java",

components/paint_preview/player/android/java/src/org/chromium/components/paintpreview/player/PlayerCompositorDelegate.java

@@ -8,9 +8,11 @@ import android.graphics.Bitmap;
 import android.graphics.Rect;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import org.chromium.base.Callback;
 import org.chromium.base.UnguessableToken;
+import org.chromium.components.paint_preview.common.proto.PaintPreview.PaintPreviewProto;
 import org.chromium.components.paintpreview.browser.NativePaintPreviewServiceProvider;
 import org.chromium.url.GURL;
 
@@ -22,7 +24,13 @@ public interface PlayerCompositorDelegate {
     /** An interface that creates an instance of {@link PlayerCompositorDelegate}. */
     interface Factory {
         PlayerCompositorDelegate create(NativePaintPreviewServiceProvider service, GURL url,
-                String directoryKey, @NonNull CompositorListener compositorListener,
+                String directoryKey, boolean mainFrameMode,
+                @NonNull CompositorListener compositorListener,
+                Callback<Integer> compositorErrorCallback);
+
+        PlayerCompositorDelegate createForProto(NativePaintPreviewServiceProvider service,
+                @Nullable PaintPreviewProto proto, GURL url, String directoryKey,
+                boolean mainFrameMode, @NonNull CompositorListener compositorListener,
                 Callback<Integer> compositorErrorCallback);
     }
 
@@ -72,7 +80,7 @@ public interface PlayerCompositorDelegate {
      * Requests a new bitmap for a frame from the Paint Preview compositor.
      * @param frameGuid The GUID of the frame.
      * @param clipRect The {@link Rect} for which the bitmap is requested.
-     * @param scaleFactor The scale factor at which the bitmap should be rendered.
+     * @param scaleFactor The scale factor at which the bitmap should be rastered.
      * @param bitmapCallback The callback that receives the bitmap once it's ready. Won't get called
      * if there are any errors.
      * @param errorCallback Gets notified if there are any errors. Won't get called otherwise.
@@ -82,6 +90,20 @@ public interface PlayerCompositorDelegate {
     int requestBitmap(UnguessableToken frameGuid, Rect clipRect, float scaleFactor,
             Callback<Bitmap> bitmapCallback, Runnable errorCallback);
 
+    /**
+     * Requests a new bitmap for a frame from the Paint Preview compositor if {@link mainFrameMode}
+     * was passed as true as a parameter in the {@link Factory}
+     * @param clipRect The {@link Rect} for which the bitmap is requested.
+     * @param scaleFactor The scale factor at which the bitmap should be rastered.
+     * @param bitmapCallback The callback that receives the bitmap once it's ready. Won't get called
+     * if there are any errors.
+     * @param errorCallback Gets notified if there are any errors. Won't get called otherwise.
+     * @return an int representing the ID for the bitmap request. Can be used with {@link
+     * cancelBitmapRequest(int)} to cancel the request if possible.
+     */
+    int requestBitmap(Rect clipRect, float scaleFactor, Callback<Bitmap> bitmapCallback,
+            Runnable errorCallback);
+
     /**
      * Cancels the bitmap request for the provided ID.
      * @param requestID the request to try to cancel.

components/paint_preview/player/android/java/src/org/chromium/components/paintpreview/player/PlayerCompositorDelegateImpl.java

@@ -9,6 +9,7 @@ import android.graphics.Rect;
 import android.text.TextUtils;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import org.chromium.base.Callback;
 import org.chromium.base.SysUtils;
@@ -16,6 +17,7 @@ import org.chromium.base.UnguessableToken;
 import org.chromium.base.annotations.CalledByNative;
 import org.chromium.base.annotations.JNINamespace;
 import org.chromium.base.annotations.NativeMethods;
+import org.chromium.components.paint_preview.common.proto.PaintPreview.PaintPreviewProto;
 import org.chromium.components.paintpreview.browser.NativePaintPreviewServiceProvider;
 import org.chromium.url.GURL;
 
@@ -34,14 +36,15 @@ class PlayerCompositorDelegateImpl implements PlayerCompositorDelegate {
     private long mNativePlayerCompositorDelegate;
     private List<Runnable> mMemoryPressureListeners = new ArrayList<>();
 
-    PlayerCompositorDelegateImpl(NativePaintPreviewServiceProvider service, GURL url,
-            String directoryKey, @NonNull CompositorListener compositorListener,
+    PlayerCompositorDelegateImpl(NativePaintPreviewServiceProvider service,
+            @Nullable PaintPreviewProto proto, GURL url, String directoryKey, boolean mainFrameMode,
+            @NonNull CompositorListener compositorListener,
             Callback<Integer> compositorErrorCallback) {
         mCompositorListener = compositorListener;
         if (service != null && service.getNativeBaseService() != 0) {
             mNativePlayerCompositorDelegate = PlayerCompositorDelegateImplJni.get().initialize(this,
-                    service.getNativeBaseService(), url.getSpec(), directoryKey,
-                    compositorErrorCallback,
+                    service.getNativeBaseService(), (proto != null) ? proto.toByteArray() : null,
+                    url.getSpec(), directoryKey, mainFrameMode, compositorErrorCallback,
                     SysUtils.amountOfPhysicalMemoryKB() < LOW_MEMORY_THRESHOLD_KB);
         }
         // TODO(crbug.com/1021590): Handle initialization errors when
@@ -80,6 +83,12 @@ class PlayerCompositorDelegateImpl implements PlayerCompositorDelegate {
                 clipRect.width(), clipRect.height());
     }
 
+    @Override
+    public int requestBitmap(Rect clipRect, float scaleFactor, Callback<Bitmap> bitmapCallback,
+            Runnable errorCallback) {
+        return requestBitmap(null, clipRect, scaleFactor, bitmapCallback, errorCallback);
+    }
+
     @Override
     public boolean cancelBitmapRequest(int requestId) {
         if (mNativePlayerCompositorDelegate == 0) {
@@ -136,8 +145,8 @@ class PlayerCompositorDelegateImpl implements PlayerCompositorDelegate {
     @NativeMethods
     interface Natives {
         long initialize(PlayerCompositorDelegateImpl caller, long nativePaintPreviewBaseService,
-                String urlSpec, String directoryKey, Callback<Integer> compositorErrorCallback,
-                boolean isLowMemory);
+                byte[] proto, String urlSpec, String directoryKey, boolean mainFrameMode,
+                Callback<Integer> compositorErrorCallback, boolean isLowMemory);
         void destroy(long nativePlayerCompositorDelegateAndroid);
         int requestBitmap(long nativePlayerCompositorDelegateAndroid, UnguessableToken frameGuid,
                 Callback<Bitmap> bitmapCallback, Runnable errorCallback, float scaleFactor,

components/paint_preview/player/android/java/src/org/chromium/components/paintpreview/player/PlayerManager.java

@@ -13,10 +13,13 @@ import android.view.ViewGroup.LayoutParams;
 import android.widget.FrameLayout;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
+import org.chromium.base.Callback;
 import org.chromium.base.TraceEvent;
 import org.chromium.base.UnguessableToken;
+import org.chromium.components.paint_preview.common.proto.PaintPreview.PaintPreviewProto;
 import org.chromium.components.paintpreview.browser.NativePaintPreviewServiceProvider;
 import org.chromium.components.paintpreview.player.frame.PlayerFrameCoordinator;
 import org.chromium.url.GURL;
@@ -71,7 +74,8 @@ public class PlayerManager {
         void onLinkClick(GURL url);
     }
 
-    private static PlayerCompositorDelegate.Factory sCompositorDelegateFactoryForTesting;
+    private static PlayerCompositorDelegate.Factory sCompositorDelegateFactory =
+            new CompositorDelegateFactory();
 
     private Context mContext;
     private PlayerCompositorDelegate mDelegate;
@@ -103,7 +107,7 @@ public class PlayerManager {
         mContext = context;
         mListener = listener;
         mDelegate = getCompositorDelegateFactory().create(nativePaintPreviewServiceProvider, url,
-                directoryKey, this::onCompositorReady, mListener::onCompositorError);
+                directoryKey, false, this::onCompositorReady, mListener::onCompositorError);
         mHostView = new FrameLayout(mContext);
         mPlayerSwipeRefreshHandler =
                 new PlayerSwipeRefreshHandler(mContext, mListener::onPullToRefresh);
@@ -241,9 +245,29 @@ public class PlayerManager {
         return mHostView;
     }
 
+    static class CompositorDelegateFactory implements PlayerCompositorDelegate.Factory {
+        @Override
+        public PlayerCompositorDelegate create(NativePaintPreviewServiceProvider service, GURL url,
+                String directoryKey, boolean mainFrameMode,
+                @NonNull PlayerCompositorDelegate.CompositorListener compositorListener,
+                Callback<Integer> compositorErrorCallback) {
+            return new PlayerCompositorDelegateImpl(service, null, url, directoryKey, mainFrameMode,
+                    compositorListener, compositorErrorCallback);
+        }
+
+        @Override
+        public PlayerCompositorDelegate createForProto(NativePaintPreviewServiceProvider service,
+                @Nullable PaintPreviewProto proto, GURL url, String directoryKey,
+                boolean mainFrameMode,
+                @NonNull PlayerCompositorDelegate.CompositorListener compositorListener,
+                Callback<Integer> compositorErrorCallback) {
+            return new PlayerCompositorDelegateImpl(service, proto, url, directoryKey,
+                    mainFrameMode, compositorListener, compositorErrorCallback);
+        }
+    }
+
     private PlayerCompositorDelegate.Factory getCompositorDelegateFactory() {
-        return (sCompositorDelegateFactoryForTesting != null) ? sCompositorDelegateFactoryForTesting
-                                                              : PlayerCompositorDelegateImpl::new;
+        return sCompositorDelegateFactory;
     }
 
     @VisibleForTesting
@@ -254,6 +278,10 @@ public class PlayerManager {
     @VisibleForTesting
     public static void overrideCompositorDelegateFactoryForTesting(
             PlayerCompositorDelegate.Factory factory) {
-        sCompositorDelegateFactoryForTesting = factory;
+        if (factory == null) {
+            sCompositorDelegateFactory = new CompositorDelegateFactory();
+            return;
+        }
+        sCompositorDelegateFactory = factory;
     }
 }

components/paint_preview/player/android/junit/src/org/chromium/components/paintpreview/player/frame/PlayerFrameMediatorTest.java

@@ -193,6 +193,14 @@ public class PlayerFrameMediatorTest {
             return requestId;
         }
 
+        @Override
+        public int requestBitmap(Rect clipRect, float scaleFactor, Callback<Bitmap> bitmapCallback,
+                Runnable errorCallback) {
+            Assert.fail("The GUIDless version of TestPlayerCompositorDelegate#requestBitmap() "
+                    + "shouldn't be called.");
+            return 0;
+        }
+
         @Override
         public boolean cancelBitmapRequest(int requestId) {
             return false;

components/paint_preview/player/android/player_compositor_delegate_android.cc

@@ -64,16 +64,18 @@ jlong JNI_PlayerCompositorDelegateImpl_Initialize(
     JNIEnv* env,
     const JavaParamRef<jobject>& j_object,
     jlong paint_preview_service,
+    const JavaParamRef<jbyteArray>& j_proto,
     const JavaParamRef<jstring>& j_url_spec,
     const JavaParamRef<jstring>& j_directory_key,
+    jboolean j_main_frame_mode,
     const JavaParamRef<jobject>& j_compositor_error_callback,
     jboolean j_is_low_mem) {
   PlayerCompositorDelegateAndroid* delegate =
       new PlayerCompositorDelegateAndroid(
           env, j_object,
           reinterpret_cast<PaintPreviewBaseService*>(paint_preview_service),
-          j_url_spec, j_directory_key, j_compositor_error_callback,
-          j_is_low_mem);
+          j_proto, j_url_spec, j_directory_key, j_main_frame_mode,
+          j_compositor_error_callback, j_is_low_mem);
   return reinterpret_cast<intptr_t>(delegate);
 }
 
@@ -81,18 +83,33 @@ PlayerCompositorDelegateAndroid::PlayerCompositorDelegateAndroid(
     JNIEnv* env,
     const JavaParamRef<jobject>& j_object,
     PaintPreviewBaseService* paint_preview_service,
+    const JavaParamRef<jbyteArray>& j_proto,
     const JavaParamRef<jstring>& j_url_spec,
     const JavaParamRef<jstring>& j_directory_key,
+    jboolean j_main_frame_mode,
     const JavaParamRef<jobject>& j_compositor_error_callback,
     jboolean j_is_low_mem)
     : PlayerCompositorDelegate(),
       request_id_(0),
       startup_timestamp_(base::TimeTicks::Now()) {
+  if (j_proto) {
+    std::string serialized_proto;
+    base::android::JavaByteArrayToString(env, j_proto, &serialized_proto);
+    auto proto = std::make_unique<PaintPreviewProto>();
+    if (!proto->ParseFromString(serialized_proto)) {
+      base::android::RunIntCallbackAndroid(
+          j_compositor_error_callback,
+          static_cast<int>(CompositorStatus::PROTOBUF_DESERIALIZATION_ERROR));
+      return;
+    }
+    PlayerCompositorDelegate::SetProto(std::move(proto));
+  }
   PlayerCompositorDelegate::Initialize(
       paint_preview_service,
       GURL(base::android::ConvertJavaStringToUTF8(env, j_url_spec)),
       DirectoryKey{
           base::android::ConvertJavaStringToUTF8(env, j_directory_key)},
+      static_cast<bool>(j_main_frame_mode),
       base::BindOnce(&base::android::RunIntCallbackAndroid,
                      ScopedJavaGlobalRef<jobject>(j_compositor_error_callback)),
       base::TimeDelta::FromSeconds(15),
@@ -224,20 +241,23 @@ jint PlayerCompositorDelegateAndroid::RequestBitmap(
   TRACE_EVENT_NESTABLE_ASYNC_BEGIN0(
       "paint_preview", "PlayerCompositorDelegateAndroid::RequestBitmap",
       TRACE_ID_LOCAL(request_id_));
-
-  int32_t id = PlayerCompositorDelegate::RequestBitmap(
-      base::android::UnguessableTokenAndroid::FromJavaUnguessableToken(
-          env, j_frame_guid),
-      gfx::Rect(j_clip_x, j_clip_y, j_clip_width, j_clip_height),
-      j_scale_factor,
-      base::BindOnce(&PlayerCompositorDelegateAndroid::OnBitmapCallback,
-                     weak_factory_.GetWeakPtr(),
-                     ScopedJavaGlobalRef<jobject>(j_bitmap_callback),
-                     ScopedJavaGlobalRef<jobject>(j_error_callback),
-                     request_id_));
+  gfx::Rect rect(j_clip_x, j_clip_y, j_clip_width, j_clip_height);
+  auto callback = base::BindOnce(
+      &PlayerCompositorDelegateAndroid::OnBitmapCallback,
+      weak_factory_.GetWeakPtr(),
+      ScopedJavaGlobalRef<jobject>(j_bitmap_callback),
+      ScopedJavaGlobalRef<jobject>(j_error_callback), request_id_);
   ++request_id_;
 
-  return static_cast<jint>(id);
+  base::Optional<base::UnguessableToken> frame_guid;
+  if (j_frame_guid) {
+    frame_guid =
+        base::android::UnguessableTokenAndroid::FromJavaUnguessableToken(
+            env, j_frame_guid);
+  }
+
+  return static_cast<jint>(PlayerCompositorDelegate::RequestBitmap(
+      frame_guid, rect, j_scale_factor, std::move(callback)));
 }
 
 jboolean PlayerCompositorDelegateAndroid::CancelBitmapRequest(

components/paint_preview/player/android/player_compositor_delegate_android.h

@@ -21,8 +21,10 @@ class PlayerCompositorDelegateAndroid : public PlayerCompositorDelegate {
       JNIEnv* env,
       const base::android::JavaParamRef<jobject>& j_object,
       PaintPreviewBaseService* paint_preview_service,
+      const base::android::JavaParamRef<jbyteArray>& j_proto,
       const base::android::JavaParamRef<jstring>& j_url_spec,
       const base::android::JavaParamRef<jstring>& j_directory_key,
+      jboolean j_main_frame_mode,
       const base::android::JavaParamRef<jobject>& j_compositor_error_callback,
       jboolean j_is_low_mem);
 

components/paint_preview/player/bitmap_request.cc

@@ -6,10 +6,11 @@
 
 namespace paint_preview {
 
-BitmapRequest::BitmapRequest(const base::UnguessableToken& frame_guid,
-                             const gfx::Rect& clip_rect,
-                             float scale_factor,
-                             BitmapRequestCallback callback)
+BitmapRequest::BitmapRequest(
+    const base::Optional<base::UnguessableToken>& frame_guid,
+    const gfx::Rect& clip_rect,
+    float scale_factor,
+    BitmapRequestCallback callback)
     : frame_guid(frame_guid),
       clip_rect(clip_rect),
       scale_factor(scale_factor),

components/paint_preview/player/bitmap_request.h

@@ -6,6 +6,7 @@
 #define COMPONENTS_PAINT_PREVIEW_PLAYER_BITMAP_REQUEST_H_
 
 #include "base/callback.h"
+#include "base/optional.h"
 #include "base/unguessable_token.h"
 #include "components/services/paint_preview_compositor/public/mojom/paint_preview_compositor.mojom.h"
 #include "third_party/skia/include/core/SkBitmap.h"
@@ -18,7 +19,7 @@ struct BitmapRequest {
       base::OnceCallback<void(mojom::PaintPreviewCompositor::BitmapStatus,
                               const SkBitmap&)>;
 
-  BitmapRequest(const base::UnguessableToken& frame_guid,
+  BitmapRequest(const base::Optional<base::UnguessableToken>& frame_guid,
                 const gfx::Rect& clip_rect,
                 float scale_factor,
                 BitmapRequestCallback callback);
@@ -27,7 +28,7 @@ struct BitmapRequest {
   BitmapRequest& operator=(BitmapRequest&& other) noexcept;
   BitmapRequest(BitmapRequest&& other) noexcept;
 
-  base::UnguessableToken frame_guid;
+  base::Optional<base::UnguessableToken> frame_guid;
   gfx::Rect clip_rect;
   float scale_factor;
   BitmapRequestCallback callback;

components/paint_preview/player/player_compositor_delegate.cc

@@ -117,6 +117,7 @@ void PlayerCompositorDelegate::Initialize(
     PaintPreviewBaseService* paint_preview_service,
     const GURL& expected_url,
     const DirectoryKey& key,
+    bool main_frame_mode,
     base::OnceCallback<void(int)> compositor_error,
     base::TimeDelta timeout_duration,
     size_t max_requests) {
@@ -141,7 +142,7 @@ void PlayerCompositorDelegate::Initialize(
           &PlayerCompositorDelegate::OnCompositorServiceDisconnected,
           weak_factory_.GetWeakPtr()));
 
-  InitializeInternal(paint_preview_service, expected_url, key,
+  InitializeInternal(paint_preview_service, expected_url, key, main_frame_mode,
                      std::move(compositor_error), timeout_duration,
                      max_requests);
 }
@@ -150,6 +151,7 @@ void PlayerCompositorDelegate::InitializeWithFakeServiceForTest(
     PaintPreviewBaseService* paint_preview_service,
     const GURL& expected_url,
     const DirectoryKey& key,
+    bool main_frame_mode,
     base::OnceCallback<void(int)> compositor_error,
     base::TimeDelta timeout_duration,
     size_t max_requests,
@@ -160,7 +162,7 @@ void PlayerCompositorDelegate::InitializeWithFakeServiceForTest(
       base::BindOnce(&PlayerCompositorDelegate::OnCompositorServiceDisconnected,
                      weak_factory_.GetWeakPtr()));
 
-  InitializeInternal(paint_preview_service, expected_url, key,
+  InitializeInternal(paint_preview_service, expected_url, key, main_frame_mode,
                      std::move(compositor_error), timeout_duration,
                      max_requests);
 }
@@ -169,10 +171,12 @@ void PlayerCompositorDelegate::InitializeInternal(
     PaintPreviewBaseService* paint_preview_service,
     const GURL& expected_url,
     const DirectoryKey& key,
+    bool main_frame_mode,
     base::OnceCallback<void(int)> compositor_error,
     base::TimeDelta timeout_duration,
     size_t max_requests) {
   max_requests_ = max_requests;
+  main_frame_mode_ = main_frame_mode;
   compositor_error_ = std::move(compositor_error);
   paint_preview_service_ = paint_preview_service;
   key_ = key;
@@ -199,12 +203,14 @@ void PlayerCompositorDelegate::InitializeInternal(
 }
 
 int32_t PlayerCompositorDelegate::RequestBitmap(
-    const base::UnguessableToken& frame_guid,
+    const base::Optional<base::UnguessableToken>& frame_guid,
     const gfx::Rect& clip_rect,
     float scale_factor,
     base::OnceCallback<void(mojom::PaintPreviewCompositor::BitmapStatus,
                             const SkBitmap&)> callback) {
   DCHECK(IsInitialized());
+  DCHECK((main_frame_mode_ && !frame_guid.has_value()) ||
+         (!main_frame_mode_ && frame_guid.has_value()));
   const int32_t request_id = next_request_id_;
   next_request_id_++;
   if (!paint_preview_compositor_client_) {
@@ -314,10 +320,15 @@ void PlayerCompositorDelegate::OnCompositorClientCreated(
   TRACE_EVENT_NESTABLE_ASYNC_END0("paint_preview",
                                   "PlayerCompositorDelegate CreateCompositor",
                                   TRACE_ID_LOCAL(this));
-  paint_preview_service_->GetFileMixin()->GetCapturedPaintPreviewProto(
-      key, base::nullopt,
-      base::BindOnce(&PlayerCompositorDelegate::OnProtoAvailable,
-                     weak_factory_.GetWeakPtr(), expected_url));
+  if (!proto_) {
+    paint_preview_service_->GetFileMixin()->GetCapturedPaintPreviewProto(
+        key, base::nullopt,
+        base::BindOnce(&PlayerCompositorDelegate::OnProtoAvailable,
+                       weak_factory_.GetWeakPtr(), expected_url));
+  } else {
+    OnProtoAvailable(expected_url, PaintPreviewFileMixin::ProtoReadStatus::kOk,
+                     std::move(proto_));
+  }
 }
 
 void PlayerCompositorDelegate::OnProtoAvailable(
@@ -394,10 +405,20 @@ void PlayerCompositorDelegate::SendCompositeRequest(
     return;
   }
 
-  paint_preview_compositor_client_->BeginSeparatedFrameComposite(
-      std::move(begin_composite_request),
-      base::BindOnce(&PlayerCompositorDelegate::OnCompositorReadyStatusAdapter,
-                     weak_factory_.GetWeakPtr()));
+  if (main_frame_mode_) {
+    paint_preview_compositor_client_->BeginMainFrameComposite(
+        std::move(begin_composite_request),
+        base::BindOnce(
+            &PlayerCompositorDelegate::OnCompositorReadyStatusAdapter,
+            weak_factory_.GetWeakPtr()));
+
+  } else {
+    paint_preview_compositor_client_->BeginSeparatedFrameComposite(
+        std::move(begin_composite_request),
+        base::BindOnce(
+            &PlayerCompositorDelegate::OnCompositorReadyStatusAdapter,
+            weak_factory_.GetWeakPtr()));
+  }
 
   // Defer building hit testers so it happens in parallel with preparing the
   // compositor.
@@ -437,9 +458,14 @@ void PlayerCompositorDelegate::ProcessBitmapRequestsFromQueue() {
     if (!paint_preview_compositor_client_)
       return;
 
-    paint_preview_compositor_client_->BitmapForSeparatedFrame(
-        request.frame_guid, request.clip_rect, request.scale_factor,
-        std::move(request.callback));
+    if (request.frame_guid.has_value()) {
+      paint_preview_compositor_client_->BitmapForSeparatedFrame(
+          request.frame_guid.value(), request.clip_rect, request.scale_factor,
+          std::move(request.callback));
+    } else {
+      paint_preview_compositor_client_->BitmapForMainFrame(
+          request.clip_rect, request.scale_factor, std::move(request.callback));
+    }
     pending_bitmap_requests_.erase(it);
   }
 }

components/paint_preview/player/player_compositor_delegate.h

@@ -49,6 +49,7 @@ class PlayerCompositorDelegate {
   void Initialize(PaintPreviewBaseService* paint_preview_service,
                   const GURL& url,
                   const DirectoryKey& key,
+                  bool main_frame_mode,
                   base::OnceCallback<void(int)> compositor_error,
                   base::TimeDelta timeout_duration,
                   size_t max_requests);
@@ -56,6 +57,10 @@ class PlayerCompositorDelegate {
   // Returns whether initialization has happened.
   bool IsInitialized() const { return paint_preview_service_; }
 
+  void SetProto(std::unique_ptr<PaintPreviewProto> proto) {
+    proto_ = std::move(proto);
+  }
+
   // Overrides whether to compress the directory when the player is closed. By
   // default compression will happen.
   void SetCompressOnClose(bool compress) { compress_on_close_ = compress; }
@@ -71,7 +76,7 @@ class PlayerCompositorDelegate {
   // Pass this ID to `CancelBitmapRequest(int32_t)` to cancel the request if it
   // hasn't already been sent.
   int32_t RequestBitmap(
-      const base::UnguessableToken& frame_guid,
+      const base::Optional<base::UnguessableToken>& frame_guid,
       const gfx::Rect& clip_rect,
       float scale_factor,
       base::OnceCallback<void(mojom::PaintPreviewCompositor::BitmapStatus,
@@ -100,6 +105,7 @@ class PlayerCompositorDelegate {
       PaintPreviewBaseService* paint_preview_service,
       const GURL& expected_url,
       const DirectoryKey& key,
+      bool main_frame_mode,
       base::OnceCallback<void(int)> compositor_error,
       base::TimeDelta timeout_duration,
       size_t max_requests,
@@ -123,6 +129,7 @@ class PlayerCompositorDelegate {
   void InitializeInternal(PaintPreviewBaseService* paint_preview_service,
                           const GURL& expected_url,
                           const DirectoryKey& key,
+                          bool main_frame_mode,
                           base::OnceCallback<void(int)> compositor_error,
                           base::TimeDelta timeout_duration,
                           size_t max_requests);
@@ -166,6 +173,7 @@ class PlayerCompositorDelegate {
 
   base::CancelableOnceClosure timeout_;
   int max_requests_{1};
+  bool main_frame_mode_{false};
 
   base::flat_map<base::UnguessableToken, std::unique_ptr<HitTester>>
       hit_testers_;

components/paint_preview/player/player_compositor_delegate_unittest.cc

@@ -84,7 +84,11 @@ class FakePaintPreviewCompositorClient : public PaintPreviewCompositorClient {
       mojom::PaintPreviewBeginCompositeRequestPtr request,
       mojom::PaintPreviewCompositor::BeginMainFrameCompositeCallback callback)
       override {
-    NOTREACHED();
+    auto response = mojom::PaintPreviewBeginCompositeResponse::New();
+    response->root_frame_guid = base::UnguessableToken::Create();
+    task_runner_->PostTask(FROM_HERE,
+                           base::BindOnce(std::move(callback), response_status_,
+                                          std::move(response)));
   }
 
   void BitmapForMainFrame(
@@ -92,7 +96,14 @@ class FakePaintPreviewCompositorClient : public PaintPreviewCompositorClient {
       float scale_factor,
       mojom::PaintPreviewCompositor::BitmapForMainFrameCallback callback)
       override {
-    NOTREACHED();
+    SkBitmap bitmap;
+    bitmap.allocPixels(
+        SkImageInfo::MakeN32Premul(clip_rect.width(), clip_rect.height()));
+    task_runner_->PostTask(
+        FROM_HERE,
+        base::BindOnce(std::move(callback),
+                       mojom::PaintPreviewCompositor::BitmapStatus::kSuccess,
+                       bitmap));
   }
 
   void SetRootFrameUrl(const GURL& url) override {
@@ -251,14 +262,16 @@ class PlayerCompositorDelegateTest : public testing::Test {
   }
 
   void SerializeProtoAndCreateRootSkp(PaintPreviewProto* proto,
-                                      const DirectoryKey& key) {
+                                      const DirectoryKey& key,
+                                      bool skip_proto_serialization = false) {
     auto file_manager = GetBaseService()->GetFileMixin()->GetFileManager();
     base::RunLoop loop;
     file_manager->GetTaskRunner()->PostTask(
         FROM_HERE,
         base::BindOnce(
             [](base::OnceClosure quit, scoped_refptr<FileManager> file_manager,
-               PaintPreviewProto* proto, const DirectoryKey& key) {
+               PaintPreviewProto* proto, const DirectoryKey& key,
+               bool skip_proto_serialization) {
               auto directory = file_manager->CreateOrGetDirectory(key, true);
 
               std::string fake_data = "Hello World!";
@@ -267,10 +280,13 @@ class PlayerCompositorDelegateTest : public testing::Test {
                   root_file.AsUTF8Unsafe());
               base::WriteFile(root_file, fake_data.data(), fake_data.size());
 
-              file_manager->SerializePaintPreviewProto(key, *proto, false);
+              if (!skip_proto_serialization) {
+                file_manager->SerializePaintPreviewProto(key, *proto, false);
+              }
               std::move(quit).Run();
             },
-            loop.QuitClosure(), file_manager, proto, key));
+            loop.QuitClosure(), file_manager, proto, key,
+            skip_proto_serialization));
     loop.Run();
   }
 
@@ -348,8 +364,9 @@ TEST_F(PlayerCompositorDelegateTest, OnClick) {
     PlayerCompositorDelegateImpl player_compositor_delegate;
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::OK);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, url, key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, url, key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
 
@@ -392,8 +409,9 @@ TEST_F(PlayerCompositorDelegateTest, BadProto) {
     player_compositor_delegate.SetExpectedStatus(
         CompositorStatus::PROTOBUF_DESERIALIZATION_ERROR);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, GURL(), key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
   }
@@ -412,8 +430,32 @@ TEST_F(PlayerCompositorDelegateTest, OldVersion) {
     PlayerCompositorDelegateImpl player_compositor_delegate;
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::OLD_VERSION);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, url, key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, url, key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
+    player_compositor_delegate.SetCompressOnClose(false);
+    env.RunUntilIdle();
+    EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
+  }
+  env.RunUntilIdle();
+}
+
+TEST_F(PlayerCompositorDelegateTest, InMemoryProto) {
+  auto* service = GetBaseService();
+  auto file_manager = service->GetFileMixin()->GetFileManager();
+  auto key = file_manager->CreateKey(1U);
+  GURL url("https://www.chromium.org/");
+  auto proto = CreateValidProto(url);
+  SerializeProtoAndCreateRootSkp(&proto, key, true);
+  {
+    PlayerCompositorDelegateImpl player_compositor_delegate;
+    player_compositor_delegate.SetProto(
+        std::make_unique<PaintPreviewProto>(proto));
+    player_compositor_delegate.SetExpectedStatus(CompositorStatus::OK);
+    player_compositor_delegate.InitializeWithFakeServiceForTest(
+        service, url, key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     player_compositor_delegate.SetCompressOnClose(false);
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
@@ -433,8 +475,9 @@ TEST_F(PlayerCompositorDelegateTest, URLMismatch) {
     player_compositor_delegate.SetExpectedStatus(
         CompositorStatus::URL_MISMATCH);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, GURL(), key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
   }
@@ -453,7 +496,7 @@ TEST_F(PlayerCompositorDelegateTest, ServiceDisconnect) {
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::OK);
     bool called = false;
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, url, key,
+        service, url, key, /*main_frame_mode=*/false,
         base::BindOnce(
             [](bool* called, int status) {
               EXPECT_EQ(static_cast<int>(
@@ -485,7 +528,7 @@ TEST_F(PlayerCompositorDelegateTest, ClientDisconnect) {
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::OK);
     bool called = false;
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, url, key,
+        service, url, key, /*main_frame_mode=*/false,
         base::BindOnce(
             [](bool* called, int status) {
               EXPECT_EQ(static_cast<int>(
@@ -527,8 +570,9 @@ TEST_F(PlayerCompositorDelegateTest, InvalidCompositeRequest) {
     player_compositor_delegate.SetExpectedStatus(
         CompositorStatus::INVALID_REQUEST);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, url, key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, url, key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
   }
@@ -547,8 +591,9 @@ TEST_F(PlayerCompositorDelegateTest, CompositorDeserializationError) {
     player_compositor_delegate.SetExpectedStatus(
         CompositorStatus::COMPOSITOR_DESERIALIZATION_ERROR);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, url, key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, url, key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     AsFakeClient(player_compositor_delegate.GetClientForTest())
         ->SetBeginSeparatedFrameResponseStatus(
             mojom::PaintPreviewCompositor::BeginCompositeStatus::
@@ -571,8 +616,9 @@ TEST_F(PlayerCompositorDelegateTest, InvalidRootSkp) {
     player_compositor_delegate.SetExpectedStatus(
         CompositorStatus::INVALID_ROOT_FRAME_SKP);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, url, key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, url, key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     AsFakeClient(player_compositor_delegate.GetClientForTest())
         ->SetBeginSeparatedFrameResponseStatus(
             mojom::PaintPreviewCompositor::BeginCompositeStatus::
@@ -606,8 +652,9 @@ TEST_F(PlayerCompositorDelegateTest, CompressOnClose) {
     PlayerCompositorDelegateImpl player_compositor_delegate;
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::NO_CAPTURE);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, GURL(), key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
   }
@@ -626,8 +673,9 @@ TEST_F(PlayerCompositorDelegateTest, RequestBitmapWithCancel) {
     PlayerCompositorDelegateImpl player_compositor_delegate;
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::NO_CAPTURE);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, GURL(), key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
 
@@ -697,8 +745,9 @@ TEST_F(PlayerCompositorDelegateTest, RequestBitmapWithCancelAll) {
     PlayerCompositorDelegateImpl player_compositor_delegate;
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::NO_CAPTURE);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, GURL(), key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
 
@@ -749,8 +798,9 @@ TEST_F(PlayerCompositorDelegateTest, RequestBitmapSuccessQueued) {
     PlayerCompositorDelegateImpl player_compositor_delegate;
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::NO_CAPTURE);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key, base::DoNothing(), base::TimeDelta::Max(),
-        kMaxParallelRequests, CreateCompositorService());
+        service, GURL(), key, /*main_frame_mode=*/false, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
     env.RunUntilIdle();
     EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
 
@@ -771,6 +821,43 @@ TEST_F(PlayerCompositorDelegateTest, RequestBitmapSuccessQueued) {
   env.RunUntilIdle();
 }
 
+TEST_F(PlayerCompositorDelegateTest, RequestMainFrameBitmapSuccess) {
+  auto* service = GetBaseService();
+  auto file_manager = service->GetFileMixin()->GetFileManager();
+  auto key = file_manager->CreateKey(1U);
+  GURL url("https://www.chromium.org/");
+  auto proto = CreateValidProto(url);
+  SerializeProtoAndCreateRootSkp(&proto, key);
+  {
+    // This test skips setting up files as the fakes don't use them. In normal
+    // execution the files are required by the service or no bitmap will be
+    // created.
+    PlayerCompositorDelegateImpl player_compositor_delegate;
+    player_compositor_delegate.SetExpectedStatus(CompositorStatus::OK);
+    player_compositor_delegate.InitializeWithFakeServiceForTest(
+        service, url, key, /*main_frame_mode=*/true, base::DoNothing(),
+        base::TimeDelta::Max(), kMaxParallelRequests,
+        CreateCompositorService());
+    env.RunUntilIdle();
+    EXPECT_TRUE(player_compositor_delegate.WasStatusChecked());
+
+    base::RunLoop loop;
+    player_compositor_delegate.RequestBitmap(
+        base::nullopt, gfx::Rect(10, 20, 30, 40), 1.0,
+        base::BindOnce(
+            [](base::OnceClosure quit,
+               mojom::PaintPreviewCompositor::BitmapStatus status,
+               const SkBitmap& bitmap) {
+              EXPECT_EQ(mojom::PaintPreviewCompositor::BitmapStatus::kSuccess,
+                        status);
+              std::move(quit).Run();
+            },
+            loop.QuitClosure()));
+    loop.Run();
+  }
+  env.RunUntilIdle();
+}
+
 TEST_F(PlayerCompositorDelegateTest, Timeout) {
   auto* service = GetBaseService();
   auto file_manager = service->GetFileMixin()->GetFileManager();
@@ -781,7 +868,7 @@ TEST_F(PlayerCompositorDelegateTest, Timeout) {
     AsFakeService(compositor_service.get())->SetTimeout();
     base::RunLoop loop;
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key,
+        service, GURL(), key, /*main_frame_mode=*/false,
         base::BindOnce(
             [](base::OnceClosure quit, int status) {
               EXPECT_EQ(static_cast<CompositorStatus>(status),
@@ -809,7 +896,7 @@ TEST_F(PlayerCompositorDelegateTest, CriticalMemoryPressure) {
     PlayerCompositorDelegateImpl player_compositor_delegate;
     player_compositor_delegate.SetExpectedStatus(CompositorStatus::NO_CAPTURE);
     player_compositor_delegate.InitializeWithFakeServiceForTest(
-        service, GURL(), key,
+        service, GURL(), key, /*main_frame_mode=*/false,
         base::BindOnce(
             [](base::OnceClosure quit, int compositor_status) {
               EXPECT_EQ(compositor_status,
@@ -846,7 +933,7 @@ TEST_F(PlayerCompositorDelegateTest, CriticalMemoryPressureBeforeStart) {
     player_compositor_delegate.SetFakeMemoryPressureMonitor(
         &memory_pressure_monitor);
     player_compositor_delegate.Initialize(
-        service, GURL(), key,
+        service, GURL(), key, /*main_frame_mode=*/false,
         base::BindOnce(
             [](base::OnceClosure quit, int compositor_status) {
               EXPECT_EQ(compositor_status,

components/services/paint_preview_compositor/paint_preview_compositor_impl.cc

@@ -137,6 +137,7 @@ void PaintPreviewCompositorImpl::BeginSeparatedFrameComposite(
 
   // Remove any previously loaded frames, in case |BeginSeparatedFrameComposite|
   // is called multiple times.
+  root_frame_ = nullptr;
   frames_.clear();
 
   auto response = mojom::PaintPreviewBeginCompositeResponse::New();
@@ -231,11 +232,15 @@ void PaintPreviewCompositorImpl::BeginMainFrameComposite(
     BeginMainFrameCompositeCallback callback) {
   TRACE_EVENT0("paint_preview",
                "PaintPreviewCompositorImpl::BeginMainFrameComposite");
+  frames_.clear();
+  auto response = mojom::PaintPreviewBeginCompositeResponse::New();
   base::Optional<PaintPreviewProto> paint_preview =
       ParsePaintPreviewProto(request->proto);
   if (!paint_preview.has_value()) {
+    response->root_frame_guid = base::UnguessableToken::Create();
     std::move(callback).Run(mojom::PaintPreviewCompositor::
-                                BeginCompositeStatus::kDeserializingFailure);
+                                BeginCompositeStatus::kDeserializingFailure,
+                            std::move(response));
     return;
   }
 
@@ -245,8 +250,10 @@ void PaintPreviewCompositorImpl::BeginMainFrameComposite(
       paint_preview->root_frame().embedding_token_low());
   if (root_frame_guid.is_empty()) {
     DVLOG(1) << "No valid root frame guid";
+    response->root_frame_guid = base::UnguessableToken::Create();
     std::move(callback).Run(mojom::PaintPreviewCompositor::
-                                BeginCompositeStatus::kDeserializingFailure);
+                                BeginCompositeStatus::kDeserializingFailure,
+                            std::move(response));
     return;
   }
 
@@ -258,15 +265,31 @@ void PaintPreviewCompositorImpl::BeginMainFrameComposite(
                                           &recording_map, &subframes_failed);
 
   if (root_frame_ == nullptr) {
+    response->root_frame_guid = base::UnguessableToken::Create();
     std::move(callback).Run(mojom::PaintPreviewCompositor::
-                                BeginCompositeStatus::kCompositingFailure);
+                                BeginCompositeStatus::kCompositingFailure,
+                            std::move(response));
     return;
   }
 
+  response->root_frame_guid = root_frame_guid;
+  auto frame_data = mojom::FrameData::New();
+  SkRect sk_rect = root_frame_->cullRect();
+  frame_data->scroll_extents = gfx::Size(sk_rect.width(), sk_rect.height());
+  frame_data->scroll_offsets =
+      gfx::Size(paint_preview->root_frame().has_scroll_offset_x()
+                    ? paint_preview->root_frame().scroll_offset_x()
+                    : 0,
+                paint_preview->root_frame().has_scroll_offset_y()
+                    ? paint_preview->root_frame().scroll_offset_y()
+                    : 0);
+  response->frames.insert({root_frame_guid, std::move(frame_data)});
+
   std::move(callback).Run(
       subframes_failed
           ? mojom::PaintPreviewCompositor::BeginCompositeStatus::kPartialSuccess
-          : mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess);
+          : mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess,
+      std::move(response));
 }
 
 void PaintPreviewCompositorImpl::BitmapForMainFrame(

components/services/paint_preview_compositor/paint_preview_compositor_impl_unittest.cc

@@ -45,7 +45,7 @@ namespace {
 // then also checks that;
 // - |response->root_frame_guid| == |expected_root_frame_guid|
 // - |response->subframe_rect_hierarchy| == |expected_data|
-void BeginSeparatedFrameCompositeCallbackImpl(
+void BeginCompositeCallbackImpl(
     mojom::PaintPreviewCompositor::BeginCompositeStatus expected_status,
     const base::UnguessableToken& expected_root_frame_guid,
     const base::flat_map<base::UnguessableToken, mojom::FrameDataPtr>&
@@ -78,12 +78,6 @@ void BeginSeparatedFrameCompositeCallbackImpl(
   }
 }
 
-void BeginMainFrameCompositeCallbackImpl(
-    mojom::PaintPreviewCompositor::BeginCompositeStatus expected_status,
-    mojom::PaintPreviewCompositor::BeginCompositeStatus status) {
-  EXPECT_EQ(status, expected_status);
-}
-
 // Checks that |status| == |expected_status|. If |expected_status| == kSuccess,
 // then it also checks that |bitmap| and |expected_bitmap| are pixel identical.
 void BitmapCallbackImpl(
@@ -280,16 +274,22 @@ class PaintPreviewCompositorBeginCompositeTest
       case CompositeType::kSeparateFrame:
         compositor_.BeginSeparatedFrameComposite(
             std::move(request),
-            base::BindOnce(&BeginSeparatedFrameCompositeCallbackImpl,
-                           expected_status, expected_root_frame_guid,
-                           std::move(expected_data)));
+            base::BindOnce(&BeginCompositeCallbackImpl, expected_status,
+                           expected_root_frame_guid, std::move(expected_data)));
         break;
-      case CompositeType::kMainFrame:
+      case CompositeType::kMainFrame: {
+        base::flat_map<base::UnguessableToken, mojom::FrameDataPtr> root_data;
+        auto it = expected_data.find(expected_root_frame_guid);
+        if (it != expected_data.end()) {
+          root_data.insert({expected_root_frame_guid, it->second.Clone()});
+          root_data.find(expected_root_frame_guid)->second->subframes.clear();
+        }
         compositor_.BeginMainFrameComposite(
             std::move(request),
-            base::BindOnce(&BeginMainFrameCompositeCallbackImpl,
-                           expected_status));
+            base::BindOnce(&BeginCompositeCallbackImpl, expected_status,
+                           expected_root_frame_guid, std::move(root_data)));
         break;
+      }
       default:
         NOTREACHED();
         break;
@@ -592,7 +592,7 @@ TEST(PaintPreviewCompositorTest, TestComposite) {
   compositor.BeginSeparatedFrameComposite(
       std::move(request),
       base::BindOnce(
-          &BeginSeparatedFrameCompositeCallbackImpl,
+          &BeginCompositeCallbackImpl,
           mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess,
           kRootFrameID, std::move(expected_data)));
   float scale_factor = 2;
@@ -664,7 +664,7 @@ TEST(PaintPreviewCompositorTest, TestCompositeWithMemoryBuffer) {
   compositor.BeginSeparatedFrameComposite(
       std::move(request),
       base::BindOnce(
-          &BeginSeparatedFrameCompositeCallbackImpl,
+          &BeginCompositeCallbackImpl,
           mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess,
           kRootFrameID, std::move(expected_data)));
   float scale_factor = 2;
@@ -712,11 +712,9 @@ TEST(PaintPreviewCompositorTest, TestCompositeMainFrameNoDependencies) {
   compositor.BeginMainFrameComposite(
       std::move(request),
       base::BindOnce(
-          [](mojom::PaintPreviewCompositor::BeginCompositeStatus status) {
-            EXPECT_EQ(
-                status,
-                mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess);
-          }));
+          &BeginCompositeCallbackImpl,
+          mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess,
+          kRootFrameID, std::move(expected_data)));
   float scale_factor = 2;
   gfx::Rect rect = gfx::ScaleToEnclosingRect(
       gfx::Rect(root_frame_scroll_extent), scale_factor);
@@ -764,14 +762,14 @@ TEST(PaintPreviewCompositorTest, TestCompositeMainFrameOneDependency) {
       mojom::PaintPreviewBeginCompositeRequest::New();
   request->recording_map = RecordingMapFromPaintPreviewProto(proto);
   request->proto = ToReadOnlySharedMemory(proto);
+  expected_data.erase(kSubframe_0_ID);
+  expected_data.find(kRootFrameID)->second->subframes.clear();
   compositor.BeginMainFrameComposite(
       std::move(request),
       base::BindOnce(
-          [](mojom::PaintPreviewCompositor::BeginCompositeStatus status) {
-            EXPECT_EQ(
-                status,
-                mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess);
-          }));
+          &BeginCompositeCallbackImpl,
+          mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess,
+          kRootFrameID, std::move(expected_data)));
   float scale_factor = 1;
   gfx::Rect rect = gfx::ScaleToEnclosingRect(
       gfx::Rect(root_frame_scroll_extent), scale_factor);
@@ -823,14 +821,14 @@ TEST(PaintPreviewCompositorTest, TestCompositeMainFrameOneDependencyScrolled) {
       mojom::PaintPreviewBeginCompositeRequest::New();
   request->recording_map = RecordingMapFromPaintPreviewProto(proto);
   request->proto = ToReadOnlySharedMemory(proto);
+  expected_data.erase(kSubframe_0_ID);
+  expected_data.find(kRootFrameID)->second->subframes.clear();
   compositor.BeginMainFrameComposite(
       std::move(request),
       base::BindOnce(
-          [](mojom::PaintPreviewCompositor::BeginCompositeStatus status) {
-            EXPECT_EQ(
-                status,
-                mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess);
-          }));
+          &BeginCompositeCallbackImpl,
+          mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess,
+          kRootFrameID, std::move(expected_data)));
   float scale_factor = 1;
   gfx::Rect rect = gfx::ScaleToEnclosingRect(
       gfx::Rect(root_frame_scroll_extent), scale_factor);
@@ -885,14 +883,14 @@ TEST(PaintPreviewCompositorTest,
       mojom::PaintPreviewBeginCompositeRequest::New();
   request->recording_map = RecordingMapFromPaintPreviewProto(proto);
   request->proto = ToReadOnlySharedMemory(proto);
+  expected_data.erase(kSubframe_0_ID);
+  expected_data.find(kRootFrameID)->second->subframes.clear();
   compositor.BeginMainFrameComposite(
       std::move(request),
       base::BindOnce(
-          [](mojom::PaintPreviewCompositor::BeginCompositeStatus status) {
-            EXPECT_EQ(
-                status,
-                mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess);
-          }));
+          &BeginCompositeCallbackImpl,
+          mojom::PaintPreviewCompositor::BeginCompositeStatus::kSuccess,
+          kRootFrameID, std::move(expected_data)));
   float scale_factor = 1;
   gfx::Rect rect =
       gfx::ScaleToEnclosingRect(root_frame_clip_rect, scale_factor);

components/services/paint_preview_compositor/public/mojom/paint_preview_compositor.mojom

@@ -114,11 +114,10 @@ interface PaintPreviewCompositor {
     (BitmapStatus status, skia.mojom.BitmapN32? bitmap);
 
   // Starts the compositing process for |BitmapForMainFrame| calls using frame
-  // data |request|. |status| will be negative on failure. If the root frame is
-  // successfully loaded but subframes fail, |status| will indicate partial
-  // success.
+  // data |request|. On success returns |response| containing metadata
+  // required for UI. |status| will be negative on failure.
   BeginMainFrameComposite(PaintPreviewBeginCompositeRequest request) =>
-    (BeginCompositeStatus status);
+    (BeginCompositeStatus status, PaintPreviewBeginCompositeResponse? response);
 
   // Requests a bitmap of the main frame with sub-frame content included. The
   // dimensions and location of the bitmap will match |clip_rect| at scale