Refactor IDMap to support safe removing of elements during iteration.

TEST=Covered by unit_tests and other automated tests.
http://crbug.com/19202

Review URL: http://codereview.chromium.org/164518

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@23572 0039d316-1c4b-4281-b951-d872f2087c98

17 files changed, 276 insertions, 119 deletions

diff --git a/chrome/browser/browser.cc b/chrome/browser/browser.cc
index ad61077..905bf11 100644
--- a/chrome/browser/browser.cc
+++ b/chrome/browser/browser.cc
@@ -128,11 +128,12 @@ class BrowserIdleTimer : public base::IdleTimer {
     process.ReduceWorkingSet();
 
     // Handle the Renderer(s).
-    RenderProcessHost::iterator renderer_iter;
-    for (renderer_iter = RenderProcessHost::begin(); renderer_iter !=
-         RenderProcessHost::end(); renderer_iter++) {
-       base::Process process = renderer_iter->second->process();
-       process.ReduceWorkingSet();
+    RenderProcessHost::iterator renderer_iter(
+        RenderProcessHost::AllHostsIterator());
+    while (!renderer_iter.IsAtEnd()) {
+      base::Process process = renderer_iter.GetCurrentValue()->process();
+      process.ReduceWorkingSet();
+      renderer_iter.Advance();
     }
 
     // Handle the child processe.  We need to iterate through them on the IO
diff --git a/chrome/browser/browser_shutdown.cc b/chrome/browser/browser_shutdown.cc
index b67542d..c762cc2 100644
--- a/chrome/browser/browser_shutdown.cc
+++ b/chrome/browser/browser_shutdown.cc
@@ -66,18 +66,19 @@ void OnShutdownStarting(ShutdownType type) {
     // shutdown path for the ones that didn't exit here.
     shutdown_num_processes_slow_ = 0;
     size_t start_rph_size = RenderProcessHost::size();
-    for (RenderProcessHost::iterator hosts = RenderProcessHost::begin();
-         hosts != RenderProcessHost::end();
-         ++hosts) {
-      RenderProcessHost* rph = hosts->second;
-      if (!rph->FastShutdownIfPossible())
+    RenderProcessHost::iterator hosts(RenderProcessHost::AllHostsIterator());
+    while (!hosts.IsAtEnd()) {
+      if (!hosts.GetCurrentValue()->FastShutdownIfPossible()) {
         // TODO(ojan): I think now that we deal with beforeunload/unload
         // higher up, it's not possible to get here. Confirm this and change
         // FastShutdownIfPossible to just be FastShutdown.
         shutdown_num_processes_slow_++;
+      }
       // The number of RPHs should not have changed as the result of invoking
       // FastShutdownIfPossible.
       CHECK(start_rph_size == RenderProcessHost::size());
+
+      hosts.Advance();
     }
   }
 }
diff --git a/chrome/browser/chrome_plugin_browsing_context.cc b/chrome/browser/chrome_plugin_browsing_context.cc
index 71dc80d..504c5c0 100644
--- a/chrome/browser/chrome_plugin_browsing_context.cc
+++ b/chrome/browser/chrome_plugin_browsing_context.cc
@@ -58,16 +58,9 @@ void CPBrowsingContextManager::Observe(NotificationType type,
   URLRequestContext* context = Source<URLRequestContext>(source).ptr();
 
   // Multiple CPBrowsingContexts may refer to the same URLRequestContext.
-  // Remove after collecting all entries, since removing may invalidate
-  // iterators.
-  std::vector<int32> ids_to_remove;
-  for (Map::const_iterator it = map_.begin(); it != map_.end(); ++it) {
-    if (it->second == context)
-      ids_to_remove.push_back(it->first);
-  }
-
-  for (size_t i = 0; i < ids_to_remove.size(); ++i) {
-    map_.Remove(ids_to_remove[i]);
+  for (Map::iterator it(&map_); !it.IsAtEnd(); it.Advance()) {
+    if (it.GetCurrentValue() == context)
+      map_.Remove(it.GetCurrentKey());
   }
 
   reverse_map_.erase(context);
diff --git a/chrome/browser/memory_details.cc b/chrome/browser/memory_details.cc
index ccf279e..df81bc3 100644
--- a/chrome/browser/memory_details.cc
+++ b/chrome/browser/memory_details.cc
@@ -196,13 +196,13 @@ void MemoryDetails::CollectChildInfoOnUIThread() {
     // check if it's a diagnostics-related process.  We skip all diagnostics
     // pages (e.g. "about:xxx" URLs).  Iterate the RenderProcessHosts to find
     // the tab contents.
-    RenderProcessHost::iterator renderer_iter;
-    for (renderer_iter = RenderProcessHost::begin(); renderer_iter !=
-         RenderProcessHost::end(); ++renderer_iter) {
-      DCHECK(renderer_iter->second);
+    RenderProcessHost::iterator renderer_iter(
+        RenderProcessHost::AllHostsIterator());
+    for (; !renderer_iter.IsAtEnd(); renderer_iter.Advance()) {
+      DCHECK(renderer_iter.GetCurrentValue());
       ProcessMemoryInformation& process =
           process_data_[CHROME_BROWSER].processes[index];
-      if (process.pid != renderer_iter->second->process().pid())
+      if (process.pid != renderer_iter.GetCurrentValue()->process().pid())
         continue;
       process.type = ChildProcessInfo::RENDER_PROCESS;
       // The RenderProcessHost may host multiple TabContents.  Any
@@ -212,16 +212,16 @@ void MemoryDetails::CollectChildInfoOnUIThread() {
       // NOTE: This is a bit dangerous.  We know that for now, listeners
       //       are always RenderWidgetHosts.  But in theory, they don't
       //       have to be.
-      RenderProcessHost::listeners_iterator iter;
-      for (iter = renderer_iter->second->listeners_begin();
-           iter != renderer_iter->second->listeners_end(); ++iter) {
-        RenderWidgetHost* widget =
-            static_cast<RenderWidgetHost*>(iter->second);
+      RenderProcessHost::listeners_iterator iter(
+          renderer_iter.GetCurrentValue()->ListenersIterator());
+      for (; !iter.IsAtEnd(); iter.Advance()) {
+        const RenderWidgetHost* widget =
+            static_cast<const RenderWidgetHost*>(iter.GetCurrentValue());
         DCHECK(widget);
         if (!widget || !widget->IsRenderView())
           continue;
 
-        RenderViewHost* host = static_cast<RenderViewHost*>(widget);
+        const RenderViewHost* host = static_cast<const RenderViewHost*>(widget);
         TabContents* contents = NULL;
         if (host->delegate())
           contents = host->delegate()->GetAsTabContents();
diff --git a/chrome/browser/metrics/metrics_service.cc b/chrome/browser/metrics/metrics_service.cc
index 71b162c..bd45fb8 100644
--- a/chrome/browser/metrics/metrics_service.cc
+++ b/chrome/browser/metrics/metrics_service.cc
@@ -957,9 +957,10 @@ void MetricsService::LogTransmissionTimerDone() {
   details->StartFetch();
 
   // Collect WebCore cache information to put into a histogram.
-  for (RenderProcessHost::iterator it = RenderProcessHost::begin();
-       it != RenderProcessHost::end(); ++it) {
-    it->second->Send(new ViewMsg_GetCacheResourceStats());
+  RenderProcessHost::iterator it(RenderProcessHost::AllHostsIterator());
+  while (!it.IsAtEnd()) {
+    it.GetCurrentValue()->Send(new ViewMsg_GetCacheResourceStats());
+    it.Advance();
   }
 }
 
diff --git a/chrome/browser/plugin_service.cc b/chrome/browser/plugin_service.cc
index fe6dfa6..c3e1b14 100644
--- a/chrome/browser/plugin_service.cc
+++ b/chrome/browser/plugin_service.cc
@@ -196,11 +196,11 @@ void PluginService::OnWaitableEventSignaled(base::WaitableEvent* waitable_event)
 
   NPAPI::PluginList::Singleton()->ResetPluginsLoaded();
 
-  for (RenderProcessHost::iterator it = RenderProcessHost::begin();
-       it != RenderProcessHost::end(); ++it) {
-    it->second->Send(new ViewMsg_PurgePluginListCache());
+  for (RenderProcessHost::iterator it(RenderProcessHost::AllHostsIterator());
+       !it.IsAtEnd(); it.Advance()) {
+    it.GetCurrentValue()->Send(new ViewMsg_PurgePluginListCache());
   }
-#endif
+#endif  // defined(OS_WIN)
 }
 
 void PluginService::Observe(NotificationType type,
diff --git a/chrome/browser/renderer_host/browser_render_process_host.cc b/chrome/browser/renderer_host/browser_render_process_host.cc
index 0a7c12b..a0b64c1 100644
--- a/chrome/browser/renderer_host/browser_render_process_host.cc
+++ b/chrome/browser/renderer_host/browser_render_process_host.cc
@@ -649,17 +649,20 @@ bool BrowserRenderProcessHost::FastShutdownIfPossible() {
 
   // Check for any external tab containers, since they may still be running even
   // though this window closed.
-  BrowserRenderProcessHost::listeners_iterator iter;
-  // NOTE: This is a bit dangerous.  We know that for now, listeners are
-  // always RenderWidgetHosts.  But in theory, they don't have to be.
-  for (iter = listeners_begin(); iter != listeners_end(); ++iter) {
-    RenderWidgetHost* widget = static_cast<RenderWidgetHost*>(iter->second);
+  BrowserRenderProcessHost::listeners_iterator iter(ListenersIterator());
+  while (!iter.IsAtEnd()) {
+    // NOTE: This is a bit dangerous.  We know that for now, listeners are
+    // always RenderWidgetHosts.  But in theory, they don't have to be.
+    const RenderWidgetHost* widget =
+        static_cast<const RenderWidgetHost*>(iter.GetCurrentValue());
     DCHECK(widget);
-    if (!widget || !widget->IsRenderView())
-      continue;
-    RenderViewHost* rvh = static_cast<RenderViewHost*>(widget);
-    if (rvh->delegate()->IsExternalTabContainer())
-      return false;
+    if (widget && widget->IsRenderView()) {
+      const RenderViewHost* rvh = static_cast<const RenderViewHost*>(widget);
+      if (rvh->delegate()->IsExternalTabContainer())
+        return false;
+    }
+
+    iter.Advance();
   }
 
   // Otherwise, we're allowed to just terminate the process. Using exit code 0
@@ -876,13 +879,11 @@ void BrowserRenderProcessHost::OnChannelError() {
 
   channel_.reset();
 
-  // This process should detach all the listeners, causing the object to be
-  // deleted. We therefore need a stack copy of the web view list to avoid
-  // crashing when checking for the termination condition the last time.
-  IDMap<IPC::Channel::Listener> local_listeners(listeners_);
-  for (listeners_iterator i = local_listeners.begin();
-       i != local_listeners.end(); ++i) {
-    i->second->OnMessageReceived(ViewHostMsg_RenderViewGone(i->first));
+  IDMap<IPC::Channel::Listener>::iterator iter(&listeners_);
+  while (!iter.IsAtEnd()) {
+    iter.GetCurrentValue()->OnMessageReceived(
+        ViewHostMsg_RenderViewGone(iter.GetCurrentKey()));
+    iter.Advance();
   }
 
   ClearTransportDIBCache();
diff --git a/chrome/browser/renderer_host/render_process_host.cc b/chrome/browser/renderer_host/render_process_host.cc
index 090ed91..4f0cbc0 100644
--- a/chrome/browser/renderer_host/render_process_host.cc
+++ b/chrome/browser/renderer_host/render_process_host.cc
@@ -125,13 +125,8 @@ void RenderProcessHost::UpdateMaxPageID(int32 page_id) {
 }
 
 // static
-RenderProcessHost::iterator RenderProcessHost::begin() {
-  return all_hosts.begin();
-}
-
-// static
-RenderProcessHost::iterator RenderProcessHost::end() {
-  return all_hosts.end();
+RenderProcessHost::iterator RenderProcessHost::AllHostsIterator() {
+  return iterator(&all_hosts);
 }
 
 // static
@@ -165,10 +160,13 @@ RenderProcessHost* RenderProcessHost::GetExistingProcessHost(Profile* profile,
   std::vector<RenderProcessHost*> suitable_renderers;
   suitable_renderers.reserve(size());
 
-  for (iterator iter = begin(); iter != end(); ++iter) {
+  iterator iter(AllHostsIterator());
+  while (!iter.IsAtEnd()) {
     if (run_renderer_in_process() ||
-        IsSuitableHost(iter->second, profile, type))
-      suitable_renderers.push_back(iter->second);
+        IsSuitableHost(iter.GetCurrentValue(), profile, type))
+      suitable_renderers.push_back(iter.GetCurrentValue());
+
+    iter.Advance();
   }
 
   // Now pick a random suitable renderer, if we have any.
diff --git a/chrome/browser/renderer_host/render_process_host.h b/chrome/browser/renderer_host/render_process_host.h
index aee189c..c906235 100644
--- a/chrome/browser/renderer_host/render_process_host.h
+++ b/chrome/browser/renderer_host/render_process_host.h
@@ -28,7 +28,7 @@ struct ViewMsg_ClosePage_Params;
 class RenderProcessHost : public IPC::Channel::Sender,
                           public IPC::Channel::Listener {
  public:
-  typedef IDMap<RenderProcessHost>::const_iterator iterator;
+  typedef IDMap<RenderProcessHost>::iterator iterator;
 
   // We classify renderers according to their highest privilege, and try
   // to group pages into renderers with similar privileges.
@@ -83,11 +83,9 @@ class RenderProcessHost : public IPC::Channel::Sender,
   // Allows iteration over this RenderProcessHost's RenderViewHost listeners.
   // Use from UI thread only.
   typedef IDMap<IPC::Channel::Listener>::const_iterator listeners_iterator;
-  listeners_iterator listeners_begin() {
-    return listeners_.begin();
-  }
-  listeners_iterator listeners_end() {
-    return listeners_.end();
+
+  listeners_iterator ListenersIterator() {
+    return listeners_iterator(&listeners_);
   }
 
   IPC::Channel::Listener* GetListenerByID(int routing_id) {
@@ -195,10 +193,7 @@ class RenderProcessHost : public IPC::Channel::Sender,
 
   // Allows iteration over all the RenderProcessHosts in the browser. Note
   // that each host may not be active, and therefore may have NULL channels.
-  // This is just a standard STL iterator, so it is not valid if the list
-  // of RenderProcessHosts changes between iterations.
-  static iterator begin();
-  static iterator end();
+  static iterator AllHostsIterator();
   static size_t size();  // TODO(brettw) rename this, it's very unclear.
 
   // Returns the RenderProcessHost given its ID.  Returns NULL if the ID does
diff --git a/chrome/browser/renderer_host/render_view_host.cc b/chrome/browser/renderer_host/render_view_host.cc
index 49c4821..8acc403 100644
--- a/chrome/browser/renderer_host/render_view_host.cc
+++ b/chrome/browser/renderer_host/render_view_host.cc
@@ -148,7 +148,7 @@ void RenderViewHost::Observe(NotificationType type,
   if (rph == process()) {
     // Try to get some debugging information on the stack.
     size_t num_hosts = RenderProcessHost::size();
-    bool no_listeners = rph->listeners_begin() == rph->listeners_end();
+    bool no_listeners = rph->ListenersIterator().IsAtEnd();
     bool live_instance = site_instance() != NULL;
     CHECK(live_instance);
     bool live_process = site_instance()->GetProcess() != NULL;
@@ -164,7 +164,7 @@ void RenderViewHost::Observe(NotificationType type,
 bool RenderViewHost::CreateRenderView() {
   DCHECK(!IsRenderViewLive()) << "Creating view twice";
   CHECK(process());
-  CHECK(process()->listeners_begin() != process()->listeners_end()) <<
+  CHECK(!process()->ListenersIterator().IsAtEnd()) <<
       "Our process should have us as a listener.";
 
   // The process may (if we're sharing a process with another host that already
diff --git a/chrome/browser/renderer_host/render_view_host.h b/chrome/browser/renderer_host/render_view_host.h
index 71a15fe..6381ef1 100644
--- a/chrome/browser/renderer_host/render_view_host.h
+++ b/chrome/browser/renderer_host/render_view_host.h
@@ -403,7 +403,7 @@ class RenderViewHost : public RenderWidgetHost,
 
   // RenderWidgetHost public overrides.
   virtual void Shutdown();
-  virtual bool IsRenderView() { return true; }
+  virtual bool IsRenderView() const { return true; }
   virtual void OnMessageReceived(const IPC::Message& msg);
   virtual void GotFocus();
   virtual bool CanBlur() const;
diff --git a/chrome/browser/renderer_host/render_widget_host.h b/chrome/browser/renderer_host/render_widget_host.h
index edace90..3b52e4a 100644
--- a/chrome/browser/renderer_host/render_widget_host.h
+++ b/chrome/browser/renderer_host/render_widget_host.h
@@ -171,7 +171,7 @@ class RenderWidgetHost : public IPC::Channel::Listener,
   virtual void Shutdown();
 
   // Manual RTTI FTW. We are not hosting a web page.
-  virtual bool IsRenderView() { return false; }
+  virtual bool IsRenderView() const { return false; }
 
   // IPC::Channel::Listener
   virtual void OnMessageReceived(const IPC::Message& msg);
diff --git a/chrome/browser/visitedlink_event_listener.cc b/chrome/browser/visitedlink_event_listener.cc
index 6b27929..86fe7d5 100644
--- a/chrome/browser/visitedlink_event_listener.cc
+++ b/chrome/browser/visitedlink_event_listener.cc
@@ -19,20 +19,20 @@ void VisitedLinkEventListener::NewTable(base::SharedMemory* table_memory) {
     return;
 
   // Send to all RenderProcessHosts.
-  for (RenderProcessHost::iterator i = RenderProcessHost::begin();
-       i != RenderProcessHost::end(); ++i) {
-    if (!i->second->HasConnection())
+  for (RenderProcessHost::iterator i(RenderProcessHost::AllHostsIterator());
+       !i.IsAtEnd(); i.Advance()) {
+    if (!i.GetCurrentValue()->HasConnection())
       continue;
 
     base::SharedMemoryHandle new_table;
-    base::ProcessHandle process = i->second->process().handle();
+    base::ProcessHandle process = i.GetCurrentValue()->process().handle();
     if (!process) {
       // process can be null if it's started with the --single-process flag.
       process = base::Process::Current().handle();
     }
 
     table_memory->ShareToProcess(process, &new_table);
-    i->second->Send(new ViewMsg_VisitedLink_NewTable(new_table));
+    i.GetCurrentValue()->Send(new ViewMsg_VisitedLink_NewTable(new_table));
   }
 }
 
@@ -51,17 +51,17 @@ void VisitedLinkEventListener::Reset() {
   coalesce_timer_.Stop();
 
   // Send to all RenderProcessHosts.
-  for (RenderProcessHost::iterator i = RenderProcessHost::begin();
-       i != RenderProcessHost::end(); ++i) {
-    i->second->ResetVisitedLinks();
+  for (RenderProcessHost::iterator i(RenderProcessHost::AllHostsIterator());
+       !i.IsAtEnd(); i.Advance()) {
+    i.GetCurrentValue()->ResetVisitedLinks();
   }
 }
 
 void VisitedLinkEventListener::CommitVisitedLinks() {
   // Send to all RenderProcessHosts.
-  for (RenderProcessHost::iterator i = RenderProcessHost::begin();
-       i != RenderProcessHost::end(); ++i) {
-    i->second->AddVisitedLinks(pending_visited_links_);
+  for (RenderProcessHost::iterator i(RenderProcessHost::AllHostsIterator());
+       !i.IsAtEnd(); i.Advance()) {
+    i.GetCurrentValue()->AddVisitedLinks(pending_visited_links_);
   }
 
   pending_visited_links_.clear();
diff --git a/chrome/chrome.gyp b/chrome/chrome.gyp
index 042c52ae..34d2c93 100644
--- a/chrome/chrome.gyp
+++ b/chrome/chrome.gyp
@@ -3968,6 +3968,7 @@
         'common/extensions/extension_unittest.cc',
         'common/extensions/url_pattern_unittest.cc',
         'common/extensions/user_script_unittest.cc',
+        'common/id_map_unittest.cc',
         'common/important_file_writer_unittest.cc',
         'common/json_value_serializer_unittest.cc',
         'common/mru_cache_unittest.cc',
diff --git a/chrome/common/histogram_synchronizer.cc b/chrome/common/histogram_synchronizer.cc
index 9ed94b0..ab46d47 100644
--- a/chrome/common/histogram_synchronizer.cc
+++ b/chrome/common/histogram_synchronizer.cc
@@ -56,9 +56,10 @@ void HistogramSynchronizer::FetchRendererHistogramsSynchronously(
 
   int sequence_number = GetNextAvaibleSequenceNumber(
       SYNCHRONOUS_HISTOGRAMS, RenderProcessHost::size());
-  for (RenderProcessHost::iterator it = RenderProcessHost::begin();
-       it != RenderProcessHost::end(); ++it) {
-    it->second->Send(new ViewMsg_GetRendererHistograms(sequence_number));
+  for (RenderProcessHost::iterator it(RenderProcessHost::AllHostsIterator());
+       !it.IsAtEnd(); it.Advance()) {
+    it.GetCurrentValue()->Send(
+        new ViewMsg_GetRendererHistograms(sequence_number));
   }
 
   TimeTicks start = TimeTicks::Now();
@@ -111,9 +112,10 @@ void HistogramSynchronizer::FetchRendererHistogramsAsynchronously(
   int sequence_number =
       current_synchronizer->GetNextAvaibleSequenceNumber(
           ASYNC_HISTOGRAMS, RenderProcessHost::size());
-  for (RenderProcessHost::iterator it = RenderProcessHost::begin();
-       it != RenderProcessHost::end(); ++it) {
-    it->second->Send(new ViewMsg_GetRendererHistograms(sequence_number));
+  for (RenderProcessHost::iterator it(RenderProcessHost::AllHostsIterator());
+       !it.IsAtEnd(); it.Advance()) {
+    it.GetCurrentValue()->Send(
+        new ViewMsg_GetRendererHistograms(sequence_number));
   }
 
   // Post a task that would be called after waiting for wait_time.
diff --git a/chrome/common/id_map.h b/chrome/common/id_map.h
index 5d8d40a..bf351b8 100644
--- a/chrome/common/id_map.h
+++ b/chrome/common/id_map.h
@@ -5,6 +5,8 @@
 #ifndef CHROME_COMMON_ID_MAP_H_
 #define CHROME_COMMON_ID_MAP_H_
 
+#include <set>
+
 #include "base/basictypes.h"
 #include "base/hash_tables.h"
 #include "base/logging.h"
@@ -21,31 +23,14 @@ template<class T>
 class IDMap {
  private:
   typedef base::hash_map<int32, T*> HashTable;
-  typedef typename HashTable::iterator iterator;
 
  public:
-  // support const iterators over the items
-  // Note, use iterator->first to get the ID, iterator->second to get the T*
-  typedef typename HashTable::const_iterator const_iterator;
-
-  IDMap() : next_id_(1), check_on_null_data_(false) {
-  }
-  IDMap(const IDMap& other)
-      : next_id_(other.next_id_),
-        data_(other.data_),
-        check_on_null_data_(other.check_on_null_data_) {
+  IDMap() : iteration_depth_(0), next_id_(1), check_on_null_data_(false) {
   }
 
   // Sets whether Add should CHECK if passed in NULL data. Default is false.
   void set_check_on_null_data(bool value) { check_on_null_data_ = value; }
 
-  const_iterator begin() const {
-    return data_.begin();
-  }
-  const_iterator end() const {
-    return data_.end();
-  }
-
   // Adds a view with an automatically generated unique ID. See AddWithID.
   int32 Add(T* data) {
     CHECK(!check_on_null_data_ || data);
@@ -67,12 +52,16 @@ class IDMap {
   }
 
   void Remove(int32 id) {
-    iterator i = data_.find(id);
+    typename HashTable::iterator i = data_.find(id);
     if (i == data_.end()) {
       NOTREACHED() << "Attempting to remove an item not in the list";
       return;
     }
-    data_.erase(i);
+
+    if (iteration_depth_ == 0)
+      data_.erase(i);
+    else
+      removed_ids_.insert(id);
   }
 
   bool IsEmpty() const {
@@ -80,7 +69,7 @@ class IDMap {
   }
 
   T* Lookup(int32 id) const {
-    const_iterator i = data_.find(id);
+    typename HashTable::const_iterator i = data_.find(id);
     if (i == data_.end())
       return NULL;
     return i->second;
@@ -90,15 +79,84 @@ class IDMap {
     return data_.size();
   }
 
- protected:
+  // It is safe to remove elements from the map during iteration. All iterators
+  // will remain valid.
+  template<class ReturnType>
+  class Iterator {
+   public:
+    Iterator(IDMap<T>* map)
+        : map_(map),
+          iter_(map_->data_.begin()) {
+      ++map_->iteration_depth_;
+      SkipRemovedEntries();
+    }
+
+    ~Iterator() {
+      if (--map_->iteration_depth_ == 0)
+        map_->Compact();
+    }
+
+    bool IsAtEnd() const {
+      return iter_ == map_->data_.end();
+    }
+
+    int32 GetCurrentKey() const {
+      return iter_->first;
+    }
+
+    ReturnType* GetCurrentValue() const {
+      return iter_->second;
+    }
+
+    void Advance() {
+      ++iter_;
+      SkipRemovedEntries();
+    }
+
+   private:
+    void SkipRemovedEntries() {
+      while (iter_ != map_->data_.end() &&
+             map_->removed_ids_.find(iter_->first) !=
+             map_->removed_ids_.end()) {
+        ++iter_;
+      }
+    }
+
+    IDMap<T>* map_;
+    typename HashTable::const_iterator iter_;
+  };
+
+  typedef Iterator<T> iterator;
+  typedef Iterator<const T> const_iterator;
+
+ private:
+  void Compact() {
+    DCHECK_EQ(0, iteration_depth_);
+    for (std::set<int32>::const_iterator i = removed_ids_.begin();
+         i != removed_ids_.end(); ++i) {
+      Remove(*i);
+    }
+    removed_ids_.clear();
+  }
+
+  // Keep track of how many iterators are currently iterating on us to safely
+  // handle removing items during iteration.
+  int iteration_depth_;
+
+  // Keep set of IDs that should be removed after the outermost iteration has
+  // finished. This way we manage to not invalidate the iterator when an element
+  // is removed.
+  std::set<int32> removed_ids_;
+
   // The next ID that we will return from Add()
   int32 next_id_;
 
   HashTable data_;
 
- private:
   // See description above setter.
   bool check_on_null_data_;
+
+  DISALLOW_COPY_AND_ASSIGN(IDMap);
 };
 
 #endif  // CHROME_COMMON_ID_MAP_H_
diff --git a/chrome/common/id_map_unittest.cc b/chrome/common/id_map_unittest.cc
new file mode 100644
index 0000000..ed15cbe
--- /dev/null
+++ b/chrome/common/id_map_unittest.cc
@@ -0,0 +1,106 @@
+// Copyright (c) 2009 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/common/id_map.h"
+
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace {
+
+class IDMapTest : public testing::Test {
+};
+
+class TestObject {
+};
+
+TEST_F(IDMapTest, Basic) {
+  IDMap<TestObject> map;
+  EXPECT_TRUE(map.IsEmpty());
+  EXPECT_EQ(0U, map.size());
+
+  TestObject obj1;
+  TestObject obj2;
+
+  int32 id1 = map.Add(&obj1);
+  EXPECT_FALSE(map.IsEmpty());
+  EXPECT_EQ(1U, map.size());
+  EXPECT_EQ(&obj1, map.Lookup(id1));
+
+  int32 id2 = map.Add(&obj2);
+  EXPECT_FALSE(map.IsEmpty());
+  EXPECT_EQ(2U, map.size());
+
+  EXPECT_EQ(&obj1, map.Lookup(id1));
+  EXPECT_EQ(&obj2, map.Lookup(id2));
+
+  map.Remove(id1);
+  EXPECT_FALSE(map.IsEmpty());
+  EXPECT_EQ(1U, map.size());
+
+  map.Remove(id2);
+  EXPECT_TRUE(map.IsEmpty());
+  EXPECT_EQ(0U, map.size());
+
+  map.AddWithID(&obj1, 1);
+  map.AddWithID(&obj2, 2);
+  EXPECT_EQ(&obj1, map.Lookup(1));
+  EXPECT_EQ(&obj2, map.Lookup(2));
+}
+
+TEST_F(IDMapTest, IteratorRemainsValidWhenRemovingCurrentElement) {
+  IDMap<TestObject> map;
+
+  TestObject obj1;
+  TestObject obj2;
+  TestObject obj3;
+
+  map.Add(&obj1);
+  map.Add(&obj2);
+  map.Add(&obj3);
+
+  for (IDMap<TestObject>::const_iterator iter(&map);
+       !iter.IsAtEnd(); iter.Advance()) {
+    map.Remove(iter.GetCurrentKey());
+  }
+}
+
+TEST_F(IDMapTest, IteratorRemainsValidWhenRemovingOtherElements) {
+  IDMap<TestObject> map;
+
+  const int kCount = 5;
+  TestObject obj[kCount];
+  int32 ids[kCount];
+
+  for (int i = 0; i < kCount; i++)
+    ids[i] = map.Add(&obj[i]);
+
+  int counter = 0;
+  for (IDMap<TestObject>::const_iterator iter(&map);
+       !iter.IsAtEnd(); iter.Advance()) {
+    switch (counter) {
+      case 0:
+        EXPECT_EQ(ids[0], iter.GetCurrentKey());
+        EXPECT_EQ(&obj[0], iter.GetCurrentValue());
+        map.Remove(ids[1]);
+        break;
+      case 1:
+        EXPECT_EQ(ids[2], iter.GetCurrentKey());
+        EXPECT_EQ(&obj[2], iter.GetCurrentValue());
+        map.Remove(ids[3]);
+        break;
+      case 2:
+        EXPECT_EQ(ids[4], iter.GetCurrentKey());
+        EXPECT_EQ(&obj[4], iter.GetCurrentValue());
+        map.Remove(ids[0]);
+        break;
+      default:
+        FAIL() << "should not have that many elements";
+        break;
+    }
+
+    counter++;
+  }
+}
+
+}  // namespace