Index: include/AdblockPlus/ReferrerMapping.h
diff --git a/include/AdblockPlus/ReferrerMapping.h b/include/AdblockPlus/ReferrerMapping.h
index 76d7cbfdfed15ca0b1e022709e15e48dbe481e53..98fd0b5c092d2bec35f1021c1c742b32fe69032e 100644
--- a/include/AdblockPlus/ReferrerMapping.h
+++ b/include/AdblockPlus/ReferrerMapping.h
@@ -22,18 +22,29 @@
 #include <map>
 #include <string>
 #include <vector>
+#include "FilterEngine.h"
 
 namespace AdblockPlus
 {
   /**
    * Stores a mapping between URLs and their referrers.
-   * This can be used to build a chain of referrers for any URL
-   * (see `BuildReferrerChain()`), which approximates the frame structure, see
-   * FilterEngine::Matches().
+   * This can be used to build a frame structure for any URL
+   * (see `BuildFrameStructure()`), it's useful for FilterEngine::Matches().
    */
   class ReferrerMapping
   {
   public:
+    /// The type for URL.
+    typedef std::string Url;
+
+    /// Contains an ordered list of URLs.
+    typedef std::vector<Url> Urls;
+
+    /// Indicates whether the URL is a frame or not.
+    enum FrameIndicator
+    {
+      FRAME_INDICATOR_NOT_FRAME = 0, FRAME_INDICATOR_FRAME = 1
+    };
     /**
      * Constructor.
      * @param maxCachedUrls Number of URL mappings to store. The higher the
@@ -46,21 +57,37 @@ namespace AdblockPlus
      * Records the refferer for a URL.
      * @param url Request URL.
      * @param referrer Request referrer.
+     * @param isFrame Indicates whether the url is a frame.
      */
-    void Add(const std::string& url, const std::string& referrer);
+    void Add(const Url& url, const Url& referrer, FrameIndicator isFrame);
 
     /**
-     * Builds a chain of referrers for the supplied URL.
+     * Builds a frame structure for the supplied URL.
      * This should reconstruct a document's parent frame URLs.
-     * @param url URL to build the chain for.
-     * @return List of URLs, starting with `url`.
+     * @param url URL to build the chain for, it is referrer as well.
+     * @return List of URLs, finishing with `url`.
      */
-    std::vector<std::string> BuildReferrerChain(const std::string& url) const;
+    Urls BuildFrameStructure(const Url& url) const;
 
   private:
     const int maxCachedUrls;
-    std::map<std::string, std::string> mapping;
-    std::list<std::string> cachedUrls;
+    struct RequestInfo
+    {
+      explicit RequestInfo(const Url& referrerArg = Url(),
+        FrameIndicator frameIndicatorArg = FrameIndicator::FRAME_INDICATOR_NOT_FRAME)
+        : referrer(referrerArg), frameIndicator(frameIndicatorArg)
+      {
+      }
+      bool IsFrame() const
+      {
+        return frameIndicator == FrameIndicator::FRAME_INDICATOR_FRAME;
+      }
+      Url referrer;
+      FrameIndicator frameIndicator;
+    };
+    typedef std::map<Url, RequestInfo> ReferrerMap;
+    ReferrerMap mapping;
+    std::list<Url> cachedUrls;
   };
 }
 
Index: src/ReferrerMapping.cpp
diff --git a/src/ReferrerMapping.cpp b/src/ReferrerMapping.cpp
index 7ac6cddff5bf2452670c6d2b986f0802fb5de4d2..0dae687e90bbd7af59ce89c117f6294411e0c937 100644
--- a/src/ReferrerMapping.cpp
+++ b/src/ReferrerMapping.cpp
@@ -24,12 +24,12 @@ ReferrerMapping::ReferrerMapping(const int maxCachedUrls)
 {
 }
 
-void ReferrerMapping::Add(const std::string& url, const std::string& referrer)
+void ReferrerMapping::Add(const Url& url, const Url& referrer, FrameIndicator isFrame)
 {
   if (mapping.find(url) != mapping.end())
     cachedUrls.remove(url);
   cachedUrls.push_back(url);
-  mapping[url] = referrer;
+  mapping[url] = RequestInfo(referrer, isFrame);
 
   const int urlsToPop = cachedUrls.size() - maxCachedUrls;
   for (int i = 0; i < urlsToPop; i++)
@@ -40,21 +40,35 @@ void ReferrerMapping::Add(const std::string& url, const std::string& referrer)
   }
 }
 
