3 files changed, 38 insertions, 12 deletions

diff --git a/webkit/dom_storage/dom_storage_map.cc b/webkit/dom_storage/dom_storage_map.cc
index 3bfbfc7..04ed753 100644
--- a/webkit/dom_storage/dom_storage_map.cc
+++ b/webkit/dom_storage/dom_storage_map.cc
@@ -6,25 +6,32 @@
 
 namespace dom_storage {
 
-DomStorageMap::DomStorageMap() {}
+DomStorageMap::DomStorageMap() {
+  ResetKeyIterator();
+}
+
 DomStorageMap::~DomStorageMap() {}
 
-unsigned DomStorageMap::Length() {
+unsigned DomStorageMap::Length() const {
   return values_.size();
 }
 
 NullableString16 DomStorageMap::Key(unsigned index) {
   if (index >= values_.size())
     return NullableString16(true);
-  ValuesMap::const_iterator it = values_.begin();
-  for (unsigned i = 0; i < index; ++i, ++it) {
-    // TODO(benm): Alter this so we don't have to walk
-    // up to the key from the beginning on each call.
+  while (last_key_index_ != index) {
+    if (last_key_index_ > index) {
+      --key_iterator_;
+      --last_key_index_;
+    } else {
+      ++key_iterator_;
+      ++last_key_index_;
+    }
   }
-  return NullableString16(it->first, false);
+  return NullableString16(key_iterator_->first, false);
 }
 
-NullableString16 DomStorageMap::GetItem(const string16& key) {
+NullableString16 DomStorageMap::GetItem(const string16& key) const {
   ValuesMap::const_iterator found = values_.find(key);
   if (found == values_.end())
     return NullableString16(true);
@@ -41,6 +48,7 @@ bool DomStorageMap::SetItem(
   else
     *old_value = found->second;
   values_[key] = NullableString16(value, false);
+  ResetKeyIterator();
   return true;
 }
 
@@ -52,17 +60,25 @@ bool DomStorageMap::RemoveItem(
     return false;
   *old_value = found->second.string();
   values_.erase(found);
+  ResetKeyIterator();
   return true;
 }
 
 void DomStorageMap::SwapValues(ValuesMap* values) {
   values_.swap(*values);
+  ResetKeyIterator();
 }
 
-DomStorageMap* DomStorageMap::DeepCopy() {
+DomStorageMap* DomStorageMap::DeepCopy() const {
   DomStorageMap* copy = new DomStorageMap;
   copy->values_ = values_;
+  copy->ResetKeyIterator();
   return copy;
 }
 
+void DomStorageMap::ResetKeyIterator() {
+  key_iterator_ = values_.begin();
+  last_key_index_ = 0;
+}
+
 }  // namespace dom_storage
diff --git a/webkit/dom_storage/dom_storage_map.h b/webkit/dom_storage/dom_storage_map.h
index 83724ce..5610824 100644
--- a/webkit/dom_storage/dom_storage_map.h
+++ b/webkit/dom_storage/dom_storage_map.h
@@ -26,21 +26,25 @@ class DomStorageMap
  public:
   DomStorageMap();
 
-  unsigned Length();
+  unsigned Length() const;
   NullableString16 Key(unsigned index);
-  NullableString16 GetItem(const string16& key);
+  NullableString16 GetItem(const string16& key) const;
   bool SetItem(const string16& key, const string16& value,
                NullableString16* old_value);
   bool RemoveItem(const string16& key, string16* old_value);
 
   void SwapValues(ValuesMap* map);
-  DomStorageMap* DeepCopy();
+  DomStorageMap* DeepCopy() const;
 
  private:
   friend class base::RefCountedThreadSafe<DomStorageMap>;
   ~DomStorageMap();
 
+  void ResetKeyIterator();
+
   ValuesMap values_;
+  ValuesMap::const_iterator key_iterator_;
+  unsigned last_key_index_;
   // TODO(benm): track usage and enforce a quota limit.
 };
 
diff --git a/webkit/dom_storage/dom_storage_map_unittest.cc b/webkit/dom_storage/dom_storage_map_unittest.cc
index cb9da2c..81043fd 100644
--- a/webkit/dom_storage/dom_storage_map_unittest.cc
+++ b/webkit/dom_storage/dom_storage_map_unittest.cc
@@ -49,11 +49,17 @@ TEST(DomStorageMapTest, DomStorageMapBasics) {
   EXPECT_TRUE(map->SetItem(kKey2, kValue2, &old_nullable_value));
   EXPECT_EQ(kValue, old_nullable_value.string());
   EXPECT_EQ(2u, map->Length());
+  EXPECT_EQ(kKey, map->Key(0).string());
+  EXPECT_EQ(kKey2, map->Key(1).string());
+  EXPECT_EQ(kKey, map->Key(0).string());
 
   copy = map->DeepCopy();
   EXPECT_EQ(2u, copy->Length());
   EXPECT_EQ(kValue, copy->GetItem(kKey).string());
   EXPECT_EQ(kValue2, copy->GetItem(kKey2).string());
+  EXPECT_EQ(kKey, copy->Key(0).string());
+  EXPECT_EQ(kKey2, copy->Key(1).string());
+  EXPECT_TRUE(copy->Key(2).is_null());
 
   map->SwapValues(&swap);
   EXPECT_EQ(2ul, swap.size());