[dom_distiller] Add support for lookup by URL for InMemoryContentStore.

components/dom_distiller/core/distilled_content_store.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 "components/dom_distiller/core/distilled_content_store.h"
 
 #include "base/message_loop/message_loop.h"
 
 namespace dom_distiller {
 
 InMemoryContentStore::InMemoryContentStore(const int max_num_entries)
-    : cache_(max_num_entries) {
+    : cache_(max_num_entries, CacheDeletor(this)) {
 }
 
 InMemoryContentStore::~InMemoryContentStore() {
+  // Clear the cache before destruction to ensure the CacheDeletor is not called
+  // after InMemoryContentStore has been destroyed.
+  cache_.Clear();
 }
 
 void InMemoryContentStore::SaveContent(
     const ArticleEntry& entry,
     const DistilledArticleProto& proto,
     InMemoryContentStore::SaveCallback callback) {
   InjectContent(entry, proto);
   if (!callback.is_null()) {
     base::MessageLoop::current()->PostTask(FROM_HERE,
                                            base::Bind(callback, true));
   }
 }
 
 void InMemoryContentStore::LoadContent(
     const ArticleEntry& entry,
     InMemoryContentStore::LoadCallback callback) {
   if (callback.is_null())
     return;
 
   ContentMap::const_iterator it = cache_.Get(entry.entry_id());
   bool success = it != cache_.end();
+  if (!success) {
+    // Could not find article by entry ID, so try looking it up by URL.
+    for (int i = 0; i < entry.pages_size(); ++i) {
+      UrlMap::const_iterator url_it = url_to_id_.find(entry.pages(i).url());
+      if (url_it != url_to_id_.end()) {
+        it = cache_.Get(url_it->second);
+        success = it != cache_.end();
+        if (success) {
+          break;
+        }
+      }
+    }
+  }
   scoped_ptr<DistilledArticleProto> distilled_article;
   if (success) {
     distilled_article.reset(new DistilledArticleProto(it->second));
   } else {
     distilled_article.reset(new DistilledArticleProto());
   }
   base::MessageLoop::current()->PostTask(
       FROM_HERE,
       base::Bind(callback, success, base::Passed(&distilled_article)));
 }
 
 void InMemoryContentStore::InjectContent(const ArticleEntry& entry,
                                          const DistilledArticleProto& proto) {
   cache_.Put(entry.entry_id(), proto);
+  AddUrlToIdMapping(entry, proto);
+}
+
+void InMemoryContentStore::AddUrlToIdMapping(
+    const ArticleEntry& entry,
+    const DistilledArticleProto& proto) {
+  for (int i = 0; i < proto.pages_size(); i++) {
+    const DistilledPageProto& page = proto.pages(i);
+    if (page.has_url()) {
+      url_to_id_[page.url()] = entry.entry_id();
+    }
+  }
+}
+
+void InMemoryContentStore::EraseUrlToIdMapping(
+    const DistilledArticleProto& proto) {
+  for (int i = 0; i < proto.pages_size(); i++) {
+    const DistilledPageProto& page = proto.pages(i);
+    if (page.has_url()) {
+      url_to_id_.erase(page.url());
+    }
+  }
+}
+
+InMemoryContentStore::CacheDeletor::CacheDeletor(InMemoryContentStore* store)
+    : store_(store) {
+}
+
+InMemoryContentStore::CacheDeletor::~CacheDeletor() {
+}
+
+void InMemoryContentStore::CacheDeletor::operator()(
+    const DistilledArticleProto& proto) {
+  // When InMemoryContentStore is deleted, the |store_| pointer becomes invalid,
+  // but since the ContentMap is cleared in the InMemoryContentStore destructor,
+  // this should never be called after the destructor.
+  store_->EraseUrlToIdMapping(proto);
 }
 
 }  // namespace dom_distiller