From 96187fc8ea0e2bcf3a316c88a309398783c600f7 Mon Sep 17 00:00:00 2001
From: Lei Zhang <[email protected]>
Date: Wed, 25 Mar 2026 11:15:27 -0700
Subject: [PATCH] [PDF Ink Signatures] Fix page unloading prevention for shapes

In PDFiumEngine, `stroked_pages_unload_preventers_` prevents page
unloading if a page has strokes or shapes. However, DiscardStroke() only
checks for strokes before removing `stroked_pages_unload_preventers_`
entries. Add in the missing check for shapes.

To help DiscardStroke() determine which pages have shapes, repurpose the
existing DCHECK-only `pages_with_loaded_v2_ink_paths_` set. Change its
semantics so LoadV2InkPathsForPage() only adds to the set when a page
has shapes. The sanity check that `pages_with_loaded_v2_ink_paths_` was
performing was not very useful anyway.

Bug: 494644471
Change-Id: I6a29720cd5464d1e7cb04fb0ed51e66d4435c32b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7699593
Reviewed-by: Andy Phan <[email protected]>
Commit-Queue: Lei Zhang <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1604965}
---

diff --git a/pdf/pdfium/pdfium_engine.cc b/pdf/pdfium/pdfium_engine.cc
index e993056..1e69739e 100644
--- a/pdf/pdfium/pdfium_engine.cc
+++ b/pdf/pdfium/pdfium_engine.cc
@@ -20,7 +20,6 @@
 #include "base/compiler_specific.h"
 #include "base/containers/flat_map.h"
 #include "base/containers/span.h"
-#include "base/dcheck_is_on.h"
 #include "base/feature_list.h"
 #include "base/functional/bind.h"
 #include "base/i18n/rtl.h"
@@ -5089,11 +5088,12 @@
   }
   ink_stroke_data_.erase(it);
 
-  bool page_still_has_strokes =
+  bool page_still_has_shapes_or_strokes =
+      pages_with_loaded_v2_ink_shapes_.contains(page_index) ||
       std::ranges::any_of(ink_stroke_data_, [page_index](const auto& it) {
         return it.second.page_index == page_index;
       });
-  if (!page_still_has_strokes) {
+  if (!page_still_has_shapes_or_strokes) {
     stroked_pages_unload_preventers_.erase(page_index);
   }
 }
@@ -5123,12 +5123,6 @@
 PDFiumEngine::LoadV2InkPathsForPage(int page_index) {
   CHECK(PageIndexInBounds(page_index));
 
-#if DCHECK_IS_ON()
-  const bool inserted =
-      pages_with_loaded_v2_ink_paths_.insert(page_index).second;
-  CHECK(inserted);
-#endif  // DCHECK_IS_ON()
-
   std::map<InkModeledShapeId, ink::PartitionedMesh> page_shape_map;
 
   PDFiumPage* page = pages_[page_index].get();
@@ -5148,9 +5142,12 @@
   // page unloads and reloads, then the loaded V2 Ink path will no longer match
   // the PDF object, and any updates to the Ink path will not be visible in the
   // PDF.
+  // Also remember the associated page has loaded shapes, so DiscardStroke()
+  // will know not to erase the `stroked_pages_unload_preventers_` entry.
   if (!page_shape_map.empty()) {
     stroked_pages_unload_preventers_.insert(
         {page_index, PDFiumPage::ScopedUnloadPreventer(page)});
+    pages_with_loaded_v2_ink_shapes_.insert(page_index);
   }
 
   return page_shape_map;
diff --git a/pdf/pdfium/pdfium_engine.h b/pdf/pdfium/pdfium_engine.h
index 0d1e554..d88ddfd 100644
--- a/pdf/pdfium/pdfium_engine.h
+++ b/pdf/pdfium/pdfium_engine.h
@@ -18,7 +18,6 @@
 
 #include "base/containers/flat_map.h"
 #include "base/containers/span.h"
-#include "base/dcheck_is_on.h"
 #include "base/functional/callback.h"
 #include "base/memory/raw_ptr.h"
 #include "base/memory/weak_ptr.h"
@@ -1392,11 +1391,11 @@
   // stroke changes.
   std::set<int> ink_stroked_pages_needing_regeneration_;
 
-#if DCHECK_IS_ON()
-  // Used to keep track of LoadV2InkPathsForPage() calls as a sanity check.
-  // Stores the 0-based page indices for pages that have been loaded.
-  std::set<int> pages_with_loaded_v2_ink_paths_;
-#endif  // DCHECK_IS_ON()
+  // Stores the 0-based page indices for pages that have loaded shapes.
+  // Unlike `ink_stroke_data_`, which is dynamic, the loaded shapes data is
+  // static. So just store this data separately from `ink_modeled_shape_map_` to
+  // make searches faster.
+  std::set<int> pages_with_loaded_v2_ink_shapes_;
 
   // Used to hand out unique IDs of type InkModeledShapeId for the V2 Ink paths
   // read out of the PDF. It is stored here as the raw type to simplify
diff --git a/pdf/pdfium/pdfium_engine_unittest.cc b/pdf/pdfium/pdfium_engine_unittest.cc
index 625a072d..71e97c6 100644
--- a/pdf/pdfium/pdfium_engine_unittest.cc
+++ b/pdf/pdfium/pdfium_engine_unittest.cc
@@ -2850,6 +2850,44 @@
             1);
 }
 
