Store NavigationIntext in CRWWKNavigationStates.

ios/web/web_state/ui/crw_wk_navigation_states.mm

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #import "ios/web/web_state/ui/crw_wk_navigation_states.h"
 
 #include "base/logging.h"
+#include "ios/web/web_state/navigation_context_impl.h"
 
 #if !defined(__has_feature) || !__has_feature(objc_arc)
 #error "This file requires ARC support."
 #endif
 
 // Holds a pair of state and creation order index.
-@interface CRWWKNavigationsStateRecord : NSObject
+@interface CRWWKNavigationsStateRecord : NSObject {
+  // Backs up |context| property.
+  std::unique_ptr<web::NavigationContextImpl> _context;
+}
 // Navigation state.
 @property(nonatomic, assign) web::WKNavigationState state;
 // Numerical index representing creation order (smaller index denotes earlier
 // navigations).
 @property(nonatomic, assign, readonly) NSUInteger index;
 
 - (instancetype)init NS_UNAVAILABLE;
 
 // Initializes record with state and index values.
 - (instancetype)initWithState:(web::WKNavigationState)state
                         index:(NSUInteger)index NS_DESIGNATED_INITIALIZER;
 
+// Initializes record with context and index values.
+- (instancetype)initWithContext:
+                    (std::unique_ptr<web::NavigationContextImpl>)context
+                          index:(NSUInteger)index NS_DESIGNATED_INITIALIZER;
+
+// web::NavigationContextImpl for this navigation.
+- (web::NavigationContextImpl*)context;
+- (void)setContext:(std::unique_ptr<web::NavigationContextImpl>)context;
+
 @end
 
 @implementation CRWWKNavigationsStateRecord
 @synthesize state = _state;
 @synthesize index = _index;
 
+#ifndef NDEBUG
 - (NSString*)description {
-  return [NSString stringWithFormat:@"state: %d, index: %ld", _state,
-                                    static_cast<long>(_index)];
+  return [NSString stringWithFormat:@"state: %d, index: %ld, context: %@",
+                                    _state, static_cast<long>(_index),
+                                    _context->GetDescription()];
 }
+#endif  // NDEBUG
 
 - (instancetype)initWithState:(web::WKNavigationState)state
                         index:(NSUInteger)index {
   if ((self = [super init])) {
     _state = state;
     _index = index;
   }
   return self;
 }
 
+- (instancetype)initWithContext:
+                    (std::unique_ptr<web::NavigationContextImpl>)context
+                          index:(NSUInteger)index {
+  if ((self = [super init])) {
+    _context = std::move(context);
+    _index = index;
+  }
+  return self;
+}
+
+- (void)setContext:(std::unique_ptr<web::NavigationContextImpl>)context {
+  _context = std::move(context);
+}
+
+- (web::NavigationContextImpl*)context {
+  return _context.get();
+}
+
 @end
 
 @interface CRWWKNavigationStates () {
   NSMapTable* _records;
   NSUInteger _lastStateIndex;
 }
 @end
 
 @implementation CRWWKNavigationStates
 
@@ skipping 36 matching lines @@
 - (void)removeNavigation:(WKNavigation*)navigation {
   if (!navigation) {
     // WKWebView may call WKNavigationDelegate callbacks with nil.
     return;
   }
 
   DCHECK([_records objectForKey:navigation]);
   [_records removeObjectForKey:navigation];
 }
 
+- (void)setContext:(std::unique_ptr<web::NavigationContextImpl>)context
+     forNavigation:(WKNavigation*)navigation {
+  if (!navigation) {
+    // WKWebView may call WKNavigationDelegate callbacks with nil.
+    return;
+  }
+  CRWWKNavigationsStateRecord* record = [_records objectForKey:navigation];
+  if (!record) {
+    record =
+        [[CRWWKNavigationsStateRecord alloc] initWithContext:std::move(context)
+                                                       index:++_lastStateIndex];
+  } else {
+    [record setContext:std::move(context)];
+  }
+  [_records setObject:record forKey:navigation];
+}
+
+- (web::NavigationContextImpl*)contextForNavigation:(WKNavigation*)navigation {
+  if (!navigation) {
+    // WKWebView may call WKNavigationDelegate callbacks with nil.
+    return nullptr;
+  }
+  CRWWKNavigationsStateRecord* record = [_records objectForKey:navigation];
+  return [record context];
+}
+
 - (WKNavigation*)lastAddedNavigation {
   WKNavigation* result = nil;
   NSUInteger lastAddedIndex = 0;  // record indices start with 1.
   for (WKNavigation* navigation in _records) {
     CRWWKNavigationsStateRecord* record = [_records objectForKey:navigation];
     if (lastAddedIndex < record.index) {
       result = navigation;
       lastAddedIndex = record.index;
     }
   }
   return result;
 }
 
 - (web::WKNavigationState)lastAddedNavigationState {
   CRWWKNavigationsStateRecord* lastAddedRecord = nil;
   WKNavigation* lastAddedNavigation = [self lastAddedNavigation];
   if (lastAddedNavigation)
     lastAddedRecord = [_records objectForKey:lastAddedNavigation];
 
   return lastAddedRecord.state;
 }
 
 @end