Traverse PDF page tree only once in CPDF_Document

core/fpdfapi/parser/cpdf_document.cpp

Index: core/fpdfapi/parser/cpdf_document.cpp
diff --git a/core/fpdfapi/parser/cpdf_document.cpp b/core/fpdfapi/parser/cpdf_document.cpp
index c5f64a790ca1d31f6cb4357f5a01b18d6b0585e1..97a430f22aa814fd6336664c427cfc9a11863cd4 100644
--- a/core/fpdfapi/parser/cpdf_document.cpp
+++ b/core/fpdfapi/parser/cpdf_document.cpp
@@ -8,6 +8,7 @@
 
 #include <memory>
 #include <set>
+#include <utility>
 #include <vector>
 
 #include "core/fpdfapi/cpdf_modulemgr.h"
@@ -262,6 +263,9 @@ int InsertDeletePDFPage(CPDF_Document* pDoc,
         }
         pPages->SetIntegerFor(
             "Count", pPages->GetIntegerFor("Count") + (bInsert ? 1 : -1));
+        // Tree will change, so reset tree transversal variables
+        pDoc->m_iLastPageTraversed = -1;
+        pDoc->m_pTreeTraversal = std::stack<std::pair<CPDF_Dictionary*, int>>();
         return 1;
       }
       nPagesToGo--;
@@ -410,6 +414,7 @@ CPDF_Dictionary* CalculateFontDesc(CPDF_Document* pDoc,
 
 CPDF_Document::CPDF_Document(std::unique_ptr<CPDF_Parser> pParser)
     : CPDF_IndirectObjectHolder(),
+      m_iLastPageTraversed(-1),
       m_pParser(std::move(pParser)),
       m_pRootDict(nullptr),
       m_pInfoDict(nullptr),
@@ -477,40 +482,63 @@ void CPDF_Document::LoadPages() {
   m_PageList.SetSize(RetrievePageCount());
 }
 
-CPDF_Dictionary* CPDF_Document::FindPDFPage(CPDF_Dictionary* pPages,
-                                            int iPage,
-                                            int nPagesToGo,
-                                            int level) {
+CPDF_Dictionary* CPDF_Document::TraversePDFPages(int iPage, int nPagesToGo) {
+  std::pair<CPDF_Dictionary*, int>* lastProc = &m_pTreeTraversal.top();
+  CPDF_Dictionary* pPages = lastProc->first;
   CPDF_Array* pKidList = pPages->GetArrayFor("Kids");
-  if (!pKidList)
-    return nPagesToGo == 0 ? pPages : nullptr;
+  if (!pKidList) {
+    m_pTreeTraversal.pop();
+    if (nPagesToGo != 1)

[npm, 2016/10/17 21:10:23]

Before nPagesToGo was the number of pages left minus 1. So compare with 1
instead of 0.

[dsinclair, 2016/10/18 13:35:59]

So, this is check if we're on the page we're looking for?

[npm, 2016/10/18 14:26:31]

The if(!pKidList) will only be entered when the root node is actually page.
Notice that we don't add nodes without Kids into the stack in this method.
If the root has no kids, it has to be a single page.

+      return nullptr;
+    m_PageList.SetAt(iPage, pPages->GetObjNum());
+    return pPages;
+  }
 
-  if (level >= FX_MAX_PAGE_LEVEL)
+  if (m_pTreeTraversal.size() >= FX_MAX_PAGE_LEVEL) {
+    m_pTreeTraversal.pop();
     return nullptr;
+  }
 
-  for (size_t i = 0; i < pKidList->GetCount(); i++) {
+  bool shouldFinish = pPages->GetIntegerFor("Count") <= nPagesToGo;

[npm, 2016/10/17 21:10:23]

This will make sure that the page is popped properly even when the number
of leafs does not match Count.

+  CPDF_Dictionary* page = nullptr;
+  for (size_t i = lastProc->second + 1; i < pKidList->GetCount(); i++) {

[npm, 2016/10/17 21:10:23]

lastProc->second is the index of the last completely processed kid, which
initially is -1, and should be pKidList->GetCount()-1 when done.

     CPDF_Dictionary* pKid = pKidList->GetDictAt(i);
     if (!pKid) {
       nPagesToGo--;
+      lastProc->second++;
       continue;
     }
-    if (pKid == pPages)
+    if (pKid == pPages) {
+      lastProc->second++;
       continue;
+    }
     if (!pKid->KeyExist("Kids")) {
-      if (nPagesToGo == 0)
-        return pKid;
-
-      m_PageList.SetAt(iPage - nPagesToGo, pKid->GetObjNum());
+      m_PageList.SetAt(iPage - nPagesToGo + 1, pKid->GetObjNum());
       nPagesToGo--;
+      lastProc->second++;
+      if (nPagesToGo == 0) {
+        page = pKid;
+        break;
+      }
     } else {
       int nPages = pKid->GetIntegerFor("Count");
-      if (nPagesToGo < nPages)
-        return FindPDFPage(pKid, iPage, nPagesToGo, level + 1);
-
+      m_pTreeTraversal.push(std::make_pair(pKid, -1));
+      CPDF_Dictionary* pageKid = TraversePDFPages(iPage, nPagesToGo);
+      // If the stack top is current element, kid was completely processed.
+      if (lastProc == &m_pTreeTraversal.top())
+        lastProc->second++;
+      if (nPagesToGo <= nPages) {
+        page = pageKid;
+        break;
+      }
       nPagesToGo -= nPages;
     }
   }
-  return nullptr;
+  if (shouldFinish ||
+      lastProc->second == static_cast<int>(pKidList->GetCount() - 1)) {
+    m_pTreeTraversal.pop();
+  }
+  return page;
 }
 
 CPDF_Dictionary* CPDF_Document::GetPagesDict() const {
@@ -534,21 +562,20 @@ CPDF_Dictionary* CPDF_Document::GetPage(int iPage) {
   }
 
   int objnum = m_PageList.GetAt(iPage);
-  if (objnum) {
-    if (CPDF_Dictionary* pDict = ToDictionary(GetOrParseIndirectObject(objnum)))
-      return pDict;
+  if (!objnum || m_iLastPageTraversed < iPage) {
+    CPDF_Dictionary* pPages = GetPagesDict();
+    if (!pPages)
+      return nullptr;
+    if (m_pTreeTraversal.empty())
+      m_pTreeTraversal.push(std::make_pair(pPages, -1));
+    CPDF_Dictionary* page = TraversePDFPages(m_iLastPageTraversed + 1,
+                                             iPage - m_iLastPageTraversed);
+    m_iLastPageTraversed = iPage;
+    return page;
   }
-
-  CPDF_Dictionary* pPages = GetPagesDict();
-  if (!pPages)
-    return nullptr;
-
-  CPDF_Dictionary* pPage = FindPDFPage(pPages, iPage, iPage, 0);
-  if (!pPage)
-    return nullptr;
-
-  m_PageList.SetAt(iPage, pPage->GetObjNum());
-  return pPage;
+  if (CPDF_Dictionary* pDict = ToDictionary(GetOrParseIndirectObject(objnum)))
+    return pDict;
+  return nullptr;
 }
 
 void CPDF_Document::SetPageObjNum(int iPage, uint32_t objNum) {