Rework some UserScriptLoader logic in preparation to removing UserScript copy.

extensions/browser/user_script_loader.cc

 // Copyright 2014 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 "extensions/browser/user_script_loader.h"
 
 #include <stddef.h>
 
 #include <set>
 #include <string>
@@ skipping 10 matching lines @@
 #include "extensions/browser/notification_types.h"
 #include "extensions/common/extension_messages.h"
 
 using content::BrowserThread;
 using content::BrowserContext;
 
 namespace extensions {
 
 namespace {
 
+#if DCHECK_IS_ON()
+bool AreScriptsUnique(const UserScriptList& scripts) {
+  std::set<int> script_ids;
+  for (const UserScript& script : scripts) {
+    if (script_ids.count(script.id()))
+      return false;
+    script_ids.insert(script.id());
+  }
+  return true;
+}
+#endif  // DCHECK_IS_ON()
+
 // Helper function to parse greasesmonkey headers
 bool GetDeclarationValue(const base::StringPiece& line,
                          const base::StringPiece& prefix,
                          std::string* value) {
   base::StringPiece::size_type index = line.find(prefix);
   if (index == base::StringPiece::npos)
     return false;
 
   std::string temp(line.data() + index + prefix.length(),
                    line.length() - index - prefix.length());
@@ skipping 113 matching lines @@
       weak_factory_(this) {
   registrar_.Add(this,
                  content::NOTIFICATION_RENDERER_PROCESS_CREATED,
                  content::NotificationService::AllBrowserContextsAndSources());
 }
 
 UserScriptLoader::~UserScriptLoader() {
   FOR_EACH_OBSERVER(Observer, observers_, OnUserScriptLoaderDestroyed(this));
 }
 
-void UserScriptLoader::AddScripts(const std::set<UserScript>& scripts) {
-  for (std::set<UserScript>::const_iterator it = scripts.begin();
-       it != scripts.end();
-       ++it) {
-    removed_scripts_.erase(*it);
-    added_scripts_.insert(*it);
+void UserScriptLoader::AddScripts(const UserScriptList& scripts) {
+#if DCHECK_IS_ON()
+  // |scripts| with non-unique IDs will work, but that would indicate we are
+  // doing something wrong somewhere, so DCHECK that.
+  DCHECK(AreScriptsUnique(scripts))
+      << "AddScripts() expects scripts with unique IDs.";
+#endif  // DCHECK_IS_ON()
+  for (const UserScript& user_script : scripts) {
+    int id = user_script.id();
+    removed_script_hosts_.erase(UserScriptIDPair(id));
+    if (added_scripts_map_.count(id) == 0)
+      added_scripts_map_[id] = user_script;
   }
   AttemptLoad();
 }
 
-void UserScriptLoader::AddScripts(const std::set<UserScript>& scripts,
+void UserScriptLoader::AddScripts(const UserScriptList& scripts,
                                   int render_process_id,
                                   int render_view_id) {
   AddScripts(scripts);
 }
 
-void UserScriptLoader::RemoveScripts(const std::set<UserScript>& scripts) {
-  for (std::set<UserScript>::const_iterator it = scripts.begin();
-       it != scripts.end();
-       ++it) {
-    added_scripts_.erase(*it);
-    removed_scripts_.insert(*it);
+void UserScriptLoader::RemoveScripts(
+    const std::set<UserScriptIDPair>& scripts) {
+  for (const UserScriptIDPair& id_pair : scripts) {
+    removed_script_hosts_.insert(UserScriptIDPair(id_pair.id, id_pair.host_id));
+    // TODO(lazyboy): We shouldn't be trying to remove scripts that were never
+    // a) added to |added_scripts_map_| or b) being loaded or has done loading
+    // through |user_scripts_|. This would reduce sending redundant IPC.
+    DCHECK(added_scripts_map_.count(id_pair.id) > 0u);

[Devlin, 2016/08/11 01:09:27]

Nit: DCHECK_GT

[lazyboy, 2016/08/11 01:51:06]

Done.

[lazyboy, 2016/08/11 02:00:21]

Actually just realized that this DCHECK is wrong because it is only checking
point a) from the TODO above, but once the script starts loading
|added_scripts_map_| will no longer contain the added script's id: see line 293
in StartLoad().

I've removed this DCHECK.

+    added_scripts_map_.erase(id_pair.id);
   }
   AttemptLoad();
 }
 
 void UserScriptLoader::ClearScripts() {
   clear_scripts_ = true;
-  added_scripts_.clear();
-  removed_scripts_.clear();
+  added_scripts_map_.clear();
+  removed_script_hosts_.clear();
   AttemptLoad();
 }
 
 void UserScriptLoader::Observe(int type,
                                const content::NotificationSource& source,
                                const content::NotificationDetails& details) {
   DCHECK_EQ(type, content::NOTIFICATION_RENDERER_PROCESS_CREATED);
   content::RenderProcessHost* process =
       content::Source<content::RenderProcessHost>(source).ptr();
   if (!ExtensionsBrowserClient::Get()->IsSameContext(
           browser_context_, process->GetBrowserContext()))
     return;
   if (scripts_ready()) {
     SendUpdate(process, shared_memory_.get(),
                std::set<HostID>());  // Include all hosts.
   }
 }
 
 bool UserScriptLoader::ScriptsMayHaveChanged() const {
   // Scripts may have changed if there are scripts added, scripts removed, or
   // if scripts were cleared and either:
   // (1) A load is in progress (which may result in a non-zero number of
   //     scripts that need to be cleared), or
   // (2) The current set of scripts is non-empty (so they need to be cleared).
-  return (added_scripts_.size() ||
-          removed_scripts_.size() ||
-          (clear_scripts_ &&
-           (is_loading() || user_scripts_->size())));
+  return (added_scripts_map_.size() || removed_script_hosts_.size() ||
+          (clear_scripts_ && (is_loading() || user_scripts_->size())));
 }
 
 void UserScriptLoader::AttemptLoad() {
   if (ready_ && ScriptsMayHaveChanged()) {
     if (is_loading())
       pending_load_ = true;
     else
       StartLoad();
   }
 }
 
 void UserScriptLoader::StartLoad() {
   DCHECK_CURRENTLY_ON(BrowserThread::UI);
   DCHECK(!is_loading());
 
   // If scripts were marked for clearing before adding and removing, then clear
   // them.
   if (clear_scripts_) {
     user_scripts_->clear();
   } else {
     for (UserScriptList::iterator it = user_scripts_->begin();
          it != user_scripts_->end();) {
-      if (removed_scripts_.count(*it))
+      UserScriptIDPair id_pair(it->id());
+      if (removed_script_hosts_.count(id_pair) > 0u)
         it = user_scripts_->erase(it);
       else
         ++it;
     }
   }
 
-  user_scripts_->insert(
-      user_scripts_->end(), added_scripts_.begin(), added_scripts_.end());
-
   std::set<int> added_script_ids;
-  for (std::set<UserScript>::const_iterator it = added_scripts_.begin();
-       it != added_scripts_.end();
-       ++it) {
-    added_script_ids.insert(it->id());
-  }
+  for (const auto& id_and_script : added_scripts_map_)
+    added_script_ids.insert(id_and_script.second.id());
 
   // Expand |changed_hosts_| for OnScriptsLoaded, which will use it in
-  // its IPC message. This must be done before we clear |added_scripts_| and
-  // |removed_scripts_| below.
-  std::set<UserScript> changed_scripts(added_scripts_);
-  changed_scripts.insert(removed_scripts_.begin(), removed_scripts_.end());
-  for (const UserScript& script : changed_scripts)
-    changed_hosts_.insert(script.host_id());
+  // its IPC message. This must be done before we clear |added_scripts_map_| and
+  // |removed_script_hosts_| below.
+  for (const auto& id_and_script : added_scripts_map_)

[Devlin, 2016/08/11 01:09:27]

Could we do this and the addition to user scripts at the same time? I.e.

for (pair : added_scripts_map_) {
  changed_hosts_.insert(pair.first);
  added_ids.insert(pair.first);
  user_scripts_.push_back(std::move(pair.second));
}

[Devlin, 2016/08/11 01:09:27]

Could we do this and the addition to user scripts at the same time? I.e.

for (pair : added_scripts_map_) {
  changed_hosts_.insert(pair.first);
  user_scripts_.push_back(std::move(pair.second));
}

[lazyboy, 2016/08/11 01:51:06]

Done.

+    changed_hosts_.insert(id_and_script.second.host_id());
+  for (const UserScriptIDPair& id_pair : removed_script_hosts_)
+    changed_hosts_.insert(id_pair.host_id);
+
+  // Done with |added_scripts_map_|, now put them into |user_scripts_|.
+  user_scripts_->reserve(user_scripts_->size() + added_scripts_map_.size());
+  for (int id : added_script_ids)
+    user_scripts_->push_back(std::move(added_scripts_map_[id]));
 
   LoadScripts(std::move(user_scripts_), changed_hosts_, added_script_ids,
               base::Bind(&UserScriptLoader::OnScriptsLoaded,
                          weak_factory_.GetWeakPtr()));
 
   clear_scripts_ = false;
-  added_scripts_.clear();
-  removed_scripts_.clear();
+  added_scripts_map_.clear();
+  removed_script_hosts_.clear();
   user_scripts_.reset();
 }
 
 // static
 std::unique_ptr<base::SharedMemory> UserScriptLoader::Serialize(
     const UserScriptList& scripts) {
   base::Pickle pickle;
   pickle.WriteUInt32(scripts.size());
   for (const UserScript& script : scripts) {
     // TODO(aa): This can be replaced by sending content script metadata to
@@ skipping 115 matching lines @@
   if (!shared_memory->ShareToProcess(handle, &handle_for_process))
     return;  // This can legitimately fail if the renderer asserts at startup.
 
   if (base::SharedMemory::IsHandleValid(handle_for_process)) {
     process->Send(new ExtensionMsg_UpdateUserScripts(
         handle_for_process, host_id(), changed_hosts, whitelisted_only));
   }
 }
 
 }  // namespace extensions