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

3 files changed, 195 insertions, 29 deletions

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