PdfStructureTreeRoot update - allowing to save Object References

openpdf/src/main/java/com/lowagie/text/pdf/PdfStructureTreeRoot.java

@@ -59,9 +59,15 @@
  */
 public class PdfStructureTreeRoot extends PdfDictionary {
 
-    private final Map<Integer, PdfArray> parentTree = new HashMap<>();
+    private final Map<Integer, PdfObject> parentTree = new HashMap<>();
     private final PdfIndirectReference reference;
 
+    /** Next key to be used for adding to the parentTree */
+    private int parentTreeNextKey = 0;
+
+    /** Map which connects [page number] with corresponding [parentTree entry key]  */
+    private final Map<Integer, Integer> pageKeysMap = new HashMap<>();
+
     /**
      * Holds value of property writer.
      */
@@ -111,14 +117,32 @@ public PdfIndirectReference getReference() {
         return this.reference;
     }
 
+    /**
+     * Adds a reference to the existing (already added to the document) object to the parentTree.
+     * This method can be used when the object is need to be referenced via /StructParent key.
+     *
+     * @return key which define the object record key (e.g. in /NUMS array)
+     */
+    public int addExistingObject(PdfIndirectReference reference) {
+        int key = parentTreeNextKey;
+        parentTree.put(key, reference);
+        parentTreeNextKey++;
+        return key;
+    }
+
     void setPageMark(int page, PdfIndirectReference reference) {
-        Integer i = page;
-        PdfArray ar = parentTree.get(i);
-        if (ar == null) {
-            ar = new PdfArray();
-            parentTree.put(i, ar);
+        Integer entryForPageArray = pageKeysMap.get(page);
+        if (entryForPageArray == null) {
+            //putting page array
+            PdfArray ar = new PdfArray();
+            entryForPageArray = parentTreeNextKey;
+            parentTree.put(entryForPageArray, ar);
+            parentTreeNextKey++;
+            pageKeysMap.put(page, entryForPageArray);
         }
-        ar.add(reference);
+
+        PdfArray pageArray = (PdfArray) parentTree.get(entryForPageArray);
+        pageArray.add(reference);
     }
 
     private void nodeProcess(PdfDictionary dictionary, PdfIndirectReference reference)
@@ -128,9 +152,15 @@ private void nodeProcess(PdfDictionary dictionary, PdfIndirectReference referenc
                 .get(0).isNumber()) {
             PdfArray ar = (PdfArray) obj;
             for (int k = 0; k < ar.size(); ++k) {
-                PdfStructureElement e = (PdfStructureElement) ar.getDirectObject(k);
-                ar.set(k, e.getReference());
-                nodeProcess(e, e.getReference());
+                PdfObject pdfObj = ar.getDirectObject(k);
+
+                if (pdfObj instanceof PdfStructureElement) {
+                    PdfStructureElement e = (PdfStructureElement) pdfObj;
+                    ar.set(k, e.getReference());
+                    nodeProcess(e, e.getReference());
+                } else if (pdfObj instanceof PdfIndirectReference) {
+                    ar.set(k, pdfObj);
+                }
             }
         }
         if (reference != null) {
@@ -141,8 +171,13 @@ private void nodeProcess(PdfDictionary dictionary, PdfIndirectReference referenc
     void buildTree() throws IOException {
         Map<Integer, PdfIndirectReference> numTree = new HashMap<>();
         for (Integer i : parentTree.keySet()) {
-            PdfArray ar = parentTree.get(i);
-            numTree.put(i, writer.addToBody(ar).getIndirectReference());
+            PdfObject ar = parentTree.get(i);
+            if (ar instanceof PdfIndirectReference) {
+                //saving the reference to the object which was already added to the body
+                numTree.put(i, (PdfIndirectReference) ar);
+            } else {
+                numTree.put(i, writer.addToBody(ar).getIndirectReference());
+            }
         }
         PdfDictionary dicTree = PdfNumberTree.writeTree(numTree, writer);
         if (dicTree != null) {

openpdf/src/test/java/com/lowagie/text/pdf/PdfStructureTreeRootTest.java

@@ -0,0 +1,91 @@
+package com.lowagie.text.pdf;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import org.junit.jupiter.api.Test;
+
+class PdfStructureTreeRootTest {
+
+    @Test
+    void shouldCreateNewInstanceSuccessfully() {
+        PdfWriter writer = mock(PdfWriter.class);
+        when(writer.getPdfIndirectReference()).thenReturn(mock(PdfIndirectReference.class));
+
+        PdfStructureTreeRoot root = new PdfStructureTreeRoot(writer);
+
+        assertNotNull(root);
+        assertEquals(PdfName.STRUCTTREEROOT, root.get(PdfName.TYPE));
+        assertNotNull(root.getReference());
+        assertSame(writer, root.getWriter());
+    }
+
+    @Test
+    void shouldMapUserTagToStandardTag() {
+        PdfStructureTreeRoot root = new PdfStructureTreeRoot(mock(PdfWriter.class));
+        PdfName userTag = new PdfName("MyTag");
+        PdfName standardTag = PdfName.H1;
+
+        root.mapRole(userTag, standardTag);
+
+        PdfDictionary roleMap = (PdfDictionary) root.get(PdfName.ROLEMAP);
+        assertNotNull(roleMap);
+        assertEquals(standardTag, roleMap.get(userTag));
+    }
+
+    @Test
+    void getWriterShouldReturnCorrectWriter() {
+        PdfWriter mockWriter = mock(PdfWriter.class);
+        PdfStructureTreeRoot treeRoot = new PdfStructureTreeRoot(mockWriter);
+
+        PdfWriter result = treeRoot.getWriter();
+
+        assertSame(mockWriter, result);
+    }
+
+    @Test
+    void addExistingObjectShouldIncreaseParentTreeNextKey() {
+        PdfWriter mockWriter = mock(PdfWriter.class);
+        PdfStructureTreeRoot treeRoot = new PdfStructureTreeRoot(mockWriter);
+        PdfIndirectReference mockRef = mock(PdfIndirectReference.class);
+
+        int firstKey = treeRoot.addExistingObject(mockRef);
+        int secondKey = treeRoot.addExistingObject(mockRef);
+
+        assertNotEquals(firstKey, secondKey);
+        assertEquals(firstKey + 1, secondKey);
+    }
+
+    @Test
+    void buildTreeShouldGenerateParentTreeWithoutException() throws IOException {
+        PdfWriter mockWriter = mock(PdfWriter.class);
+        when(mockWriter.addToBody(any(PdfObject.class))).thenAnswer(invocation -> {
+            PdfObject arg = invocation.getArgument(0);
+            return new PdfIndirectObject(0, arg, mockWriter);
+        });
+
+        PdfStructureTreeRoot treeRoot = new PdfStructureTreeRoot(mockWriter);
+
+        assertDoesNotThrow(() -> treeRoot.buildTree());
+    }
+
+    @Test
+    void buildTreeShouldHandleIOException() throws IOException {
+        PdfWriter mockWriter = mock(PdfWriter.class);
+        doThrow(IOException.class).when(mockWriter).addToBody(any(PdfObject.class));
+
+        PdfStructureTreeRoot treeRoot = new PdfStructureTreeRoot(mockWriter);
+        treeRoot.setPageMark(1, mock(PdfIndirectReference.class));
+
+        assertThrows(IOException.class, () -> treeRoot.buildTree());
+    }
+}