1 files changed, 60 insertions, 0 deletions
diff --git a/core/fpdfapi/edit/cpdf_pagecontentmanager.cpp b/core/fpdfapi/edit/cpdf_pagecontentmanager.cpp
index e9ade27bf8..46033bc0d1 100644
--- a/core/fpdfapi/edit/cpdf_pagecontentmanager.cpp
+++ b/core/fpdfapi/edit/cpdf_pagecontentmanager.cpp
@@ -4,6 +4,11 @@
 
 #include "core/fpdfapi/edit/cpdf_pagecontentmanager.h"
 
+#include <map>
+#include <numeric>
+#include <vector>
+
+#include "core/fpdfapi/page/cpdf_pageobject.h"
 #include "core/fpdfapi/parser/cpdf_array.h"
 #include "core/fpdfapi/parser/cpdf_dictionary.h"
 #include "core/fpdfapi/parser/cpdf_document.h"
@@ -85,3 +90,58 @@ size_t CPDF_PageContentManager::AddStream(std::ostringstream* buf) {
   contents_stream_ = new_stream;
   return 0;
 }
+
+void CPDF_PageContentManager::ScheduleRemoveStreamByIndex(size_t stream_index) {
+  streams_to_remove_.insert(stream_index);
+}
+
+void CPDF_PageContentManager::ExecuteScheduledRemovals() {
+  // This method assumes there are no dirty streams in the
+  // CPDF_PageObjectHolder. If there were any, their indexes would need to be
+  // updated.
+  // Since this is only called by CPDF_PageContentGenerator::GenerateContent(),
+  // which cleans up the dirty streams first, this should always be true.
+  ASSERT(obj_holder_->GetDirtyStreams()->empty());
+
+  if (contents_stream_) {
+    // Only stream that can be removed is 0.
+    if (streams_to_remove_.find(0) != streams_to_remove_.end()) {
+      CPDF_Dictionary* page_dict = obj_holder_->GetDict();
+      page_dict->RemoveFor("Contents");
+      contents_stream_ = nullptr;
+    }
+  } else if (contents_array_) {
+    // Initialize a vector with the old stream indexes. This will be used to
+    // build a map from the old to the new indexes.
+    std::vector<size_t> streams_left(contents_array_->GetCount());
+    std::iota(streams_left.begin(), streams_left.end(), 0);
+
+    // In reverse order so as to not change the indexes in the middle of the
+    // loop, remove the streams.
+    for (auto it = streams_to_remove_.rbegin(); it != streams_to_remove_.rend();
+         ++it) {
+      size_t stream_index = *it;
+      contents_array_->RemoveAt(stream_index);
+      streams_left.erase(streams_left.begin() + stream_index);
+    }
+
+    // Create a mapping from the old to the new stream indexes, shifted due to
+    // the deletion of the |streams_to_remove_|.
+    std::map<int32_t, size_t> stream_index_mapping;
+    for (size_t i = 0; i < streams_left.size(); ++i)
+      stream_index_mapping[streams_left[i]] = i;
+
+    // Update the page objects' content stream indexes.
+    for (const auto& obj : *obj_holder_->GetPageObjectList()) {
+      int32_t old_stream_index = obj->GetContentStream();
+      size_t new_stream_index = stream_index_mapping[old_stream_index];
+      obj->SetContentStream(new_stream_index);
+    }
+
+    // Even if there is a single content stream now, keep the array with a
+    // single element. It's valid, a second stream might be added soon, and the
+    // complexity of removing it is not worth it.
+  }
+
+  streams_to_remove_.clear();
+}