[ios] Move snapshot image during drag and drop

Previously, there was only one snapshot cache so the image did
not need to be moved during drag and drop. Recently, there
was a change to have a snapshot cache per browser.

This CL moves the snapshot image from the source cache to the
destination cache when the corresponding WebState is moved.

Bug: 1132306
Change-Id: I30599d17e33f9610adc7d715675bcbf829eadf8d
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2438312
Commit-Queue: edchin <edchin@chromium.org>
Auto-Submit: edchin <edchin@chromium.org>
Reviewed-by: Sergio Collazos <sczs@chromium.org>
Cr-Commit-Position: refs/heads/master@{#811943}

ios/chrome/browser/main/BUILD.gn

@@ -104,6 +104,7 @@ source_set("unit_tests") {
     "//base",
     "//base/test:test_support",
     "//ios/chrome/browser/browser_state:test_support",
+    "//ios/chrome/browser/snapshots",
     "//ios/chrome/browser/tabs",
     "//ios/chrome/browser/web:tab_id_tab_helper",
     "//ios/chrome/browser/web_state_list",

ios/chrome/browser/main/browser_util.mm

@@ -13,6 +13,8 @@
 #import "ios/chrome/browser/main/browser.h"
 #import "ios/chrome/browser/main/browser_list.h"
 #import "ios/chrome/browser/main/browser_list_factory.h"
+#import "ios/chrome/browser/snapshots/snapshot_browser_agent.h"
+#import "ios/chrome/browser/snapshots/snapshot_cache.h"
 #import "ios/chrome/browser/web/tab_id_tab_helper.h"
 #import "ios/chrome/browser/web_state_list/web_state_list.h"
 #import "ios/chrome/browser/web_state_list/web_state_opener.h"
@@ -58,6 +60,28 @@ Browser* FindBrowser(NSString* tab_id,
                                  tab_index);
 }
 
+// Moves snapshot associated with |snapshot_id| from |source_browser| to
+// |destination_browser|'s snapshot cache.
+void MoveSnapshot(NSString* snapshot_id,
+                  Browser* source_browser,
+                  Browser* destination_browser) {
+  DCHECK(snapshot_id.length);
+  SnapshotCache* source_cache =
+      SnapshotBrowserAgent::FromBrowser(source_browser)->snapshot_cache();
+  SnapshotCache* destination_cache =
+      SnapshotBrowserAgent::FromBrowser(destination_browser)->snapshot_cache();
+  [source_cache
+      retrieveImageForSnapshotID:snapshot_id
+                        callback:^(UIImage* snapshot) {
+                          if (snapshot) {
+                            [destination_cache setImage:snapshot
+                                         withSnapshotID:snapshot_id];
+                            [source_cache
+                                removeImageWithSnapshotID:snapshot_id];
+                          }
+                        }];
+}
+
 }  // namespace
 
 void MoveTabToBrowser(NSString* tab_id,
@@ -81,6 +105,7 @@ void MoveTabToBrowser(NSString* tab_id,
         source_tab_index, destination_tab_index);
     return;
   }
+  MoveSnapshot(tab_id, source_browser, destination_browser);
   std::unique_ptr<web::WebState> web_state =
       source_browser->GetWebStateList()->DetachWebStateAt(source_tab_index);
   destination_browser->GetWebStateList()->InsertWebState(

ios/chrome/browser/main/browser_util_unittest.mm

@@ -10,6 +10,8 @@
 #import "ios/chrome/browser/main/browser_list.h"
 #import "ios/chrome/browser/main/browser_list_factory.h"
 #include "ios/chrome/browser/main/test_browser.h"
+#import "ios/chrome/browser/snapshots/snapshot_browser_agent.h"
+#import "ios/chrome/browser/snapshots/snapshot_tab_helper.h"
 #import "ios/chrome/browser/web/tab_id_tab_helper.h"
 #import "ios/chrome/browser/web_state_list/web_state_list.h"
 #import "ios/chrome/browser/web_state_list/web_state_opener.h"
@@ -43,6 +45,11 @@ class BrowserUtilTest : public PlatformTest {
     browser_list_->AddIncognitoBrowser(incognito_browser_.get());
     browser_list_->AddIncognitoBrowser(other_incognito_browser_.get());
 
+    SnapshotBrowserAgent::CreateForBrowser(browser_.get());
+    SnapshotBrowserAgent::CreateForBrowser(other_browser_.get());
+    SnapshotBrowserAgent::CreateForBrowser(incognito_browser_.get());
+    SnapshotBrowserAgent::CreateForBrowser(other_incognito_browser_.get());
+
     AppendNewWebState(browser_.get());
     AppendNewWebState(browser_.get());
     AppendNewWebState(browser_.get());
@@ -53,10 +60,13 @@ class BrowserUtilTest : public PlatformTest {
   web::TestWebState* AppendNewWebState(Browser* browser) {
     auto test_web_state = std::make_unique<web::TestWebState>();
     web::TestWebState* inserted_web_state = test_web_state.get();
+    TabIdTabHelper::CreateForWebState(inserted_web_state);
+    NSString* tab_id =
+        TabIdTabHelper::FromWebState(inserted_web_state)->tab_id();
+    SnapshotTabHelper::CreateForWebState(inserted_web_state, tab_id);
     browser->GetWebStateList()->InsertWebState(
         WebStateList::kInvalidIndex, std::move(test_web_state),
         WebStateList::INSERT_ACTIVATE, WebStateOpener());
-    TabIdTabHelper::CreateForWebState(inserted_web_state);
     return inserted_web_state;
   }