Merge pull request #224 from InusualZ/object-copier

Writer: New API to copy token

src/UglyToad.PdfPig/Writer/Copier/IObjectCopier.cs

@@ -0,0 +1,24 @@
+namespace UglyToad.PdfPig.Writer.Copier
+{
+    using System;
+    using Tokens;
+
+    /// <summary>
+    /// An interface for copying token
+    /// </summary>
+    internal interface IObjectCopier
+    {
+        /// <summary>
+        /// Copy the token to the destination stream
+        /// </summary>
+        /// <param name="sourceToken">Token to copy</param>
+        /// <param name="tokenScanner">Function to resolve indirect reference identified in the token to copy</param>
+        /// <returns></returns>
+        public IToken CopyObject(IToken sourceToken, Func<IndirectReferenceToken, IToken> tokenScanner);
+
+        /// <summary>
+        /// Clear the references of the previously copied object
+        /// </summary>
+        public void ClearReference();
+    }
+}

src/UglyToad.PdfPig/Writer/Copier/MultiCopier.cs

@@ -0,0 +1,75 @@
+namespace UglyToad.PdfPig.Writer.Copier
+{
+    using System;
+    using System.Collections.Generic;
+    using Tokens;
+    using Writer;
+
+    /// <inheritdoc/>
+    internal class MultiCopier : ObjectCopier
+    {
+        private readonly List<IObjectCopier> copiers;
+
+        /// <inheritdoc/>
+        public MultiCopier(PdfStreamWriter destinationStream) : base(destinationStream)
+        {
+            copiers = new List<IObjectCopier>();
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="copier"></param>
+        public void AddCopier(IObjectCopier copier)
+        {
+            copiers.Add(copier);
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="copier"></param>
+        /// <returns></returns>
+        public bool RemoveCopier(IObjectCopier copier)
+        {
+            return copiers.Remove(copier);
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <returns></returns>
+        public IReadOnlyList<IObjectCopier> GetCopiers()
+        {
+            return copiers;
+        }
+
+        /// <inheritdoc/>
+        public override IToken CopyObject(IToken sourceToken, Func<IndirectReferenceToken, IToken> tokenScanner)
+        {
+            // We give the token to the child copiers, to see if they have a better way of copying the token
+            foreach (var copier in copiers)
+            {
+                var newToken = copier.CopyObject(sourceToken, tokenScanner);
+                if (newToken != null)
+                {
+                    return newToken;
+                }
+            }
+
+            // If the token did not found a suitable copier, let just do a simple copy of the token
+            return base.CopyObject(sourceToken, tokenScanner);
+        }
+
+        /// <inheritdoc/>
+        public override void ClearReference()
+        {
+            foreach (var copier in copiers)
+            {
+                copier.ClearReference();
+            }
+
+            base.ClearReference();
+        }
+    }
+}

src/UglyToad.PdfPig/Writer/Copier/ObjectCopier.cs

@@ -0,0 +1,167 @@
+namespace UglyToad.PdfPig.Writer.Copier
+{
+    using System;
+    using System.Collections.Generic;
+    using PdfPig;
+    using Tokenization.Scanner;
+    using Tokens;
+    using Writer;
+
+    /// <inheritdoc/>
+    internal class ObjectCopier : IObjectCopier
+    {
+        private readonly PdfStreamWriter pdfStream;
+
+        private readonly Dictionary<IndirectReferenceToken, IndirectReferenceToken> newReferenceMap;
+
+        /// <inheritdoc/>
+        public ObjectCopier(PdfStreamWriter destinationStream)
+        {
+            pdfStream = destinationStream ?? throw new ArgumentNullException(nameof(destinationStream));
+            newReferenceMap = new Dictionary<IndirectReferenceToken, IndirectReferenceToken>();
+        }
+
+        /// <inheritdoc/>
+        public IToken CopyObject(IToken sourceToken, PdfDocument sourceDocument)
+        {
+            IToken tokenScanner(IndirectReferenceToken referenceToken)
+            {
+                var objToken = sourceDocument.Structure.GetObject(referenceToken.Data);
+                return objToken.Data;
+            }
+
+            return CopyObject(sourceToken, tokenScanner);
+        }
+
+        /// <inheritdoc/>
+        public IToken CopyObject(IToken sourceToken, IPdfTokenScanner tokenScanner)
+        {
+            IToken tokenGetter(IndirectReferenceToken referenceToken)
+            {
+                var objToken = tokenScanner.Get(referenceToken.Data);
+                return objToken.Data;
+            }
+
+            return CopyObject(sourceToken, tokenGetter);
+        }
+
+        /// <inheritdoc/>
+        public virtual IToken CopyObject(IToken sourceToken, Func<IndirectReferenceToken, IToken> tokenScanner)
+        {
+            // This token need to be deep copied, because they could contain reference. So we have to update them.
+            switch (sourceToken)
+            {
+                case DictionaryToken dictionaryToken:
+                    {
+                        var newContent = new Dictionary<NameToken, IToken>();
+                        foreach (var setPair in dictionaryToken.Data)
+                        {
+                            var name = setPair.Key;
+                            var token = setPair.Value;
+
+                            newContent.Add(NameToken.Create(name), CopyObject(token, tokenScanner));
+                        }
+
+                        return new DictionaryToken(newContent);
+                    }
+                case ArrayToken arrayToken:
+                    {
+                        var newArray = new List<IToken>(arrayToken.Length);
+                        foreach (var token in arrayToken.Data)
+                        {
+                            newArray.Add(CopyObject(token, tokenScanner));
+                        }
+
+                        return new ArrayToken(newArray);
+                    }
+                case IndirectReferenceToken referenceToken:
+                    {
+                        if (TryGetNewReference(referenceToken, out var newReferenceToken))
+                        {
+                            return newReferenceToken;
+                        }
+
+                        var referencedToken = tokenScanner(referenceToken);
+                        var newReferencedToken = CopyObject(referencedToken, tokenScanner);
+
+                        var newToken = WriteToken(newReferencedToken);
+                        SetNewReference(referenceToken, newToken);
+                        return newToken;
+                    }
+
+                case StreamToken streamToken:
+                    {
+                        var properties = CopyObject(streamToken.StreamDictionary, tokenScanner);
+                        var bytes = streamToken.Data;
+                        return new StreamToken(properties as DictionaryToken, bytes);
+                    }
+
+                case ObjectToken _:
+                    {
+                        throw new NotSupportedException("Copying a Object Token is not supported");
+                    }
+            }
+
+            return sourceToken;
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="sourceReferenceToken"></param>
+        /// <param name="newReferenceToken"></param>
+        /// <returns></returns>
+        public virtual bool TryGetNewReference(IndirectReferenceToken sourceReferenceToken, out IndirectReferenceToken newReferenceToken)
+        {
+            newReferenceToken = default;
+            foreach (var referenceSet in newReferenceMap)
+            {
+                if (!referenceSet.Key.Equals(sourceReferenceToken))
+                {
+                    continue;
+                }
+
+                newReferenceToken = referenceSet.Value;
+                return true;
+            }
+
+            return false;
+        }
+
+        /// <inheritdoc/>
+        public virtual void ClearReference()
+        {
+            newReferenceMap.Clear();
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="oldToken"></param>
+        /// <param name="newToken"></param>
+        public void SetNewReference(IndirectReferenceToken oldToken, IndirectReferenceToken newToken)
+        {
+            newReferenceMap.Add(oldToken, newToken);
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <returns></returns>
+        public int ReserveTokenNumber()
+        {
+            return pdfStream.ReserveNumber();
+        }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="token"></param>
+        /// <param name="reservedNumber"></param>
+        /// <returns></returns>
+        public IndirectReferenceToken WriteToken(IToken token, int? reservedNumber = null)
+        {
+            return pdfStream.WriteToken(token, reservedNumber);
+        }
+    }
+}

src/UglyToad.PdfPig/Writer/Copier/Page/PagesCopier.cs

@@ -0,0 +1,96 @@
+namespace UglyToad.PdfPig.Writer.Copier.Page
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using Core;
+    using Tokens;
+
+    /// <inheritdoc/>
+    internal class PagesCopier : IObjectCopier
+    {
+        private readonly ObjectCopier copier;
+
+        private readonly IndirectReferenceToken rootPagesReferenceToken;
+
+        /// <inheritdoc/>
+        public PagesCopier(ObjectCopier mainCopier, IndirectReferenceToken rootPagesToken = null)
+        {
+            copier = mainCopier;
+            rootPagesReferenceToken = rootPagesToken;
+        }
+
+        /// <inheritdoc/>
+        public IToken CopyObject(IToken sourceToken, Func<IndirectReferenceToken, IToken> tokenScanner)
+        {
+            if (!(sourceToken is IndirectReferenceToken sourceReferenceToken))
+            {
+                return null;
+            }
+
+            // Check if this token haven't been copied before
+            if (copier.TryGetNewReference(sourceReferenceToken, out var newReferenceToken))
+            {
+                return newReferenceToken;
+            }
+
+            // Make sure that we are copying a DictionaryToken
+            var token = tokenScanner(sourceReferenceToken);
+            if (!(token is DictionaryToken dictionaryToken))
+            {
+                return null;
+            }
+
+            // Make sure we are copying a `/Pages` Dictionary
+            if (!dictionaryToken.TryGet(NameToken.Type, out var nameTypeToken) || !nameTypeToken.Equals(NameToken.Pages))
+            {
+                return null;
+            }
+
+            // We have to reserve the reference before hand, because if we don't, we would fall in a loop.
+            // The child `/Page` have a reference to the parent
+            var tokenNumber = copier.ReserveTokenNumber();
+            copier.SetNewReference(sourceReferenceToken, new IndirectReferenceToken(new IndirectReference(tokenNumber, 0)));
+
+            // If `/Pages` is not the root page node, copy the token normally
+            // We are testing for one:
+            //  * If @rootPagesReferenceToken is null, just do a normal copy of the tree
+            //  * If the tree have a Parent NameToken, it means the tree is not a root tree so we don't have to assign him
+            //    a new parent
+            if (rootPagesReferenceToken == null || dictionaryToken.TryGet(NameToken.Parent, out IndirectReferenceToken _))
+            {
+                return copier.WriteToken(copier.CopyObject(dictionaryToken, tokenScanner), tokenNumber);
+            }
+
+            // Since the tree is a root tree, it means that the tree comes from another document, we have to make sure
+            // that the new tree is a child of the new root tree, this we do by adding a Parent NameToken to the tree,
+            // that point to @rootPagesReferenceToken
+            return CopyPagesTree(dictionaryToken, tokenNumber, tokenScanner);
+        }
+
+        private IndirectReferenceToken CopyPagesTree(DictionaryToken pagesDictionary, int reservedNumber, Func<IndirectReferenceToken, IToken> tokenScanner)
+        {
+            Debug.Assert(rootPagesReferenceToken != null);
+
+            var newContent = new Dictionary<NameToken, IToken>()
+            {
+                {NameToken.Parent, rootPagesReferenceToken}
+            };
+
+            foreach (var dataSet in pagesDictionary.Data)
+            {
+                newContent.Add(NameToken.Create(dataSet.Key), copier.CopyObject(dataSet.Value, tokenScanner));
+            }
+
+            var newPagesTree = new DictionaryToken(newContent);
+
+            return copier.WriteToken(newPagesTree, reservedNumber);
+        }
+
+        /// <inheritdoc/>
+        public void ClearReference()
+        {
+            // Nothing to do
+        }
+    }
+}

src/UglyToad.PdfPig/Writer/Copier/TokenHelper.cs

@@ -0,0 +1,34 @@
+namespace UglyToad.PdfPig.Writer.Copier
+{
+    using System;
+    using Tokens;
+
+    internal static class TokenHelper
+    {
+        // This is to avoid infinite loop in production. Although, it should never happen
+        const int MAX_ITERATIONS = 10;
+
+        public static T GetTokenAs<T>(IToken token, Func<IndirectReferenceToken, IToken> lookupFunc) where T : IToken
+        {
+            var iterations = 0;
+
+            var original = token;
+            while (iterations++ < MAX_ITERATIONS)
+            {
+                switch (token)
+                {
+                    case T result:
+                        return result;
+                    case IndirectReferenceToken tokenReference:
+                        token = lookupFunc(tokenReference);
+                        continue;
+                    case ObjectToken tokenObject:
+                        token = tokenObject.Data;
+                        continue;
+                }
+            }
+
+            throw new InvalidOperationException($"Unable to extract a {typeof(T)} token from {original}");
+        }
+    }
+}

src/UglyToad.PdfPig/Writer/PdfMerger.cs

@@ -2,9 +2,10 @@
 {
     using System;
     using System.Collections.Generic;
-    using System.Diagnostics;
     using System.IO;
     using Content;
+    using Copier;
+    using Copier.Page;
     using Core;
     using CrossReference;
     using Encryption;
@@ -149,31 +150,36 @@
         private class DocumentMerger
         {
             private const decimal DefaultVersion = 1.2m;
-            
-            private readonly PdfStreamWriter context = new PdfStreamWriter();
-            private readonly List<IndirectReferenceToken> pagesTokenReferences = new List<IndirectReferenceToken>();
+
+            private readonly PdfStreamWriter context;
+            private readonly List<IndirectReferenceToken> pagesTokenReferences;
             private readonly IndirectReferenceToken rootPagesReference;
+            private readonly MultiCopier copier;
 
             private decimal currentVersion = DefaultVersion;
             private int pageCount = 0;
-            
-            private readonly Dictionary<IndirectReferenceToken, IndirectReferenceToken> referencesFromDocument =
-                new Dictionary<IndirectReferenceToken, IndirectReferenceToken>();
 
             public DocumentMerger()
             {
+                context = new PdfStreamWriter();
+                pagesTokenReferences = new List<IndirectReferenceToken>();
+
                 rootPagesReference = context.ReserveNumberToken();
+
+                copier = new MultiCopier(context);
+                copier.AddCopier(new PagesCopier(copier, rootPagesReference));
             }
- 
+
             public void AppendDocument(Catalog documentCatalog, decimal version, IPdfTokenScanner tokenScanner)
             {
                 currentVersion = Math.Max(version, currentVersion);
 
-                var (pagesReference, count) = CopyPagesTree(documentCatalog.PageTree, rootPagesReference, tokenScanner);
-                pageCount += count;
-                pagesTokenReferences.Add(pagesReference);
+                var copiedPages = copier.CopyObject(new IndirectReferenceToken(documentCatalog.PageTree.Reference), tokenScanner) as IndirectReferenceToken;
+                pagesTokenReferences.Add(copiedPages);
 
-                referencesFromDocument.Clear();
+                pageCount += documentCatalog.PagesDictionary.Get<NumericToken>(NameToken.Count, tokenScanner).Int;
+
+                copier.ClearReference();
             }
 
             public byte[] Build()
@@ -190,7 +196,7 @@
                     { NameToken.Count, new NumericToken(pageCount) }
                 });
 
-                var pagesRef = context.WriteToken( pagesDictionary, (int)rootPagesReference.Data.ObjectNumber);
+                var pagesRef = context.WriteToken(pagesDictionary, (int)rootPagesReference.Data.ObjectNumber);
 
                 var catalog = new DictionaryToken(new Dictionary<NameToken, IToken>
                 {
@@ -199,9 +205,9 @@
                 });
 
                 var catalogRef = context.WriteToken(catalog);
-                
+
                 context.Flush(currentVersion, catalogRef);
-                
+
                 var bytes = context.ToArray();
 
                 Close();
@@ -209,165 +215,10 @@
                 return bytes;
             }
 
-            public void Close()
+            private void Close()
             {
                 context.Dispose();
             }
-
-            private (IndirectReferenceToken, int) CopyPagesTree(PageTreeNode treeNode, IndirectReferenceToken treeParentReference, IPdfTokenScanner tokenScanner)
-            {
-                Debug.Assert(!treeNode.IsPage);
-                
-                var currentNodeReference = context.ReserveNumberToken();
-
-                var pageReferences = new List<IndirectReferenceToken>();
-                var nodeCount = 0;
-                foreach (var pageNode in treeNode.Children)
-                {
-                    IndirectReferenceToken newEntry;
-                    if (!pageNode.IsPage)
-                    {
-                        var count = 0;
-                        (newEntry, count) = CopyPagesTree(pageNode, currentNodeReference, tokenScanner);
-                        nodeCount += count;
-                    }
-                    else
-                    {
-                        newEntry = CopyPageNode(pageNode, currentNodeReference, tokenScanner);
-                        ++nodeCount;
-                    }
-                    
-                    pageReferences.Add(newEntry);
-                }
-
-                var newPagesNode = new Dictionary<NameToken, IToken>
-                {
-                    { NameToken.Type, NameToken.Pages },
-                    { NameToken.Kids, new ArrayToken(pageReferences) },
-                    { NameToken.Count, new NumericToken(nodeCount) },
-                    { NameToken.Parent, treeParentReference }
-                };
-
-                foreach (var pair in treeNode.NodeDictionary.Data)
-                {
-                    if (IgnoreKeyForPagesNode(pair))
-                    {
-                        continue;
-                    }
-
-                    newPagesNode[NameToken.Create(pair.Key)] = CopyToken(pair.Value, tokenScanner);
-                }
-
-                var pagesDictionary = new DictionaryToken(newPagesNode);
-
-                return (context.WriteToken(pagesDictionary, (int)currentNodeReference.Data.ObjectNumber), nodeCount);
-            }
-
-            private IndirectReferenceToken CopyPageNode(PageTreeNode pageNode, IndirectReferenceToken parentPagesObject, IPdfTokenScanner tokenScanner)
-            {
-                Debug.Assert(pageNode.IsPage);
-
-                var pageDictionary = new Dictionary<NameToken, IToken>
-                {
-                    {NameToken.Parent, parentPagesObject},
-                };
-
-                foreach (var setPair in pageNode.NodeDictionary.Data)
-                {
-                    var name = setPair.Key;
-                    var token = setPair.Value;
-
-                    if (name == NameToken.Parent)
-                    {
-                        // Skip Parent token, since we have to reassign it
-                        continue;
-                    }
-
-                    pageDictionary.Add(NameToken.Create(name), CopyToken(token, tokenScanner));
-                }
-
-                return context.WriteToken(new DictionaryToken(pageDictionary));
-            }
-            
-            /// <summary>
-            /// The purpose of this method is to resolve indirect reference. That mean copy the reference's content to the new document's stream
-            /// and replace the indirect reference with the correct/new one
-            /// </summary>
-            /// <param name="tokenToCopy">Token to inspect for reference</param>
-            /// <param name="tokenScanner">scanner get the content from the original document</param>
-            /// <returns>A reference of the token that was copied. With all the reference updated</returns>
-            private IToken CopyToken(IToken tokenToCopy, IPdfTokenScanner tokenScanner)
-            {
-                // This token need to be deep copied, because they could contain reference. So we have to update them.
-                switch (tokenToCopy)
-                {
-                    case DictionaryToken dictionaryToken:
-                    {
-                        var newContent = new Dictionary<NameToken, IToken>();
-                        foreach (var setPair in dictionaryToken.Data)
-                        {
-                            var name = setPair.Key;
-                            var token = setPair.Value;
-                            newContent.Add(NameToken.Create(name), CopyToken(token, tokenScanner));
-                        }
-
-                        return new DictionaryToken(newContent);
-                    }
-                    case ArrayToken arrayToken:
-                    {
-                        var newArray = new List<IToken>(arrayToken.Length);
-                        foreach (var token in arrayToken.Data)
-                        {
-                            newArray.Add(CopyToken(token, tokenScanner));
-                        }
-
-                        return new ArrayToken(newArray);
-                    }
-                    case IndirectReferenceToken referenceToken:
-                    {
-                        if (referencesFromDocument.TryGetValue(referenceToken, out var newReferenceToken))
-                        {
-                            return newReferenceToken;
-                        }
-                        
-                        var tokenObject = DirectObjectFinder.Get<IToken>(referenceToken.Data, tokenScanner);
-
-                        Debug.Assert(!(tokenObject is IndirectReferenceToken));
-
-                        var newToken = CopyToken(tokenObject, tokenScanner);
-                        newReferenceToken = context.WriteToken(newToken);
-                        
-                        referencesFromDocument.Add(referenceToken, newReferenceToken);
-                        
-                        return newReferenceToken;
-                    }
-                    case StreamToken streamToken:
-                    {
-                        var properties = CopyToken(streamToken.StreamDictionary, tokenScanner) as DictionaryToken;
-                        Debug.Assert(properties != null);
-
-                        var bytes = streamToken.Data;
-                        return new StreamToken(properties, bytes);
-                    }
-                    
-                    case ObjectToken _:
-                    {
-                        // Since we don't write token directly to the stream.
-                        // We can't know the offset. Therefore the token would be invalid
-                        throw new NotSupportedException("Copying a Object token is not supported");
-                    }
-                }
-
-                return tokenToCopy;
-            }
-
-            private static bool IgnoreKeyForPagesNode(KeyValuePair<string, IToken> token)
-            {
-                return string.Equals(token.Key, NameToken.Type.Data, StringComparison.OrdinalIgnoreCase)
-                    || string.Equals(token.Key, NameToken.Kids.Data, StringComparison.OrdinalIgnoreCase)
-                    || string.Equals(token.Key, NameToken.Count.Data, StringComparison.OrdinalIgnoreCase)
-                    || string.Equals(token.Key, NameToken.Parent.Data, StringComparison.OrdinalIgnoreCase);
-            }
         }
     }
 }

src/UglyToad.PdfPig/Writer/PdfStreamWriter.cs

@@ -9,7 +9,8 @@
     using Tokens;
 
     /// <summary>
-    /// This class would lazily flush all token. Allowing us to make changes to references without need to rewrite the whole stream
+    /// This class would lazily flush all token.
+    /// Allowing us to make changes to references without need to rewrite the whole stream
     /// </summary>
     internal class PdfStreamWriter : IDisposable
     {
@@ -17,20 +18,34 @@
 
         private readonly Dictionary<IndirectReferenceToken, IToken> tokenReferences = new Dictionary<IndirectReferenceToken, IToken>();
 
-        public int CurrentNumber { get; private set; } = 1;
+        private int currentNumber = 1;
 
-        public Stream Stream { get; private set; }
+        private Stream stream;
 
+        /// <summary>
+        /// Flag to set whether or not we want to dispose the stream
+        /// </summary>
         public bool DisposeStream { get; set; }
 
+        /// <summary>
+        /// Construct a PdfStreamWriter with a memory stream
+        /// </summary>
         public PdfStreamWriter() : this(new MemoryStream()) { }
 
+        /// <summary>
+        /// Construct a PdfStreamWriter
+        /// </summary>
         public PdfStreamWriter(Stream baseStream, bool disposeStream = true)
         {
-            Stream = baseStream ?? throw new ArgumentNullException(nameof(baseStream));
+            stream = baseStream ?? throw new ArgumentNullException(nameof(baseStream));
             DisposeStream = disposeStream;
         }
 
+        /// <summary>
+        /// Flush the document with all the token that we have accumulated
+        /// </summary>
+        /// <param name="version">Pdf Version that we are targeting</param>
+        /// <param name="catalogReference">Catalog's indirect reference token to which the token are related</param>
         public void Flush(decimal version, IndirectReferenceToken catalogReference)
         {
             if (catalogReference == null)
@@ -38,14 +53,14 @@
                 throw new ArgumentNullException(nameof(catalogReference));
             }
 
-            WriteString($"%PDF-{version.ToString("0.0", CultureInfo.InvariantCulture)}", Stream);
+            WriteString($"%PDF-{version.ToString("0.0", CultureInfo.InvariantCulture)}", stream);
 
-            Stream.WriteText("%");
-            Stream.WriteByte(169);
-            Stream.WriteByte(205);
-            Stream.WriteByte(196);
-            Stream.WriteByte(210);
-            Stream.WriteNewLine();
+            stream.WriteText("%");
+            stream.WriteByte(169);
+            stream.WriteByte(205);
+            stream.WriteByte(196);
+            stream.WriteByte(210);
+            stream.WriteNewLine();
 
             var offsets = new Dictionary<IndirectReference, long>();
             ObjectToken catalogToken = null;
@@ -53,10 +68,10 @@
             {
                 var referenceToken = pair.Key;
                 var token = pair.Value;
-                var offset = Stream.Position;
+                var offset = stream.Position;
                 var obj = new ObjectToken(offset, referenceToken.Data, token);
 
-                TokenWriter.WriteToken(obj, Stream);
+                TokenWriter.WriteToken(obj, stream);
 
                 offsets.Add(referenceToken.Data, offset);
 
@@ -72,75 +87,105 @@
             }
 
             // TODO: Support document information
-            TokenWriter.WriteCrossReferenceTable(offsets, catalogToken, Stream, null);
+            TokenWriter.WriteCrossReferenceTable(offsets, catalogToken, stream, null);
         }
 
+        /// <summary>
+        /// Push a new token to be written
+        /// </summary>
+        /// <param name="token"></param>
+        /// <param name="reservedNumber"></param>
+        /// <returns></returns>
         public IndirectReferenceToken WriteToken(IToken token, int? reservedNumber = null)
         {
-            // if you can't consider deduplicating the token. 
-            // It must be because it's referenced by his child element, so you must have reserved a number before hand
-            // Example /Pages Obj
-            var canBeDuplicated = !reservedNumber.HasValue;
-            if (!canBeDuplicated)
+            if (reservedNumber.HasValue)
             {
                 if (!reservedNumbers.Remove(reservedNumber.Value))
                 {
                     throw new InvalidOperationException("You can't reuse a reserved number");
                 }
 
-                // When we end up writing this token, all of his child would already have been added and checked for duplicate
                 return AddToken(token, reservedNumber.Value);
             }
 
-            var reference = FindToken(token);
-            if (reference == null)
-            {
-                return AddToken(token, CurrentNumber++);
-            }
-
-            return reference;
+            return AddToken(token, currentNumber++);
         }
 
+        /// <summary>
+        /// Get a token based on his indirect reference
+        /// </summary>
+        /// <param name="referenceToken"></param>
+        /// <returns></returns>
+        public IToken GetToken(IndirectReferenceToken referenceToken)
+        {
+            return tokenReferences.TryGetValue(referenceToken, out var token) ? token : null;
+        }
+
+        /// <summary>
+        /// Replace a token base on his indirect reference
+        /// </summary>
+        /// <param name="referenceToken"></param>
+        /// <param name="newToken"></param>
+        public void ReplaceToken(IndirectReferenceToken referenceToken, IToken newToken)
+        {
+            tokenReferences[referenceToken] = newToken;
+        }
+
+        /// <summary>
+        /// Reserve a number for a token
+        /// </summary>
+        /// <returns></returns>
         public int ReserveNumber()
         {
-            var reserved = CurrentNumber;
+            var reserved = currentNumber;
             reservedNumbers.Add(reserved);
-            CurrentNumber++;
+            currentNumber++;
             return reserved;
         }
 
+        /// <summary>
+        /// Reserve a number and create a token with it
+        /// </summary>
+        /// <returns></returns>
         public IndirectReferenceToken ReserveNumberToken()
         {
             return new IndirectReferenceToken(new IndirectReference(ReserveNumber(), 0));
         }
 
+        /// <summary>
+        /// Return the bytes that have been flushed to the stream
+        /// </summary>
+        /// <returns></returns>
         public byte[] ToArray()
         {
-            var currentPosition = Stream.Position;
-            Stream.Seek(0, SeekOrigin.Begin);
+            var currentPosition = stream.Position;
+            stream.Seek(0, SeekOrigin.Begin);
 
-            var bytes = new byte[Stream.Length];
+            var bytes = new byte[stream.Length];
 
-            if (Stream.Read(bytes, 0, bytes.Length) != bytes.Length)
+            if (stream.Read(bytes, 0, bytes.Length) != bytes.Length)
             {
-                throw new Exception("Unable to read all the bytes from stream");
+                throw new IOException("Unable to read all the bytes from stream");
             }
 
-            Stream.Seek(currentPosition, SeekOrigin.Begin);
+            stream.Seek(currentPosition, SeekOrigin.Begin);
 
             return bytes;
         }
 
+        /// <summary>
+        /// Dispose the stream if the PdfStreamWriter#DisposeStream flag is set
+        /// </summary>
         public void Dispose()
         {
             if (!DisposeStream)
             {
-                Stream = null;
+                stream = null;
                 return;
             }
 
-            Stream?.Dispose();
-            Stream = null;
+            stream?.Dispose();
+            stream = null;
         }
 
         private IndirectReferenceToken AddToken(IToken token, int reservedNumber)
@@ -151,21 +196,6 @@
             return referenceToken;
         }
 
-        private IndirectReferenceToken FindToken(IToken token)
-        {
-            foreach (var pair in tokenReferences)
-            {
-                var reference = pair.Key;
-                var storedToken = pair.Value;
-                if (storedToken.Equals(token))
-                {
-                    return reference;
-                }
-            }
-
-            return null;
-        }
-
         private static void WriteString(string text, Stream stream)
         {
             var bytes = OtherEncodings.StringAsLatin1Bytes(text);