Update the defaultLayout property on the Grid View Controller.

Refactor |willTransitionToLayout:completion| and update |defaultLayout|
after a layout transition.

Bug: 1127604
Change-Id: I46832b073db56b189cb46f046229e7e99c3f48ee
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2450278
Auto-Submit: Roberto Moura <mouraroberto@google.com>
Commit-Queue: Gauthier Ambard <gambard@chromium.org>
Reviewed-by: Gauthier Ambard <gambard@chromium.org>
Cr-Commit-Position: refs/heads/master@{#814580}

ios/chrome/browser/ui/tab_grid/grid/grid_view_controller.mm

@@ -696,22 +696,23 @@ NSIndexPath* CreateIndexPath(NSInteger index) {
 - (void)willTransitionToLayout:(LayoutSwitcherState)nextState
                     completion:
                         (void (^)(BOOL completed, BOOL finished))completion {
-  auto completionBlock = ^(BOOL completed, BOOL finished) {
-    completion(completed, finished);
-    self.collectionView.scrollEnabled = YES;
-  };
+  TabSwitcherLayout* nextLayout;
   switch (nextState) {
     case LayoutSwitcherState::Horizontal:
-      self.gridHorizontalTransitionLayout = [self.collectionView
-          startInteractiveTransitionToCollectionViewLayout:self.horizontalLayout
-                                                completion:completionBlock];
+      nextLayout = self.horizontalLayout;
       break;
     case LayoutSwitcherState::Full:
-      self.gridHorizontalTransitionLayout = [self.collectionView
-          startInteractiveTransitionToCollectionViewLayout:self.gridLayout
-                                                completion:completionBlock];
+      nextLayout = self.gridLayout;
       break;
   }
+  auto completionBlock = ^(BOOL completed, BOOL finished) {
+    completion(completed, finished);
+    self.collectionView.scrollEnabled = YES;
+    self.defaultLayout = nextLayout;
+  };
+  self.gridHorizontalTransitionLayout = [self.collectionView
+      startInteractiveTransitionToCollectionViewLayout:nextLayout
+                                            completion:completionBlock];
 
   // Stops collectionView scrolling when the animation starts.
   [self.collectionView setContentOffset:self.collectionView.contentOffset