Traverse PDF page tree only once in CPDF_Document Try 3

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 ad84a15a9751a76ef770cd24279d5c3cd83c70d5..efc213ed1e605b0e10f0022d23477fa7229b451a 100644
--- a/core/fpdfapi/parser/cpdf_document.cpp
+++ b/core/fpdfapi/parser/cpdf_document.cpp
@@ -336,6 +336,7 @@ CPDF_Document::CPDF_Document(std::unique_ptr<CPDF_Parser> pParser)
       m_pParser(std::move(pParser)),
       m_pRootDict(nullptr),
       m_pInfoDict(nullptr),
+      m_iNextPageToTraverse(0),
       m_bLinearized(false),
       m_iFirstPageNo(0),
       m_dwFirstPageObjNum(0),
@@ -400,40 +401,71 @@ void CPDF_Document::LoadPages() {
   m_PageList.SetSize(RetrievePageCount());
 }
 
-CPDF_Dictionary* CPDF_Document::FindPDFPage(CPDF_Dictionary* pPages,
-                                            int iPage,
-                                            int nPagesToGo,
-                                            int level) {
+// When this is called, m_pTreeTraversal[level] exists

[Lei Zhang, 2016/11/04 01:38:52]

Move documentation to the header, or make it an ASSERT().

[npm, 2016/11/04 19:33:33]

Done.

+CPDF_Dictionary* CPDF_Document::TraversePDFPages(int iPage,
+                                                 int& nPagesToGo,
+                                                 int level) {
+  if (nPagesToGo < 0)
+    return nullptr;
+  CPDF_Dictionary* pPages = m_pTreeTraversal[level].first;
   CPDF_Array* pKidList = pPages->GetArrayFor("Kids");
-  if (!pKidList)
-    return nPagesToGo == 0 ? pPages : nullptr;
+  if (!pKidList) {
+    if (nPagesToGo != 1)
+      return nullptr;
+    m_PageList.SetAt(iPage, pPages->GetObjNum());
+    return pPages;
+  }
 
-  if (level >= FX_MAX_PAGE_LEVEL)
+  if (level >= FX_MAX_PAGE_LEVEL) {
+    m_pTreeTraversal.pop_back();
     return nullptr;
+  }
 
-  for (size_t i = 0; i < pKidList->GetCount(); i++) {
+  CPDF_Dictionary* page = nullptr;
+  for (size_t i = m_pTreeTraversal[level].second; i < pKidList->GetCount();
+       i++) {
+    if (nPagesToGo == 0)
+      break;
     CPDF_Dictionary* pKid = pKidList->GetDictAt(i);
     if (!pKid) {
       nPagesToGo--;
+      m_pTreeTraversal[level].second++;
       continue;
     }
-    if (pKid == pPages)
+    if (pKid == pPages) {
+      m_pTreeTraversal[level].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--;
+      m_pTreeTraversal[level].second++;
+      if (nPagesToGo == 0) {
+        page = pKid;
+        break;
+      }
     } else {
-      int nPages = pKid->GetIntegerFor("Count");
-      if (nPagesToGo < nPages)
-        return FindPDFPage(pKid, iPage, nPagesToGo, level + 1);
-
-      nPagesToGo -= nPages;
+      // If the vector has size level+1, the child is not in yet
+      if (static_cast<int>(m_pTreeTraversal.size()) == level + 1)

[Lei Zhang, 2016/11/04 01:38:52]

Should |level| be a size_t since it's an index into an array, so you don't need
all these casts?

[npm, 2016/11/04 19:33:33]

Yes, I remember I tried this and failed, maybe forgot to change the other to
size_t too. Done.

+        m_pTreeTraversal.push_back(std::make_pair(pKid, 0));
+      // Now m_pTreeTraversal[level+1] should exist and be equal to pKid.
+      CPDF_Dictionary* pageKid = TraversePDFPages(iPage, nPagesToGo, level + 1);
+      // Check if child was completely processed, i.e. it popped itself out
+      if (static_cast<int>(m_pTreeTraversal.size()) == level + 1)
+        m_pTreeTraversal[level].second++;
+      // If child did not finish or if no pages to go, we are done
+      if (static_cast<int>(m_pTreeTraversal.size()) != level + 1 ||

[npm, 2016/11/03 22:37:57]

Just noticed this 'if' might be causing confusion.  It is sufficient to have
if (nPagesToGo == 0)
because, if the child did not finish, then it must have hit a break, so
nPagesToGo must have reached 0.
I guess I was being extra cautious to make sure to only continue in the
for loop when I'm sure the m_pTreeeTraversal[level].second is increased.
But I can simplify the if, WDYT?

+          nPagesToGo == 0) {
+        page = pageKid;
+        break;
+      }
     }
   }
-  return nullptr;
+  if (m_pTreeTraversal[level].second ==
+      static_cast<int>(pKidList->GetCount())) {
+    m_pTreeTraversal.pop_back();
+  }
+  return page;
 }
 
 CPDF_Dictionary* CPDF_Document::GetPagesDict() const {
@@ -460,17 +492,18 @@ CPDF_Dictionary* CPDF_Document::GetPage(int iPage) {
   if (objnum) {
     if (CPDF_Dictionary* pDict = ToDictionary(GetOrParseIndirectObject(objnum)))
       return pDict;
+    return nullptr;
   }
 
   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());
+  if (m_pTreeTraversal.empty())
+    m_pTreeTraversal.push_back(std::make_pair(pPages, 0));
+  int nPagesToGo = iPage - m_iNextPageToTraverse + 1;
+  CPDF_Dictionary* pPage = TraversePDFPages(iPage, nPagesToGo, 0);
+  m_iNextPageToTraverse = iPage + 1;
   return pPage;
 }
 
@@ -664,6 +697,9 @@ bool CPDF_Document::InsertDeletePDFPage(CPDF_Dictionary* pPages,
       }
       pPages->SetIntegerFor(
           "Count", pPages->GetIntegerFor("Count") + (bInsert ? 1 : -1));
+      // Tree will change, so reset tree transversal variables

[Tom Sepez, 2016/11/03 22:31:00]

nit: maybe a ResetTraversal() helper method that does both of these together,
and lose the comment?

[npm, 2016/11/04 19:33:33]

Done.

+      m_iNextPageToTraverse = 0;
+      m_pTreeTraversal.clear();
       break;
     }
     int nPages = pKid->GetIntegerFor("Count");
@@ -704,6 +740,9 @@ bool CPDF_Document::InsertNewPage(int iPage, CPDF_Dictionary* pPageDict) {
     pPagesList->Add(new CPDF_Reference(this, pPageDict->GetObjNum()));
     pPages->SetIntegerFor("Count", nPages + 1);
     pPageDict->SetReferenceFor("Parent", this, pPages->GetObjNum());
+    // Reset tree transversal variables
+    m_iNextPageToTraverse = 0;
+    m_pTreeTraversal.clear();
   } else {
     std::set<CPDF_Dictionary*> stack = {pPages};
     if (!InsertDeletePDFPage(pPages, iPage, pPageDict, true, &stack))