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

components/dom_distiller/core/distilled_content_store.cc

Index: components/dom_distiller/core/distilled_content_store.cc
diff --git a/components/dom_distiller/core/distilled_content_store.cc b/components/dom_distiller/core/distilled_content_store.cc
index 9a92920c75a8d57b60b33b9243b7e3122bf3db6e..11ff967a809905d8a87f9048c8d0c9c344fd5a49 100644
--- a/components/dom_distiller/core/distilled_content_store.cc
+++ b/components/dom_distiller/core/distilled_content_store.cc
@@ -9,10 +9,13 @@
 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(
@@ -34,6 +37,19 @@ void InMemoryContentStore::LoadContent(
 
   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));
@@ -48,6 +64,43 @@ void InMemoryContentStore::LoadContent(
 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