Order the script badges in the location bar consistently

chrome/browser/extensions/script_badge_controller.cc

 // Copyright (c) 2012 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.
 
 #include "chrome/browser/extensions/script_badge_controller.h"
 
 #include "chrome/browser/extensions/extension_service.h"
 #include "chrome/browser/extensions/extension_system.h"
 #include "chrome/browser/ui/tab_contents/tab_contents.h"
 #include "chrome/common/extensions/extension.h"
 #include "chrome/common/extensions/extension_action.h"
 #include "chrome/common/extensions/extension_messages.h"
 #include "chrome/common/extensions/extension_set.h"
 #include "chrome/common/chrome_notification_types.h"
 #include "content/public/browser/navigation_controller.h"
 #include "content/public/browser/navigation_entry.h"
 #include "content/public/browser/notification_service.h"
 #include "content/public/browser/web_contents.h"
 #include "ipc/ipc_message.h"
 #include "ipc/ipc_message_macros.h"
 
 namespace extensions {
 
 ScriptBadgeController::ScriptBadgeController(TabContents* tab_contents)
     : content::WebContentsObserver(tab_contents->web_contents()),
       script_executor_(tab_contents->web_contents()),
-      tab_contents_(tab_contents) {}
+      tab_contents_(tab_contents) {
+  registrar_.Add(this,
+                 chrome::NOTIFICATION_EXTENSION_UNLOADED,
+                 content::Source<Profile>(tab_contents->profile()));
+}
 
 ScriptBadgeController::~ScriptBadgeController() {}
 
-scoped_ptr<std::vector<ExtensionAction*> >
-ScriptBadgeController::GetCurrentActions() {
-  scoped_ptr<std::vector<ExtensionAction*> > current_actions(
-      new std::vector<ExtensionAction*>());
-
-  ExtensionService* service = GetExtensionService();
-  if (!service)
-    return current_actions.Pass();
-
-  for (ExtensionSet::const_iterator it = service->extensions()->begin();
-       it != service->extensions()->end(); ++it) {
-    const Extension* extension = *it;
-    if (extensions_executing_scripts_.count(extension->id()))
-      current_actions->push_back(extension->GetScriptBadge());
-  }
-  return current_actions.Pass();
+std::vector<ExtensionAction*> ScriptBadgeController::GetCurrentActions() {
+  return current_actions_;
 }
 
 LocationBarController::Action ScriptBadgeController::OnClicked(
     const std::string& extension_id, int mouse_button) {
   ExtensionService* service = GetExtensionService();
   if (!service)
     return ACTION_NONE;
 
   const Extension* extension = service->extensions()->GetByID(extension_id);
   CHECK(extension);
@@ skipping 34 matching lines @@
                                  this_callback);
 }
 
 void ScriptBadgeController::OnExecuteScriptFinished(
     const std::string& extension_id,
     const ExecuteScriptCallback& callback,
     bool success,
     int32 page_id,
     const std::string& error) {
   if (success && page_id == GetPageID()) {
-    extensions_executing_scripts_.insert(extension_id);
-    Notify();
+    if (InsertExtension(extension_id))
+      Notify();
   }
 
   callback.Run(success, page_id, error);
 }
 
 ExtensionService* ScriptBadgeController::GetExtensionService() {
   return ExtensionSystem::Get(tab_contents_->profile())->extension_service();
 }
 
 int32 ScriptBadgeController::GetPageID() {
   return tab_contents_->web_contents()->GetController().GetActiveEntry()->
       GetPageID();
 }
 
 void ScriptBadgeController::Notify() {
   content::NotificationService::current()->Notify(
       chrome::NOTIFICATION_EXTENSION_LOCATION_BAR_UPDATED,
       content::Source<Profile>(tab_contents_->profile()),
       content::Details<TabContents>(tab_contents_));
 }
 
 void ScriptBadgeController::DidNavigateMainFrame(
     const content::LoadCommittedDetails& details,
     const content::FrameNavigateParams& params) {
   extensions_executing_scripts_.clear();
+  current_actions_.clear();
+}
+
+void ScriptBadgeController::Observe(
+    int type,
+    const content::NotificationSource& source,
+    const content::NotificationDetails& details) {
+  DCHECK_EQ(type, chrome::NOTIFICATION_EXTENSION_UNLOADED);
+  const Extension* extension =
+      content::Details<UnloadedExtensionInfo>(details)->extension;
+  if (EraseExtension(extension))
+    Notify();
 }
 
 bool ScriptBadgeController::OnMessageReceived(const IPC::Message& message) {
   bool handled = true;
   IPC_BEGIN_MESSAGE_MAP(ScriptBadgeController, message)
     IPC_MESSAGE_HANDLER(ExtensionHostMsg_ContentScriptsExecuting,
                         OnContentScriptsExecuting)
     IPC_MESSAGE_UNHANDLED(handled = false)
   IPC_END_MESSAGE_MAP()
   return handled;
 }
 
 void ScriptBadgeController::OnContentScriptsExecuting(
     const std::set<std::string>& extension_ids, int32 page_id) {
-  if (page_id == GetPageID()) {
-    size_t original_size = extensions_executing_scripts_.size();
-    extensions_executing_scripts_.insert(extension_ids.begin(),
-                                         extension_ids.end());
-    if (extensions_executing_scripts_.size() > original_size)
-      Notify();
+  if (page_id != GetPageID())
+    return;
+
+  bool changed = false;
+  for (std::set<std::string>::const_iterator it = extension_ids.begin();
+       it != extension_ids.end(); ++it) {
+    changed |= InsertExtension(*it);
   }
+  if (changed)
+    Notify();
+}
+
+bool ScriptBadgeController::InsertExtension(const std::string& extension_id) {
+  if (!extensions_executing_scripts_.insert(extension_id).second)
+    return false;
+
+  ExtensionService* service = GetExtensionService();
+  if (!service)
+    return  false;
+
+  const Extension* extension = service->extensions()->GetByID(extension_id);
+  if (!extension)
+    return false;
+
+  current_actions_.push_back(extension->GetScriptBadge());
+  return true;
+}
+
+
+bool ScriptBadgeController::EraseExtension(const Extension* extension) {
+  if (extensions_executing_scripts_.erase(extension->id()) == 0)
+    return false;
+
+  size_t size_before = current_actions_.size();
+
+  for (std::vector<ExtensionAction*>::iterator it = current_actions_.begin();
+       it != current_actions_.end(); ++it) {
+    // Safe to -> the extension action because we still have a handle to the
+    // owner Extension.
+    //
+    // Also note that this means that when extensions are uninstalled their
+    // script badges will disappear, even though they're still acting on the
+    // page (since they would have already acted).
+    if ((*it)->extension_id() == extension->id()) {
+      current_actions_.erase(it);
+      break;
+    }
+  }
+
+  CHECK_EQ(size_before, current_actions_.size() + 1);
+  return true;
 }
 
 }  // namespace extensions