Add the possibility to save added annotations when reorganizing a pdf (bug 2023086)

Author: calixteman

src/core/annotation.js

@@ -354,12 +354,12 @@ class AnnotationFactory {
 
   static async saveNewAnnotations(
     evaluator,
+    xref,
     task,
     annotations,
     imagePromises,
     changes
   ) {
-    const xref = evaluator.xref;
     let baseFontRef;
     const promises = [];
     const { isOffscreenCanvasSupported } = evaluator.options;

src/core/document.js

@@ -151,6 +151,10 @@ class Page {
     });
   }
 
+  createAnnotationEvaluator(handler) {
+    return this.#createPartialEvaluator(handler);
+  }
+
   #getInheritableProperty(key, getArray = false) {
     const value = getInheritableProperty({
       dict: this.pageDict,
@@ -386,6 +390,7 @@ class Page {
     );
     const newData = await AnnotationFactory.saveNewAnnotations(
       partialEvaluator,
+      this.xref,
       task,
       annotations,
       imagePromises,

src/core/editor/pdf_editor.js

@@ -16,16 +16,21 @@
 /** @typedef {import("../document.js").PDFDocument} PDFDocument */
 /** @typedef {import("../document.js").Page} Page */
 /** @typedef {import("../xref.js").XRef} XRef */
+/** @typedef {import("../worker.js").WorkerTask} WorkerTask */
+// eslint-disable-next-line max-len
+/** @typedef {import("../../shared/message_handler.js").MessageHandler} MessageHandler */
 
 import {
   deepCompare,
   getInheritableProperty,
+  getNewAnnotationsMap,
   stringToAsciiOrUTF16BE,
 } from "../core_utils.js";
 import { Dict, isName, Name, Ref, RefSet, RefSetCache } from "../primitives.js";
 import { getModificationDate, stringToPDFString } from "../../shared/util.js";
 import { incrementalUpdate, writeValue } from "../writer.js";
 import { NameTree, NumberTree } from "../name_number_tree.js";
+import { AnnotationFactory } from "../annotation.js";
 import { BaseStream } from "../base_stream.js";
 import { StringStream } from "../stream.js";
 
@@ -75,8 +80,9 @@ class DocumentData {
 }
 
 class XRefWrapper {
-  constructor(entries) {
+  constructor(entries, getNewRef) {
     this.entries = entries;
+    this._getNewRef = getNewRef;
   }
 
   fetch(ref) {
@@ -94,11 +100,17 @@ class XRefWrapper {
   fetchAsync(ref) {
     return Promise.resolve(this.fetch(ref));
   }
+
+  getNewTemporaryRef() {
+    return this._getNewRef();
+  }
 }
 
 class PDFEditor {
   hasSingleFile = false;
 
+  #newAnnotationsParams = null;
+
   currentDocument = null;
 
   oldPages = [];
@@ -107,7 +119,7 @@ class PDFEditor {
 
   xref = [null];
 
-  xrefWrapper = new XRefWrapper(this.xref);
+  xrefWrapper = new XRefWrapper(this.xref, () => this.newRef);
 
   newRefCount = 1;
 
@@ -535,13 +547,33 @@ class PDFEditor {
   /**
    * Extract pages from the given documents.
    * @param {Array<PageInfo>} pageInfos
+   * @param {Object} annotationStorage - The annotation storage containing the
+   *  annotations to be merged into the new document.
+   * @param {MessageHandler} handler - The message handler to use for processing
+   *  the annotations.
+   * @param {WorkerTask} task - The worker task to use for reporting progress
+   *  and cancellation.
    * @return {Promise<void>}
    */
-  async extractPages(pageInfos) {
+  async extractPages(pageInfos, annotationStorage, handler, task) {
     const promises = [];
     let newIndex = 0;
     this.hasSingleFile = pageInfos.length === 1;
     const allDocumentData = [];
+
+    if (annotationStorage) {
+      this.#newAnnotationsParams = {
+        handler,
+        task,
+        newAnnotationsByPage: getNewAnnotationsMap(annotationStorage),
+        imagesPromises: AnnotationFactory.generateImages(
+          annotationStorage.values(),
+          this.xrefWrapper,
+          true
+        ),
+      };
+    }
+
     for (const {
       document,
       includePages,
@@ -1930,16 +1962,44 @@ class PDFEditor {
       await this.#collectDependencies(resources, true, xref)
     );
 
+    let newAnnots = null;
+
     if (annotations) {
       const newAnnotations = await this.#collectDependencies(
         annotations,
         true,
         xref
       );
       this.#fixNamedDestinations(newAnnotations, dedupNamedDestinations);
-      pageDict.setIfArray("Annots", newAnnotations);
+      if (Array.isArray(newAnnotations) && newAnnotations.length > 0) {
+        newAnnots = newAnnotations;
+      }
+    }
+
+    const newAnnotations =
+      this.#newAnnotationsParams?.newAnnotationsByPage.get(pageIndex);
+    if (newAnnotations) {
+      const { handler, task, imagesPromises } = this.#newAnnotationsParams;
+      const changes = new RefSetCache();
+      const newData = await AnnotationFactory.saveNewAnnotations(
+        page.createAnnotationEvaluator(handler),
+        this.xrefWrapper,
+        task,
+        newAnnotations,
+        imagesPromises,
+        changes
+      );
+      for (const [ref, { data }] of changes.items()) {
+        this.xref[ref.num] = data;
+      }
+      newAnnots ||= [];
+      for (const { ref } of newData.annotations) {
+        newAnnots.push(ref);
+      }
     }
 
+    pageDict.setIfArray("Annots", newAnnots);
+
     if (this.useObjectStreams) {
       const newLastRef = this.newRefCount;
       const pageObjectRefs = [];

src/core/worker.js

@@ -551,96 +551,111 @@ class WorkerMessageHandler {
       return pdfManager.ensureDoc("calculationOrderIds");
     });
 
-    handler.on("ExtractPages", async function ({ pageInfos }) {
-      if (!pageInfos) {
-        warn("extractPages: nothing to extract.");
-        return null;
-      }
-      if (!Array.isArray(pageInfos)) {
-        pageInfos = [pageInfos];
-      }
-      let newDocumentId = 0;
-      for (const pageInfo of pageInfos) {
-        if (pageInfo.document === null) {
-          pageInfo.document = pdfManager.pdfDocument;
-        } else if (ArrayBuffer.isView(pageInfo.document)) {
-          const manager = new LocalPdfManager({
-            source: pageInfo.document,
-            docId: `${docId}_extractPages_${newDocumentId++}`,
-            handler,
-            password: pageInfo.password ?? null,
-            evaluatorOptions: Object.assign({}, pdfManager.evaluatorOptions),
-          });
-          let recoveryMode = false;
-          let isValid = true;
-          while (true) {
-            try {
-              await manager.requestLoadedStream();
-              await manager.ensureDoc("checkHeader");
-              await manager.ensureDoc("parseStartXRef");
-              await manager.ensureDoc("parse", [recoveryMode]);
-              break;
-            } catch (e) {
-              if (e instanceof XRefParseException) {
-                if (recoveryMode === false) {
-                  recoveryMode = true;
-                  continue;
+    handler.on(
+      "ExtractPages",
+      async function ({ pageInfos, annotationStorage }) {
+        if (!pageInfos) {
+          warn("extractPages: nothing to extract.");
+          return null;
+        }
+        if (!Array.isArray(pageInfos)) {
+          pageInfos = [pageInfos];
+        }
+        let newDocumentId = 0;
+        for (const pageInfo of pageInfos) {
+          if (pageInfo.document === null) {
+            pageInfo.document = pdfManager.pdfDocument;
+          } else if (ArrayBuffer.isView(pageInfo.document)) {
+            const manager = new LocalPdfManager({
+              source: pageInfo.document,
+              docId: `${docId}_extractPages_${newDocumentId++}`,
+              handler,
+              password: pageInfo.password ?? null,
+              evaluatorOptions: Object.assign({}, pdfManager.evaluatorOptions),
+            });
+            let recoveryMode = false;
+            let isValid = true;
+            while (true) {
+              try {
+                await manager.requestLoadedStream();
+                await manager.ensureDoc("checkHeader");
+                await manager.ensureDoc("parseStartXRef");
+                await manager.ensureDoc("parse", [recoveryMode]);
+                break;
+              } catch (e) {
+                if (e instanceof XRefParseException) {
+                  if (recoveryMode === false) {
+                    recoveryMode = true;
+                    continue;
+                  } else {
+                    isValid = false;
+                    warn("extractPages: XRefParseException.");
+                  }
+                } else if (e instanceof PasswordException) {
+                  const task = new WorkerTask(
+                    `PasswordException: response ${e.code}`
+                  );
+
+                  startWorkerTask(task);
+
+                  try {
+                    const { password } = await handler.sendWithPromise(
+                      "PasswordRequest",
+                      e
+                    );
+                    manager.updatePassword(password);
+                  } catch {
+                    isValid = false;
+                    warn("extractPages: invalid password.");
+                  } finally {
+                    finishWorkerTask(task);
+                  }
                 } else {
                   isValid = false;
-                  warn("extractPages: XRefParseException.");
+                  warn("extractPages: invalid document.");
                 }
-              } else if (e instanceof PasswordException) {
-                const task = new WorkerTask(
-                  `PasswordException: response ${e.code}`
-                );
-
-                startWorkerTask(task);
-
-                try {
-                  const { password } = await handler.sendWithPromise(
-                    "PasswordRequest",
-                    e
-                  );
-                  manager.updatePassword(password);
-                } catch {
-                  isValid = false;
-                  warn("extractPages: invalid password.");
-                } finally {
-                  finishWorkerTask(task);
+                if (!isValid) {
+                  break;
                 }
-              } else {
-                isValid = false;
-                warn("extractPages: invalid document.");
-              }
-              if (!isValid) {
-                break;
               }
             }
-          }
-          if (!isValid) {
-            pageInfo.document = null;
-          }
-          const isPureXfa = await manager.ensureDoc("isPureXfa");
-          if (isPureXfa) {
-            pageInfo.document = null;
-            warn("extractPages does not support pure XFA documents.");
+            if (!isValid) {
+              pageInfo.document = null;
+            }
+            const isPureXfa = await manager.ensureDoc("isPureXfa");
+            if (isPureXfa) {
+              pageInfo.document = null;
+              warn("extractPages does not support pure XFA documents.");
+            } else {
+              pageInfo.document = manager.pdfDocument;
+            }
           } else {
-            pageInfo.document = manager.pdfDocument;
+            warn("extractPages: invalid document.");
+          }
+        }
+        let task;
+        try {
+          const pdfEditor = new PDFEditor();
+          task = new WorkerTask(`ExtractPages: ${pageInfos.length} page(s)`);
+          startWorkerTask(task);
+          const buffer = await pdfEditor.extractPages(
+            pageInfos,
+            annotationStorage,
+            handler,
+            task
+          );
+          return buffer;
+        } catch (reason) {
+          // eslint-disable-next-line no-console
+          console.error(reason);
+          return null;
+        } finally {
+          if (task) {
+            finishWorkerTask(task);
           }
-        } else {
-          warn("extractPages: invalid document.");
         }
       }
-      try {
-        const pdfEditor = new PDFEditor();
-        const buffer = await pdfEditor.extractPages(pageInfos);
-        return buffer;
-      } catch (reason) {
-        // eslint-disable-next-line no-console
-        console.error(reason);
-        return null;
-      }
-    });
+    );
 
     handler.on(
       "SaveDocument",

src/display/api.js

@@ -2973,7 +2973,20 @@ class WorkerTransport {
   }
 
   extractPages(pageInfos) {
-    return this.messageHandler.sendWithPromise("ExtractPages", { pageInfos });
+    const params = {
+      pageInfos,
+    };
+    let transfer;
+    if (this.annotationStorage.size > 0) {
+      const { map, transfer: t } = this.annotationStorage.serializable;
+      params.annotationStorage = map;
+      transfer = t;
+    }
+    return this.messageHandler
+      .sendWithPromise("ExtractPages", params, transfer)
+      .finally(() => {
+        this.annotationStorage.resetModified();
+      });
   }
 
   getPage(pageNumber) {