-std::vector<std::string> ReferrerMapping::BuildReferrerChain(
-  const std::string& url) const
+ReferrerMapping::Urls ReferrerMapping::BuildFrameStructure(const Url& url) const
 {
-  std::vector<std::string> referrerChain;
-  referrerChain.push_back(url);
+  Urls frames;
   // We need to limit the chain length to ensure we don't block indefinitely
   // if there's a referrer loop.
   const int maxChainLength = 10;
-  std::map<std::string, std::string>::const_iterator currentEntry =
+  ReferrerMap::const_iterator currentEntry = url.empty() ? mapping.end() :
     mapping.find(url);
+  ReferrerMap::const_iterator prevEntry = mapping.end();
   for (int i = 0; i < maxChainLength && currentEntry != mapping.end(); i++)
   {
-    const std::string& currentUrl = currentEntry->second;
-    referrerChain.insert(referrerChain.begin(), currentUrl);
-    currentEntry = mapping.find(currentUrl);
+    if (currentEntry->second.IsFrame())
+    {
+      frames.insert(frames.begin(), currentEntry->first);
+    }
+    prevEntry = currentEntry;
+    currentEntry = mapping.find(currentEntry->second.referrer);
   }
-  return referrerChain;
+  if (prevEntry == mapping.end())
+  {
+    if (!url.empty())
+    {
+      // No data, just assume that the url is a frame.
+      frames.push_back(url);
+    }
+  }
+  else if (!prevEntry->second.referrer.empty())
+  {
+    frames.insert(frames.begin(), prevEntry->second.referrer);
+  }
+  return frames;
 }
Index: test/ReferrerMapping.cpp
diff --git a/test/ReferrerMapping.cpp b/test/ReferrerMapping.cpp
index 63d8e16511aec5b3bf7c2332c6a193a3034d9852..f645aedd351dd2c7fdc100d364d7920cc6c2facd 100644
--- a/test/ReferrerMapping.cpp
+++ b/test/ReferrerMapping.cpp
@@ -18,21 +18,55 @@
 #include <AdblockPlus.h>
 #include <gtest/gtest.h>
 
