commit c4f3b9666c92920174b23c399ec6d8aedc12464b
author Tom Sepez <tsepez@chromium.org> Mon Sep 26 19:32:42 2022 +0000
committer Pdfium LUCI CQ <pdfium-scoped@luci-project-accounts.iam.gserviceaccount.com> Mon Sep 26 19:32:42 2022 +0000
tree e8ef729f02fba3bd3f98e6dd6e542b266d9daaca
parent b7ee8d6c9f27d7e94f12243385ed29062bd17fab

Store retained keys in CPDF_StructTree::StructElementMap.

Change-Id: Ic361535d77c54cf4628036616666f69513b30981
Reviewed-on: https://pdfium-review.googlesource.com/c/pdfium/+/98472
Commit-Queue: Tom Sepez <tsepez@chromium.org>
Reviewed-by: Lei Zhang <thestig@chromium.org>

core/fpdfdoc/cpdf_structelement.cpp

diff --git a/core/fpdfdoc/cpdf_structelement.cpp b/core/fpdfdoc/cpdf_structelement.cpp
index 11e47fe..b03189a 100644
--- a/core/fpdfdoc/cpdf_structelement.cpp
+++ b/core/fpdfdoc/cpdf_structelement.cpp
@@ -40,9 +40,9 @@
 CPDF_StructElement::Kid::~Kid() = default;
 
 CPDF_StructElement::CPDF_StructElement(const CPDF_StructTree* pTree,
-                                       const CPDF_Dictionary* pDict)
+                                       RetainPtr<const CPDF_Dictionary> pDict)
     : m_pTree(pTree),
-      m_pDict(pDict),
+      m_pDict(std::move(pDict)),
       m_Type(GetStructElementType(m_pTree.Get(), m_pDict.Get())) {
   LoadKids(m_pDict.Get());
 }

core/fpdfdoc/cpdf_structelement.h

diff --git a/core/fpdfdoc/cpdf_structelement.h b/core/fpdfdoc/cpdf_structelement.h
index c34f000..7521e01 100644
--- a/core/fpdfdoc/cpdf_structelement.h
+++ b/core/fpdfdoc/cpdf_structelement.h
@@ -57,7 +57,7 @@
   };
 
   CPDF_StructElement(const CPDF_StructTree* pTree,
-                     const CPDF_Dictionary* pDict);
+                     RetainPtr<const CPDF_Dictionary> pDict);
   ~CPDF_StructElement() override;
 
   void LoadKids(const CPDF_Dictionary* pDict);

core/fpdfdoc/cpdf_structtree.cpp

diff --git a/core/fpdfdoc/cpdf_structtree.cpp b/core/fpdfdoc/cpdf_structtree.cpp
index 8650e65..1294ba5 100644
--- a/core/fpdfdoc/cpdf_structtree.cpp
+++ b/core/fpdfdoc/cpdf_structtree.cpp
@@ -6,6 +6,8 @@
 
 #include "core/fpdfdoc/cpdf_structtree.h"
 
+#include <utility>
+
 #include "core/fpdfapi/parser/cpdf_array.h"
 #include "core/fpdfapi/parser/cpdf_dictionary.h"
 #include "core/fpdfapi/parser/cpdf_document.h"
@@ -80,15 +82,13 @@
   StructElementMap element_map;
   for (size_t i = 0; i < pParentArray->size(); i++) {
     RetainPtr<const CPDF_Dictionary> pParent = pParentArray->GetDictAt(i);
-    if (pParent) {
-      // TODO(tsepez): pass moved retained object.
-      AddPageNode(pParent.Get(), &element_map, 0);
-    }
+    if (pParent)
+      AddPageNode(std::move(pParent), &element_map, 0);
   }
 }
 
 RetainPtr<CPDF_StructElement> CPDF_StructTree::AddPageNode(
-    const CPDF_Dictionary* pDict,
+    RetainPtr<const CPDF_Dictionary> pDict,
     StructElementMap* map,
     int nLevel) {
   static constexpr int kStructTreeMaxRecursion = 32;
@@ -103,17 +103,17 @@
   (*map)[pDict] = pElement;
   RetainPtr<const CPDF_Dictionary> pParent = pDict->GetDictFor("P");
   if (!pParent || pParent->GetNameFor("Type") == "StructTreeRoot") {
-    if (!AddTopLevelNode(pDict, pElement))
+    if (!AddTopLevelNode(pDict.Get(), pElement))
       map->erase(pDict);
     return pElement;
   }
 
   RetainPtr<CPDF_StructElement> pParentElement =
-      AddPageNode(pParent.Get(), map, nLevel + 1);
+      AddPageNode(std::move(pParent), map, nLevel + 1);
   if (!pParentElement)
     return pElement;
 
-  if (!pParentElement->UpdateKidIfElement(pDict, pElement.Get()))
+  if (!pParentElement->UpdateKidIfElement(pDict.Get(), pElement.Get()))
     map->erase(pDict);
 
   pElement->SetParent(pParentElement.Get());

core/fpdfdoc/cpdf_structtree.h

diff --git a/core/fpdfdoc/cpdf_structtree.h b/core/fpdfdoc/cpdf_structtree.h
index b0eafba..f99205f 100644
--- a/core/fpdfdoc/cpdf_structtree.h
+++ b/core/fpdfdoc/cpdf_structtree.h
@@ -7,6 +7,7 @@
 #ifndef CORE_FPDFDOC_CPDF_STRUCTTREE_H_
 #define CORE_FPDFDOC_CPDF_STRUCTTREE_H_
 
+#include <functional>
 #include <map>
 #include <memory>
 #include <vector>
@@ -33,13 +34,15 @@
   const CPDF_Dictionary* GetTreeRoot() const { return m_pTreeRoot.Get(); }
 
  private:
-  using StructElementMap =
-      std::map<const CPDF_Dictionary*, RetainPtr<CPDF_StructElement>>;
+  using StructElementMap = std::map<RetainPtr<const CPDF_Dictionary>,
+                                    RetainPtr<CPDF_StructElement>,
+                                    std::less<>>;
 
   void LoadPageTree(const CPDF_Dictionary* pPageDict);
-  RetainPtr<CPDF_StructElement> AddPageNode(const CPDF_Dictionary* pDict,
-                                            StructElementMap* map,
-                                            int nLevel);
+  RetainPtr<CPDF_StructElement> AddPageNode(
+      RetainPtr<const CPDF_Dictionary> pDict,
+      StructElementMap* map,
+      int nLevel);
   bool AddTopLevelNode(const CPDF_Dictionary* pDict,
                        const RetainPtr<CPDF_StructElement>& pElement);