Return removed value from dictionary.

Change-Id: I218179565cd991b71fd9c909b94f967bbf80c74d
Reviewed-on: https://pdfium-review.googlesource.com/8912
Commit-Queue: Lei Zhang <thestig@chromium.org>
Reviewed-by: Lei Zhang <thestig@chromium.org>

3 files changed, 20 insertions, 3 deletions

diff --git a/core/fpdfapi/parser/cpdf_dictionary.cpp b/core/fpdfapi/parser/cpdf_dictionary.cpp
index cbefb5d56d..a339a8becd 100644
--- a/core/fpdfapi/parser/cpdf_dictionary.cpp
+++ b/core/fpdfapi/parser/cpdf_dictionary.cpp
@@ -197,8 +197,15 @@ void CPDF_Dictionary::ConvertToIndirectObjectFor(
   it->second = pdfium::MakeUnique<CPDF_Reference>(pHolder, pObj->GetObjNum());
 }
 
-void CPDF_Dictionary::RemoveFor(const CFX_ByteString& key) {
-  m_Map.erase(key);
+std::unique_ptr<CPDF_Object> CPDF_Dictionary::RemoveFor(
+    const CFX_ByteString& key) {
+  std::unique_ptr<CPDF_Object> result;
+  auto it = m_Map.find(key);
+  if (it != m_Map.end()) {
+    result = std::move(it->second);
+    m_Map.erase(it);
+  }
+  return result;
 }
 
 void CPDF_Dictionary::ReplaceKey(const CFX_ByteString& oldkey,
diff --git a/core/fpdfapi/parser/cpdf_dictionary.h b/core/fpdfapi/parser/cpdf_dictionary.h
index f741299d2d..31a09d954d 100644
--- a/core/fpdfapi/parser/cpdf_dictionary.h
+++ b/core/fpdfapi/parser/cpdf_dictionary.h
@@ -92,7 +92,7 @@ class CPDF_Dictionary : public CPDF_Object {
                                   CPDF_IndirectObjectHolder* pHolder);
 
   // Invalidates iterators for the element with the key |key|.
-  void RemoveFor(const CFX_ByteString& key);
+  std::unique_ptr<CPDF_Object> RemoveFor(const CFX_ByteString& key);
 
   // Invalidates iterators for the element with the key |oldkey|.
   void ReplaceKey(const CFX_ByteString& oldkey, const CFX_ByteString& newkey);
diff --git a/core/fpdfapi/parser/cpdf_object_unittest.cpp b/core/fpdfapi/parser/cpdf_object_unittest.cpp
index 24a989665f..7d474b2e7d 100644
--- a/core/fpdfapi/parser/cpdf_object_unittest.cpp
+++ b/core/fpdfapi/parser/cpdf_object_unittest.cpp
@@ -883,3 +883,13 @@ TEST(PDFDictionaryTest, ConvertIndirect) {
   EXPECT_EQ(pObj, pNum);
   EXPECT_EQ(42, dict->GetIntegerFor("clams"));
 }
+
+TEST(PDFDictionaryTest, ExtractObjectOnRemove) {
+  auto dict = pdfium::MakeUnique<CPDF_Dictionary>();
+  CPDF_Object* pObj = dict->SetNewFor<CPDF_Number>("child", 42);
+  auto extracted_object = dict->RemoveFor("child");
+  EXPECT_EQ(pObj, extracted_object.get());
+
+  extracted_object = dict->RemoveFor("non_exists_object");
+  EXPECT_FALSE(extracted_object);
+}