+typedef AdblockPlus::FilterEngine::ContentType ContentType;
+typedef AdblockPlus::ReferrerMapping::FrameIndicator FrameIndicator;
+
 TEST(ReferrerMappingTest, EmptyReferrerChain)
 {
   AdblockPlus::ReferrerMapping referrerMapping;
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("first");
+    referrerMapping.BuildFrameStructure("first");
   ASSERT_EQ(1, referrerChain.size());
-  ASSERT_EQ("first", referrerChain[0]);
+  EXPECT_EQ("first", referrerChain[0]);
+}
+
+TEST(ReferrerMappingTest, EmptyUrl)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  std::vector<std::string> referrerChain =
+    referrerMapping.BuildFrameStructure("");
+  EXPECT_EQ(0, referrerChain.size());
+
+  referrerMapping.Add("", "referrer", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerChain = referrerMapping.BuildFrameStructure("");
+  EXPECT_EQ(0, referrerChain.size());
+}
+
+TEST(ReferrerMappingTest, EmptyReferrer)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  referrerMapping.Add("first", "", FrameIndicator::FRAME_INDICATOR_FRAME);
+  std::vector<std::string> referrerChain =
+    referrerMapping.BuildFrameStructure("first");
+  ASSERT_EQ(1, referrerChain.size());
+  EXPECT_EQ("first", referrerChain[0]);
+}
+
+TEST(ReferrerMappingTest, NotAFrame)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  referrerMapping.Add("first", "", FrameIndicator::FRAME_INDICATOR_NOT_FRAME);
+  std::vector<std::string> referrerChain =
+    referrerMapping.BuildFrameStructure("first");
+  ASSERT_EQ(0, referrerChain.size());
 }
 
 TEST(ReferrerMappingTest, TwoElementReferrerChain)
 {
   AdblockPlus::ReferrerMapping referrerMapping;
-  referrerMapping.Add("second", "first");
+  referrerMapping.Add("second", "first", FrameIndicator::FRAME_INDICATOR_FRAME);
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("second");
+    referrerMapping.BuildFrameStructure("second");
   ASSERT_EQ(2, referrerChain.size());
   ASSERT_EQ("first", referrerChain[0]);
   ASSERT_EQ("second", referrerChain[1]);
@@ -41,17 +75,17 @@ TEST(ReferrerMappingTest, TwoElementReferrerChain)
 TEST(ReferrerMappingTest, TenElementReferrerChain)
 {
   AdblockPlus::ReferrerMapping referrerMapping;
-  referrerMapping.Add("second", "first");
-  referrerMapping.Add("third", "second");
-  referrerMapping.Add("fourth", "third");
-  referrerMapping.Add("fifth", "fourth");
-  referrerMapping.Add("sixth", "fifth");
-  referrerMapping.Add("seventh", "sixth");
-  referrerMapping.Add("eighth", "seventh");
-  referrerMapping.Add("ninth", "eighth");
-  referrerMapping.Add("tenth", "ninth");
+  referrerMapping.Add("second", "first", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("third", "second", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fourth", "third", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fifth", "fourth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("sixth", "fifth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("seventh", "sixth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("eighth", "seventh", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("ninth", "eighth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("tenth", "ninth", FrameIndicator::FRAME_INDICATOR_FRAME);
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("tenth");
+    referrerMapping.BuildFrameStructure("tenth");
   ASSERT_EQ(10, referrerChain.size());
   ASSERT_EQ("first", referrerChain[0]);
   ASSERT_EQ("second", referrerChain[1]);
@@ -68,14 +102,15 @@ TEST(ReferrerMappingTest, TenElementReferrerChain)
 TEST(ReferrerMappingTest, CacheOnlyFiveUrls)
 {
   AdblockPlus::ReferrerMapping referrerMapping(5);
-  referrerMapping.Add("second", "first");
-  referrerMapping.Add("third", "second");
-  referrerMapping.Add("fourth", "third");
-  referrerMapping.Add("fifth", "fourth");
-  referrerMapping.Add("sixth", "fifth");
-  referrerMapping.Add("seventh", "sixth");
+  std::string outOfCacheUrl = "first";
+  referrerMapping.Add("second", outOfCacheUrl, FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("third", "second", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fourth", "third", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fifth", "fourth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("sixth", "fifth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("seventh", "sixth", FrameIndicator::FRAME_INDICATOR_FRAME);
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("seventh");
+    referrerMapping.BuildFrameStructure("seventh");
   ASSERT_EQ(6, referrerChain.size());
   ASSERT_EQ("second", referrerChain[0]);
   ASSERT_EQ("third", referrerChain[1]);
@@ -84,3 +119,26 @@ TEST(ReferrerMappingTest, CacheOnlyFiveUrls)
   ASSERT_EQ("sixth", referrerChain[4]);
   ASSERT_EQ("seventh", referrerChain[5]);
 }
+
+TEST(ReferrerMappingTest, OnlyFrames)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  referrerMapping.Add("iframe1", "root-page", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("not a frame", "iframe1", FrameIndicator::FRAME_INDICATOR_NOT_FRAME);
+  referrerMapping.Add("iframe2", "not a frame", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("media", "iframe2", FrameIndicator::FRAME_INDICATOR_NOT_FRAME);
+
+  std::vector<std::string> referrerChain = referrerMapping.BuildFrameStructure("media");
+  ASSERT_EQ(3, referrerChain.size());
+  EXPECT_EQ("root-page", referrerChain[0]);
+  EXPECT_EQ("iframe1", referrerChain[1]);
+  EXPECT_EQ("iframe2", referrerChain[2]);
+
+  referrerMapping.Add("iframe3", "media", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerChain = referrerMapping.BuildFrameStructure("iframe3");
+  ASSERT_EQ(4, referrerChain.size());
+  EXPECT_EQ("root-page", referrerChain[0]);
+  EXPECT_EQ("iframe1", referrerChain[1]);
+  EXPECT_EQ("iframe2", referrerChain[2]);
+  EXPECT_EQ("iframe3", referrerChain[3]);
+}
\ No newline at end of file