[ios] Specify size of snapshot

Previously, WKWebView snapshots were defaulting to the size of the
web view, rather than the visible frame size. This caused the snapshot
to look out of place when animating.

Cq-Include-Trybots: luci.chromium.try:ios-simulator-cronet;luci.chromium.try:ios-simulator-full-configs
Change-Id: I533f570c996320ab9ae1d33cffe05a63303f7f36
Reviewed-on: https://chromium-review.googlesource.com/1247881
Commit-Queue: edchin <edchin@chromium.org>
Reviewed-by: Justin Cohen <justincohen@chromium.org>
Reviewed-by: Eugene But <eugenebut@chromium.org>
Reviewed-by: edchin <edchin@chromium.org>
Cr-Commit-Position: refs/heads/master@{#594811}

ios/chrome/browser/snapshots/snapshot_generator.mm

@@ -11,7 +11,9 @@
 
 #include <algorithm>
 
+#include "base/bind.h"
 #include "base/logging.h"
+#include "base/task/post_task.h"
 #include "ios/chrome/browser/browser_state/chrome_browser_state.h"
 #import "ios/chrome/browser/snapshots/snapshot_cache.h"
 #import "ios/chrome/browser/snapshots/snapshot_cache_factory.h"
@@ -22,6 +24,8 @@
 #import "ios/web/public/features.h"
 #import "ios/web/public/web_state/web_state.h"
 #import "ios/web/public/web_state/web_state_observer_bridge.h"
+#include "ios/web/public/web_task_traits.h"
+#include "ios/web/public/web_thread.h"
 #include "ui/gfx/image/image.h"
 
 #if !defined(__has_feature) || !__has_feature(objc_arc)
@@ -216,40 +220,57 @@ BOOL ViewHierarchyContainsWKWebView(UIView* view) {
 - (void)updateWebViewSnapshotWithCompletion:(void (^)(UIImage*))completion {
   DCHECK(_webState);
   CGRect frame = [self snapshotFrameVisibleFrameOnly:YES];
-  if (CGRectIsEmpty(frame))
+  if (CGRectIsEmpty(frame)) {
+    if (completion) {
+      base::PostTaskWithTraits(FROM_HERE, {web::WebThread::UI},
+                               base::BindOnce(^{
+                                 completion(nil);
+                               }));
+    }
     return;
+  }
+  DCHECK(std::isnormal(frame.size.width) && (frame.size.width > 0))
+      << ": frame.size.width=" << frame.size.width;
+  DCHECK(std::isnormal(frame.size.height) && (frame.size.height > 0))
+      << ": frame.size.height=" << frame.size.height;
+
   NSArray<SnapshotOverlay*>* overlays =
       [_delegate snapshotOverlaysForWebState:_webState];
   UIImage* snapshot =
       [_coalescingSnapshotContext cachedSnapshotWithOverlays:overlays
                                             visibleFrameOnly:YES];
   if (snapshot) {
-    if (completion)
-      completion(snapshot);
+    if (completion) {
+      base::PostTaskWithTraits(FROM_HERE, {web::WebThread::UI},
+                               base::BindOnce(^{
+                                 completion(nil);
+                               }));
+    }
     return;
   }
 
   [_delegate willUpdateSnapshotForWebState:_webState];
   __weak SnapshotGenerator* weakSelf = self;
-  _webState->TakeSnapshot(base::BindOnce(^(gfx::Image image) {
-    SnapshotGenerator* strongSelf = weakSelf;
-    if (!strongSelf)
-      return;
-    UIImage* snapshot = image.ToUIImage();
-    if (overlays.count > 0) {
-      snapshot = [strongSelf snapshotWithOverlays:overlays
-                                         snapshot:snapshot
-                                            frame:frame];
-    }
-    [strongSelf.snapshotCache setImage:snapshot
-                         withSessionID:_snapshotSessionId];
-    [_coalescingSnapshotContext setCachedSnapshot:snapshot
-                                     withOverlays:overlays
-                                 visibleFrameOnly:YES];
-    [_delegate didUpdateSnapshotForWebState:_webState withImage:snapshot];
-    if (completion)
-      completion(snapshot);
-  }));
+  _webState->TakeSnapshot(
+      frame, base::BindOnce(^(gfx::Image image) {
+        SnapshotGenerator* strongSelf = weakSelf;
+        if (!strongSelf)
+          return;
+        UIImage* snapshot = image.ToUIImage();
+        if (overlays.count > 0) {
+          snapshot = [strongSelf snapshotWithOverlays:overlays
+                                             snapshot:snapshot
+                                                frame:frame];
+        }
+        [strongSelf.snapshotCache setImage:snapshot
+                             withSessionID:_snapshotSessionId];
+        [_coalescingSnapshotContext setCachedSnapshot:snapshot
+                                         withOverlays:overlays
+                                     visibleFrameOnly:YES];
+        [_delegate didUpdateSnapshotForWebState:_webState withImage:snapshot];
+        if (completion)
+          completion(snapshot);
+      }));
 }
 
 - (UIImage*)generateSnapshotWithOverlays:(BOOL)shouldAddOverlay

ios/web/public/test/fakes/test_web_state.h

@@ -86,7 +86,7 @@ class TestWebState : public WebState {
   void DidChangeVisibleSecurityState() override {}
   bool HasOpener() const override;
   void SetHasOpener(bool has_opener) override;
-  void TakeSnapshot(SnapshotCallback callback) override;
+  void TakeSnapshot(CGRect rect, SnapshotCallback callback) override;
 
   // Setters for test data.
   void SetBrowserState(BrowserState* browser_state);

ios/web/public/test/fakes/test_web_state.mm

@@ -389,7 +389,7 @@ void TestWebState::SetHasOpener(bool has_opener) {
   has_opener_ = has_opener;
 }
 
-void TestWebState::TakeSnapshot(SnapshotCallback callback) {
+void TestWebState::TakeSnapshot(CGRect rect, SnapshotCallback callback) {
   std::move(callback).Run(gfx::Image([[UIImage alloc] init]));
 }
 

ios/web/public/web_state/web_state.h

@@ -7,6 +7,7 @@
 
 #include <stdint.h>
 
+#import <CoreGraphics/CoreGraphics.h>
 #include <memory>
 #include <string>
 #include <utility>
@@ -311,10 +312,10 @@ class WebState : public base::SupportsUserData {
   // Callback used to handle snapshots. The parameter is the snapshot image.
   typedef base::OnceCallback<void(gfx::Image)> SnapshotCallback;
 
-  // Takes a snapshot of this WebState. |callback| is
-  // asynchronously invoked after performing the snapshot. Prior to iOS 11, the
-  // callback is invoked with a nil snapshot.
-  virtual void TakeSnapshot(SnapshotCallback callback) = 0;
+  // Takes a snapshot of this WebState with |rect|. |callback| is asynchronously
+  // invoked after performing the snapshot. Prior to iOS 11, the callback is
+  // invoked with a nil snapshot.
+  virtual void TakeSnapshot(CGRect rect, SnapshotCallback callback) = 0;
 
   // Adds and removes observers for page navigation notifications. The order in
   // which notifications are sent to observers is undefined. Clients must be

ios/web/web_state/ui/crw_web_controller.h

@@ -211,10 +211,11 @@ class WebStateImpl;
             (web::NavigationInitiationType)type
                                           hasUserGesture:(BOOL)hasUserGesture;
 
-// Takes snapshot of web view. |completion| is always called,
-// but |snapshot| may be nil. Prior to iOS 11, |completion| is called with a nil
+// Takes snapshot of web view with |rect|. |completion| is always called, but
+// |snapshot| may be nil. Prior to iOS 11, |completion| is called with a nil
 // snapshot.
-- (void)takeSnapshotWithCompletion:(void (^)(UIImage* snapshot))completion;
+- (void)takeSnapshotWithRect:(CGRect)rect
+                  completion:(void (^)(UIImage* snapshot))completion;
 
 @end
 

ios/web/web_state/ui/crw_web_controller.mm

@@ -2282,11 +2282,15 @@ registerLoadRequestForURL:(const GURL&)requestURL
     self.navigationManagerImpl->DiscardNonCommittedItems();
 }
 
-- (void)takeSnapshotWithCompletion:(void (^)(UIImage*))completion {
+- (void)takeSnapshotWithRect:(CGRect)rect
+                  completion:(void (^)(UIImage*))completion {
   if (@available(iOS 11, *)) {
     if (_webView) {
+      WKSnapshotConfiguration* configuration =
+          [[WKSnapshotConfiguration alloc] init];
+      configuration.rect = rect;
       [_webView
-          takeSnapshotWithConfiguration:nil
+          takeSnapshotWithConfiguration:configuration
                       completionHandler:^(UIImage* snapshot, NSError* error) {
                         if (error)
                           DLOG(ERROR) << "WKWebView snapshot error: "

ios/web/web_state/web_state_impl.h

@@ -234,7 +234,7 @@ class WebStateImpl : public WebState, public NavigationManagerDelegate {
       mojo::ScopedMessagePipeHandle interface_pipe) override;
   bool HasOpener() const override;
   void SetHasOpener(bool has_opener) override;
-  void TakeSnapshot(SnapshotCallback callback) override;
+  void TakeSnapshot(CGRect rect, SnapshotCallback callback) override;
   void AddObserver(WebStateObserver* observer) override;
   void RemoveObserver(WebStateObserver* observer) override;
 

ios/web/web_state/web_state_impl.mm

@@ -727,11 +727,13 @@ void WebStateImpl::SetHasOpener(bool has_opener) {
   created_with_opener_ = has_opener;
 }
 
-void WebStateImpl::TakeSnapshot(SnapshotCallback callback) {
+void WebStateImpl::TakeSnapshot(CGRect rect, SnapshotCallback callback) {
   __block SnapshotCallback shared_callback = std::move(callback);
-  [web_controller_ takeSnapshotWithCompletion:^(UIImage* snapshot) {
-    std::move(shared_callback).Run(gfx::Image(snapshot));
-  }];
+  [web_controller_
+      takeSnapshotWithRect:rect
+                completion:^(UIImage* snapshot) {
+                  std::move(shared_callback).Run(gfx::Image(snapshot));
+                }];
 }
 
 void WebStateImpl::OnNavigationStarted(web::NavigationContext* context) {

ios/web/web_state/web_state_unittest.mm

@@ -210,27 +210,29 @@ TEST_P(WebStateTest, Snapshot) {
       addSubview:web_state()->GetView()];
   // The subview is added but not immediately painted, so a small delay is
   // necessary.
+  CGRect rect = [web_state()->GetView() bounds];
   base::test::ios::SpinRunLoopWithMinDelay(base::TimeDelta::FromSecondsD(0.2));
-  web_state()->TakeSnapshot(base::BindOnce(^(gfx::Image snapshot) {
-    if (@available(iOS 11, *)) {
-      ASSERT_FALSE(snapshot.IsEmpty());
-      EXPECT_GT(snapshot.Width(), 0);
-      EXPECT_GT(snapshot.Height(), 0);
-      int red_pixel_x = (snapshot.Width() / 2) - 10;
-      int white_pixel_x = (snapshot.Width() / 2) + 10;
-      // Test a pixel on the left (red) side.
-      gfx::test::CheckColors(
-          gfx::test::GetPlatformImageColor(gfx::test::ToPlatformType(snapshot),
-                                           red_pixel_x, 50),
-          SK_ColorRED);
-      // Test a pixel on the right (white) side.
-      gfx::test::CheckColors(
-          gfx::test::GetPlatformImageColor(gfx::test::ToPlatformType(snapshot),
-                                           white_pixel_x, 50),
-          SK_ColorWHITE);
-    }
-    snapshot_complete = true;
-  }));
+  web_state()->TakeSnapshot(
+      rect, base::BindOnce(^(gfx::Image snapshot) {
+        if (@available(iOS 11, *)) {
+          ASSERT_FALSE(snapshot.IsEmpty());
+          EXPECT_GT(snapshot.Width(), 0);
+          EXPECT_GT(snapshot.Height(), 0);
+          int red_pixel_x = (snapshot.Width() / 2) - 10;
+          int white_pixel_x = (snapshot.Width() / 2) + 10;
+          // Test a pixel on the left (red) side.
+          gfx::test::CheckColors(
+              gfx::test::GetPlatformImageColor(
+                  gfx::test::ToPlatformType(snapshot), red_pixel_x, 50),
+              SK_ColorRED);
+          // Test a pixel on the right (white) side.
+          gfx::test::CheckColors(
+              gfx::test::GetPlatformImageColor(
+                  gfx::test::ToPlatformType(snapshot), white_pixel_x, 50),
+              SK_ColorWHITE);
+        }
+        snapshot_complete = true;
+      }));
   WaitForCondition(^{
     return snapshot_complete;
   });