[iOS] Use InfoBarIOS::Observer to update InfobarBadgeTabModel state.

This CL updates InfobarBadgeTabModel to use InfoBarIOS::Observer
callbacks to detect infobar accept/revert events, deprecating the
public API previously used to relay this information.

Additionally, this CL decomposes observer functionality into internal
classes in order to allow easily hooking up multiple scoped observers.
Using nested classes also helps to minimize the publicly-accessible
InfobarBadgeTabHelper API, preventing it from being misused.

This CL also updates the badge models to be stored in an std::map to
avoid the extra step of converting the InfobarType to an NSNumber to use
for NSDictionary lookups.

Bug: 1030357
Change-Id: I5f22f81fc8f1838b1cf235584390f7d5af54ea30
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1977026
Commit-Queue: Kurt Horimoto <kkhorimoto@chromium.org>
Reviewed-by: Sergio Collazos <sczs@chromium.org>
Cr-Commit-Position: refs/heads/master@{#727273}

ios/chrome/browser/infobars/infobar_badge_tab_helper.h

@@ -5,10 +5,11 @@
 #ifndef IOS_CHROME_BROWSER_INFOBARS_INFOBAR_BADGE_TAB_HELPER_H_
 #define IOS_CHROME_BROWSER_INFOBARS_INFOBAR_BADGE_TAB_HELPER_H_
 
-#include <unordered_map>
+#include <map>
 
 #include "base/scoped_observer.h"
 #include "components/infobars/core/infobar_manager.h"
+#include "ios/chrome/browser/infobars/infobar_ios.h"
 #import "ios/chrome/browser/infobars/infobar_type.h"
 #import "ios/chrome/browser/ui/badges/badge_item.h"
 #import "ios/web/public/web_state_user_data.h"
@@ -24,19 +25,12 @@ class WebState;
 // TabHelper that observes InfoBarManager. It updates an InfobarBadge delegate
 // for relevant Infobar changes.
 class InfobarBadgeTabHelper
-    : public infobars::InfoBarManager::Observer,
-      public web::WebStateUserData<InfobarBadgeTabHelper> {
+    : public web::WebStateUserData<InfobarBadgeTabHelper> {
  public:
-  // Creates the tab helper for |web_state| if it does not exist.
-  static void CreateForWebState(web::WebState* web_state);
+  ~InfobarBadgeTabHelper() override;
+
   // Sets the InfobarBadgeTabHelperDelegate to |delegate|.
   void SetDelegate(id<InfobarBadgeTabHelperDelegate> delegate);
-  // Updates Infobar badge for the case where Infobar of |infobar_type| was
-  // accepted.
-  void UpdateBadgeForInfobarAccepted(InfobarType infobar_type);
-  // Updates Infobar badge for the case where Infobar action of |infobar_type|
-  // was reverted. This will negate the accepted state of the badge.
-  void UpdateBadgeForInfobarReverted(InfobarType infobar_type);
   // Updates Infobar badge for the case where an Infobar banner of
   // |infobar_type| was presented.
   void UpdateBadgeForInfobarBannerPresented(InfobarType infobar_type);
@@ -47,30 +41,78 @@ class InfobarBadgeTabHelper
   // Returns all BadgeItems for the TabHelper Webstate.
   NSArray<id<BadgeItem>>* GetInfobarBadgeItems();
 
-  ~InfobarBadgeTabHelper() override;
+  // DEPRECATED: The accept state of an infobar is now stored directly in
+  // InfoBarIOS, and should be updated there rather than using these functions.
+  void UpdateBadgeForInfobarAccepted(InfobarType infobar_type);
+  void UpdateBadgeForInfobarReverted(InfobarType infobar_type);
 
- protected:
-  // Constructor.
+ private:
+  friend class web::WebStateUserData<InfobarBadgeTabHelper>;
   explicit InfobarBadgeTabHelper(web::WebState* web_state);
+
+  // Notifies the tab helper to reset state for added or removed infobars with
+  // |infobar_type|.  Only called for infobars that support badges.
+  void ResetStateForAddedInfobar(InfobarType infobar_type);
+  void ResetStateForRemovedInfobar(InfobarType infobar_type);
+  // Notifies the tab helper that an infobar with |type| was accepted or
+  // reverted.
+  void OnInfobarAcceptanceStateChanged(InfobarType infobar_type, bool accepted);
+
+  // Helper object that listens for accept and revert events for an InfoBarIOS.
+  class InfobarAcceptanceObserver : public InfoBarIOS::Observer {
+   public:
+    explicit InfobarAcceptanceObserver(InfobarBadgeTabHelper* tab_helper);
+    ~InfobarAcceptanceObserver() override;
+
+    // Returns a reference to the scoped observer.
+    ScopedObserver<InfoBarIOS, InfoBarIOS::Observer>& scoped_observer() {
+      return scoped_observer_;
+    }
+
+   private:
+    // InfoBarIOS::Observer:
+    void DidUpdateAcceptedState(InfoBarIOS* infobar) override;
+    void InfobarDestroyed(InfoBarIOS* infobar) override;
+
+    // The owning tab helper.
+    InfobarBadgeTabHelper* tab_helper_ = nullptr;
+    // Scoped observer that facilitates observing InfoBarIOS objects.
+    ScopedObserver<InfoBarIOS, InfoBarIOS::Observer> scoped_observer_;
+  };
+
+  // Helper object that updates state and adds an InfobarAcceptanceObserver
+  // when each infobar is added or removed.
+  class InfobarManagerObserver : public infobars::InfoBarManager::Observer {
+   public:
+    InfobarManagerObserver(InfobarBadgeTabHelper* tab_helper,
+                           web::WebState* web_state,
+                           InfobarAcceptanceObserver* infobar_accept_observer);
+    ~InfobarManagerObserver() override;
+
+   private:
+    // InfoBarManagerObserver:
+    void OnInfoBarAdded(infobars::InfoBar* infobar) override;
+    void OnInfoBarRemoved(infobars::InfoBar* infobar, bool animate) override;
+    void OnManagerShuttingDown(infobars::InfoBarManager* manager) override;
+
+    // The owning tab helper.
+    InfobarBadgeTabHelper* tab_helper_ = nullptr;
+    // The infobar acceptance observer for |tab_helper_|.  Added to each infobar
+    // in the observed manager.
+    InfobarAcceptanceObserver* infobar_accept_observer_ = nullptr;
+    // Scoped observer that facilitates observing an InfoBarManager.
+    ScopedObserver<infobars::InfoBarManager, infobars::InfoBarManager::Observer>
+        scoped_observer_;
+  };
+
   // Delegate which displays the Infobar badge.
   __weak id<InfobarBadgeTabHelperDelegate> delegate_ = nil;
-  // Holds the state of each displaying badge keyed by its InfobarType.
-  NSMutableDictionary<NSNumber*, InfobarBadgeModel*>* infobar_badge_models_;
-
- private:
-  friend class web::WebStateUserData<InfobarBadgeTabHelper>;
-  // Helper method to get Obj-C key for InfobarType.
-  NSNumber* GetNumberKeyForInfobarType(InfobarType infobar_type);
-  // InfoBarManagerObserver implementation.
-  void OnInfoBarAdded(infobars::InfoBar* infobar) override;
-  void OnInfoBarRemoved(infobars::InfoBar* infobar, bool animate) override;
-  void OnManagerShuttingDown(infobars::InfoBarManager* manager) override;
-  // Updates the badge delegate for |infobar|.
-  void UpdateBadgeForInfobar(infobars::InfoBar* infobar, bool display);
-
-  // Manages this object as an observer of infobars.
-  ScopedObserver<infobars::InfoBarManager, infobars::InfoBarManager::Observer>
-      infobar_observer_;
+  // The infobar accept/revert observer.
+  InfobarAcceptanceObserver infobar_accept_observer_;
+  // The infobar manager observer.
+  InfobarManagerObserver infobar_manager_observer_;
+  // Map storing the badge models for each InfobarType.
+  std::map<InfobarType, InfobarBadgeModel*> infobar_badge_models_;
 
   WEB_STATE_USER_DATA_KEY_DECL();
   DISALLOW_COPY_AND_ASSIGN(InfobarBadgeTabHelper);

ios/chrome/browser/infobars/infobar_badge_tab_helper.mm

@@ -6,7 +6,6 @@
 
 #include "ios/chrome/browser/infobars/infobar_badge_model.h"
 #include "ios/chrome/browser/infobars/infobar_badge_tab_helper_delegate.h"
-#include "ios/chrome/browser/infobars/infobar_ios.h"
 #include "ios/chrome/browser/infobars/infobar_manager_impl.h"
 #import "ios/chrome/browser/ui/infobars/infobar_feature.h"
 #import "ios/chrome/browser/ui/infobars/infobar_ui_delegate.h"
@@ -15,16 +14,28 @@
 #error "This file requires ARC support."
 #endif
 
-#pragma mark - Public
-
-// static
-void InfobarBadgeTabHelper::CreateForWebState(web::WebState* web_state) {
-  DCHECK(web_state);
-  if (!FromWebState(web_state)) {
-    web_state->SetUserData(
-        UserDataKey(), base::WrapUnique(new InfobarBadgeTabHelper(web_state)));
-  }
+namespace {
+// Returns |infobar|'s InfobarType.
+InfobarType GetInfobarType(infobars::InfoBar* infobar) {
+  return static_cast<InfoBarIOS*>(infobar)->InfobarUIDelegate().infobarType;
 }
+// Returns whether |infobar| supports badges.
+bool SupportsBadges(infobars::InfoBar* infobar) {
+  return static_cast<InfoBarIOS*>(infobar)->InfobarUIDelegate().hasBadge;
+}
+}  // namespace
+
+#pragma mark - InfobarBadgeTabHelper
+
+WEB_STATE_USER_DATA_KEY_IMPL(InfobarBadgeTabHelper)
+
+InfobarBadgeTabHelper::InfobarBadgeTabHelper(web::WebState* web_state)
+    : infobar_accept_observer_(this),
+      infobar_manager_observer_(this, web_state, &infobar_accept_observer_) {}
+
+InfobarBadgeTabHelper::~InfobarBadgeTabHelper() = default;
+
+#pragma mark Public
 
 void InfobarBadgeTabHelper::SetDelegate(
     id<InfobarBadgeTabHelperDelegate> delegate) {
@@ -33,98 +44,124 @@ void InfobarBadgeTabHelper::SetDelegate(
 
 void InfobarBadgeTabHelper::UpdateBadgeForInfobarAccepted(
     InfobarType infobar_type) {
-  NSNumber* infobar_type_key = GetNumberKeyForInfobarType(infobar_type);
-  infobar_badge_models_[infobar_type_key].badgeState |=
-      BadgeStateAccepted | BadgeStateRead;
-  [delegate_ updateInfobarBadge:infobar_badge_models_[infobar_type_key]];
+  OnInfobarAcceptanceStateChanged(infobar_type, /*accepted=*/true);
 }
 
 void InfobarBadgeTabHelper::UpdateBadgeForInfobarReverted(
     InfobarType infobar_type) {
-  NSNumber* infobar_type_key = GetNumberKeyForInfobarType(infobar_type);
-  infobar_badge_models_[infobar_type_key].badgeState &= ~BadgeStateAccepted;
-  [delegate_ updateInfobarBadge:infobar_badge_models_[infobar_type_key]];
+  OnInfobarAcceptanceStateChanged(infobar_type, /*accepted=*/false);
 }
 
 void InfobarBadgeTabHelper::UpdateBadgeForInfobarBannerPresented(
     InfobarType infobar_type) {
-  NSNumber* infobar_type_key = GetNumberKeyForInfobarType(infobar_type);
-  infobar_badge_models_[infobar_type_key].badgeState |= BadgeStatePresented;
-  [delegate_ updateInfobarBadge:infobar_badge_models_[infobar_type_key]];
+  infobar_badge_models_[infobar_type].badgeState |= BadgeStatePresented;
+  [delegate_ updateInfobarBadge:infobar_badge_models_[infobar_type]];
 }
 
 void InfobarBadgeTabHelper::UpdateBadgeForInfobarBannerDismissed(
     InfobarType infobar_type) {
-  NSNumber* infobar_type_key = GetNumberKeyForInfobarType(infobar_type);
-  infobar_badge_models_[infobar_type_key].badgeState &= ~BadgeStatePresented;
-  [delegate_ updateInfobarBadge:infobar_badge_models_[infobar_type_key]];
+  infobar_badge_models_[infobar_type].badgeState &= ~BadgeStatePresented;
+  [delegate_ updateInfobarBadge:infobar_badge_models_[infobar_type]];
 }
 
 NSArray<id<BadgeItem>>* InfobarBadgeTabHelper::GetInfobarBadgeItems() {
-  return [infobar_badge_models_ allValues];
+  NSMutableArray* badge_items = [NSMutableArray array];
+  for (auto& infobar_type_badge_model_pair : infobar_badge_models_) {
+    id<BadgeItem> badge = infobar_type_badge_model_pair.second;
+    if (badge)
+      [badge_items addObject:badge];
+  }
+  return badge_items;
 }
 
-InfobarBadgeTabHelper::~InfobarBadgeTabHelper() = default;
+#pragma mark Private
 
-#pragma mark - Private
+void InfobarBadgeTabHelper::ResetStateForAddedInfobar(
+    InfobarType infobar_type) {
+  InfobarBadgeModel* new_badge =
+      [[InfobarBadgeModel alloc] initWithInfobarType:infobar_type];
+  infobar_badge_models_[infobar_type] = new_badge;
+  [delegate_ addInfobarBadge:new_badge];
+}
 
-InfobarBadgeTabHelper::InfobarBadgeTabHelper(web::WebState* web_state)
-    : infobar_observer_(this) {
-  infobars::InfoBarManager* infoBarManager =
-      InfoBarManagerImpl::FromWebState(web_state);
-  if (infoBarManager) {
-    infobar_observer_.Add(infoBarManager);
+void InfobarBadgeTabHelper::ResetStateForRemovedInfobar(
+    InfobarType infobar_type) {
+  InfobarBadgeModel* removed_badge = infobar_badge_models_[infobar_type];
+  infobar_badge_models_[infobar_type] = nil;
+  [delegate_ removeInfobarBadge:removed_badge];
+}
+
+void InfobarBadgeTabHelper::OnInfobarAcceptanceStateChanged(
+    InfobarType infobar_type,
+    bool accepted) {
+  id<BadgeItem> item = infobar_badge_models_[infobar_type];
+  if (accepted) {
+    item.badgeState |= BadgeStateAccepted | BadgeStateRead;
+  } else {
+    item.badgeState &= ~BadgeStateAccepted;
   }
-  // Since the TabHelper is being created, |infobar_badge_models_| needs to be
-  // properly initialized.
-  infobar_badge_models_ = [NSMutableDictionary dictionary];
+  [delegate_ updateInfobarBadge:item];
 }
 
-NSNumber* InfobarBadgeTabHelper::GetNumberKeyForInfobarType(
-    InfobarType infobar_type) {
-  NSNumber* infobar_type_key = [NSNumber numberWithInt:(int)infobar_type];
-  DCHECK(infobar_type_key);
-  return infobar_type_key;
+#pragma mark - InfobarBadgeTabHelper::InfobarAcceptanceObserver
+
+InfobarBadgeTabHelper::InfobarAcceptanceObserver::InfobarAcceptanceObserver(
+    InfobarBadgeTabHelper* tab_helper)
+    : tab_helper_(tab_helper), scoped_observer_(this) {
+  DCHECK(tab_helper_);
 }
 
-#pragma mark InfobarObserver
+InfobarBadgeTabHelper::InfobarAcceptanceObserver::~InfobarAcceptanceObserver() =
+    default;
 
-void InfobarBadgeTabHelper::OnInfoBarAdded(infobars::InfoBar* infobar) {
-  this->UpdateBadgeForInfobar(infobar, true);
+void InfobarBadgeTabHelper::InfobarAcceptanceObserver::DidUpdateAcceptedState(
+    InfoBarIOS* infobar) {
+  tab_helper_->OnInfobarAcceptanceStateChanged(GetInfobarType(infobar),
+                                               infobar->accepted());
 }
 
-void InfobarBadgeTabHelper::OnInfoBarRemoved(infobars::InfoBar* infobar,
-                                             bool animate) {
-  this->UpdateBadgeForInfobar(infobar, false);
+void InfobarBadgeTabHelper::InfobarAcceptanceObserver::InfobarDestroyed(
+    InfoBarIOS* infobar) {
+  scoped_observer_.Remove(infobar);
 }
 
-void InfobarBadgeTabHelper::OnManagerShuttingDown(
-    infobars::InfoBarManager* manager) {
-  infobar_observer_.Remove(manager);
-}
-
-#pragma mark Helpers
-
-void InfobarBadgeTabHelper::UpdateBadgeForInfobar(infobars::InfoBar* infobar,
-                                                  bool display) {
-  InfoBarIOS* infobar_ios = static_cast<InfoBarIOS*>(infobar);
-  id<InfobarUIDelegate> controller_ = infobar_ios->InfobarUIDelegate();
-  if (IsInfobarUIRebootEnabled() && controller_.hasBadge) {
-    InfobarType infobar_type = controller_.infobarType;
-    if (display) {
-      InfobarBadgeModel* new_badge =
-          [[InfobarBadgeModel alloc] initWithInfobarType:infobar_type];
-      NSNumber* infobar_type_key = [NSNumber numberWithInt:(int)infobar_type];
-      infobar_badge_models_[infobar_type_key] = new_badge;
-      [delegate_ addInfobarBadge:new_badge];
-    } else {
-      NSNumber* infobar_type_key = [NSNumber numberWithInt:(int)infobar_type];
-      InfobarBadgeModel* removed_badge =
-          infobar_badge_models_[infobar_type_key];
-      [infobar_badge_models_ removeObjectForKey:infobar_type_key];
-      [delegate_ removeInfobarBadge:removed_badge];
-    }
+#pragma mark - InfobarBadgeTabHelper::InfobarManagerObserver
+
+InfobarBadgeTabHelper::InfobarManagerObserver::InfobarManagerObserver(
+    InfobarBadgeTabHelper* tab_helper,
+    web::WebState* web_state,
+    InfobarAcceptanceObserver* infobar_accept_observer)
+    : tab_helper_(tab_helper),
+      infobar_accept_observer_(infobar_accept_observer),
+      scoped_observer_(this) {
+  DCHECK(tab_helper_);
+  DCHECK(infobar_accept_observer_);
+  scoped_observer_.Add(InfoBarManagerImpl::FromWebState(web_state));
+}
+
+InfobarBadgeTabHelper::InfobarManagerObserver::~InfobarManagerObserver() =
+    default;
+
+void InfobarBadgeTabHelper::InfobarManagerObserver::OnInfoBarAdded(
+    infobars::InfoBar* infobar) {
+  if (SupportsBadges(infobar)) {
+    tab_helper_->ResetStateForAddedInfobar(GetInfobarType(infobar));
+    infobar_accept_observer_->scoped_observer().Add(
+        static_cast<InfoBarIOS*>(infobar));
   }
 }
 
-WEB_STATE_USER_DATA_KEY_IMPL(InfobarBadgeTabHelper)
+void InfobarBadgeTabHelper::InfobarManagerObserver::OnInfoBarRemoved(
+    infobars::InfoBar* infobar,
+    bool animate) {
+  if (SupportsBadges(infobar)) {
+    tab_helper_->ResetStateForRemovedInfobar(GetInfobarType(infobar));
+    infobar_accept_observer_->scoped_observer().Remove(
+        static_cast<InfoBarIOS*>(infobar));
+  }
+}
+
+void InfobarBadgeTabHelper::InfobarManagerObserver::OnManagerShuttingDown(
+    infobars::InfoBarManager* manager) {
+  scoped_observer_.Remove(manager);
+}

ios/chrome/browser/ui/infobars/infobar_container_coordinator_unittest.mm

@@ -89,9 +89,9 @@ class InfobarContainerCoordinatorTest : public PlatformTest {
     web_state_list_->ActivateWebStateAt(0);
 
     // Setup InfobarBadgeTabHelper and InfoBarManager
+    InfoBarManagerImpl::CreateForWebState(web_state_list_->GetActiveWebState());
     InfobarBadgeTabHelper::CreateForWebState(
         web_state_list_->GetActiveWebState());
-    InfoBarManagerImpl::CreateForWebState(web_state_list_->GetActiveWebState());
 
     // Setup the InfobarContainerCoordinator.
     infobar_container_coordinator_ = [[InfobarContainerCoordinator alloc]
@@ -188,10 +188,11 @@ class InfobarContainerCoordinatorTest : public PlatformTest {
   void AddSecondWebstate() {
     std::unique_ptr<web::TestWebState> second_web_state =
         std::make_unique<web::TestWebState>();
+    InfoBarManagerImpl::CreateForWebState(second_web_state.get());
+    InfobarBadgeTabHelper::CreateForWebState(second_web_state.get());
     web_state_list_->InsertWebState(1, std::move(second_web_state),
                                     WebStateList::INSERT_NO_FLAGS,
                                     WebStateOpener());
-    InfobarBadgeTabHelper::CreateForWebState(web_state_list_->GetWebStateAt(1));
   }
 
   // Adds a Legacy Infobar to the InfobarManager, triggering an InfobarBanner