+TEST_P(PDFiumEngineInkDrawTest, LoadedV2InkPathsAndApplyAndDiscardStroke) {
+  TestClient client(/*use_skia_renderer=*/GetParam());
+  std::unique_ptr<PDFiumEngine> engine =
+      InitializeEngine(&client, FILE_PATH_LITERAL("ink_v2.pdf"));
+  ASSERT_TRUE(engine);
+  ASSERT_EQ(1, engine->GetNumberOfPages());
+
+  // Check the initial loaded PDF.
+  constexpr int kPageIndex = 0;
+  std::map<InkModeledShapeId, ink::PartitionedMesh> ink_shapes =
+      engine->LoadV2InkPathsForPage(kPageIndex);
+  ASSERT_EQ(1u, ink_shapes.size());
+  ASSERT_EQ(GetPdfMarkObjCountForTesting(engine->doc(),
+                                         kInkAnnotationIdentifierKeyV2),
+            1);
+  ASSERT_TRUE(engine->stroked_pages_unload_preventers_for_testing().contains(
+      kPageIndex));
+
+  // Draw a stroke and immediately discard the stroke to undo.
+  auto brush = std::make_unique<PdfInkBrush>(PdfInkBrush::Type::kPen,
+                                             SK_ColorRED, /*size=*/4.0f);
+  constexpr auto kInputs0 = std::to_array<PdfInkInputData>({
+      {{5.0f, 5.0f}, base::Seconds(0.0f)},
+      {{50.0f, 5.0f}, base::Seconds(0.1f)},
+  });
+  std::optional<ink::StrokeInputBatch> batch = CreateInkInputBatch(kInputs0);
+  ASSERT_TRUE(batch.has_value());
+  ink::Stroke stroke0(brush->ink_brush(), batch.value());
+  constexpr InkStrokeId kStrokeId(1);
+  engine->ApplyStroke(kPageIndex, kStrokeId, stroke0);
+  engine->DiscardStroke(kPageIndex, kStrokeId);
+
+  // The page at `kPageIndex` should still not be allowed to unload, since
+  // `engine` is holding onto page objects within that page.
+  EXPECT_TRUE(engine->stroked_pages_unload_preventers_for_testing().contains(
+      kPageIndex));
+}
+
 TEST_P(PDFiumEngineInkDrawTest, ThumbnailsDoNotContainStrokes) {
   TestClient client(/*use_skia_renderer=*/GetParam());
   std::unique_ptr<PDFiumEngine> engine =

diff --git a/pdf/pdfium/pdfium_engine_unittest.cc b/pdf/pdfium/pdfium_engine_unittest.cc
index 625a072d..71e97c6 100644
--- a/pdf/pdfium/pdfium_engine_unittest.cc
+++ b/pdf/pdfium/pdfium_engine_unittest.cc
@@ -2850,6 +2850,44 @@
             1);
 }
 
+TEST_P(PDFiumEngineInkDrawTest, LoadedV2InkPathsAndApplyAndDiscardStroke) {
+  TestClient client(/*use_skia_renderer=*/GetParam());
+  std::unique_ptr<PDFiumEngine> engine =
+      InitializeEngine(&client, FILE_PATH_LITERAL("ink_v2.pdf"));
+  ASSERT_TRUE(engine);
+  ASSERT_EQ(1, engine->GetNumberOfPages());
+
+  // Check the initial loaded PDF.
+  constexpr int kPageIndex = 0;
+  std::map<InkModeledShapeId, ink::PartitionedMesh> ink_shapes =
+      engine->LoadV2InkPathsForPage(kPageIndex);
+  ASSERT_EQ(1u, ink_shapes.size());
+  ASSERT_EQ(GetPdfMarkObjCountForTesting(engine->doc(),
+                                         kInkAnnotationIdentifierKeyV2),
+            1);
+  ASSERT_TRUE(engine->stroked_pages_unload_preventers_for_testing().contains(
+      kPageIndex));
+
+  // Draw a stroke and immediately discard the stroke to undo.
+  auto brush = std::make_unique<PdfInkBrush>(PdfInkBrush::Type::kPen,
+                                             SK_ColorRED, /*size=*/4.0f);
+  constexpr auto kInputs0 = std::to_array<PdfInkInputData>({
+      {{5.0f, 5.0f}, base::Seconds(0.0f)},
+      {{50.0f, 5.0f}, base::Seconds(0.1f)},
+  });
+  std::optional<ink::StrokeInputBatch> batch = CreateInkInputBatch(kInputs0);
+  ASSERT_TRUE(batch.has_value());
+  ink::Stroke stroke0(brush->ink_brush(), batch.value());
+  constexpr InkStrokeId kStrokeId(1);
+  engine->ApplyStroke(kPageIndex, kStrokeId, stroke0);
+  engine->DiscardStroke(kPageIndex, kStrokeId);
+
+  // The page at `kPageIndex` should still not be allowed to unload, since
+  // `engine` is holding onto page objects within that page.
+  EXPECT_TRUE(engine->stroked_pages_unload_preventers_for_testing().contains(
+      kPageIndex));
+}
+
 TEST_P(PDFiumEngineInkDrawTest, ThumbnailsDoNotContainStrokes) {
   TestClient client(/*use_skia_renderer=*/GetParam());
   std::unique_ptr<PDFiumEngine> engine =