PdfStreamWriter: Error Checking and Code Formatting

src/UglyToad.PdfPig/Writer/PdfStreamWriter.cs

@@ -1,38 +1,37 @@
-using System;
-using System.Collections.Generic;
-using System.Diagnostics;
-using System.IO;
-using UglyToad.PdfPig.Core;
-using UglyToad.PdfPig.Graphics.Operations;
-using UglyToad.PdfPig.Tokens;
-
-namespace UglyToad.PdfPig.Writer
+namespace UglyToad.PdfPig.Writer
 {
+    using System;
+    using System.Collections.Generic;
+    using System.IO;
+    using Core;
+    using Graphics.Operations;
+    using Tokens;
+    
     /// <summary>
     /// 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
     {
-        private readonly SortedSet<int> reservedNumbers = new SortedSet<int>();
+        private readonly List<int> reservedNumbers = new List<int>();
 
         private readonly Dictionary<IndirectReferenceToken, IToken> tokenReferences = new Dictionary<IndirectReferenceToken, IToken>();
 
         public int CurrentNumber { get; private set; } = 1;
 
-        public Stream Stream { get; }
+        public Stream Stream { get; private set; }
 
         public PdfStreamWriter() : this(new MemoryStream()) { }
 
         public PdfStreamWriter(Stream baseStream)
         {
-            Stream = baseStream;
+            Stream = baseStream ?? throw new ArgumentNullException(nameof(baseStream));
         }
         
         public void Flush(decimal version, IndirectReferenceToken catalogReference)
         {
             if (catalogReference == null)
                 throw new ArgumentNullException(nameof(catalogReference));
-
+            
             WriteString($"%PDF-{version:0.0}", Stream);
 
             Stream.WriteText("%");
@@ -57,7 +56,7 @@ namespace UglyToad.PdfPig.Writer
 
                 if (catalogToken == null && referenceToken == catalogReference)
                 {
-                    catalogToken = new ObjectToken(offset, referenceToken.Data, token);
+                    catalogToken = obj;
                 }
             }
 
@@ -70,9 +69,9 @@ namespace UglyToad.PdfPig.Writer
             TokenWriter.WriteCrossReferenceTable(offsets, catalogToken, Stream, null);
         }
 
-        public IndirectReferenceToken WriteObject(IToken token, int? reservedNumber = null)
+        public IndirectReferenceToken WriteToken(IToken token, int? reservedNumber = null)
         {
-            // if you can't consider deduplicating a token. 
+            // 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;
@@ -84,20 +83,57 @@ namespace UglyToad.PdfPig.Writer
                 }
 
                 // When we end up writing this token, all of his child would already have been added and checked for duplicate
-                return AddObject(token, reservedNumber.Value);
+                return AddToken(token, reservedNumber.Value);
             }
 
             var reference = FindToken(token);
             if (reference == null)
             {
-                // TODO: Check his children
-                return AddObject(token, CurrentNumber++);
+                return AddToken(token, CurrentNumber++);
             }
 
             return reference;
         }
+        
+        public int ReserveNumber()
+        {
+            var reserved = CurrentNumber;
+            reservedNumbers.Add(reserved);
+            CurrentNumber++;
+            return reserved;
+        }
 
-        private IndirectReferenceToken AddObject(IToken token, int reservedNumber)
+        public IndirectReferenceToken ReserveNumberToken()
+        { 
+            return new IndirectReferenceToken(new IndirectReference(ReserveNumber(), 0));
+        }
+        
+        public byte[] ToArray()
+        {
+            if (!Stream.CanSeek)
+                throw new NotSupportedException($"{Stream.GetType()} can't seek");
+            
+            var currentPosition = Stream.Position;
+            Stream.Seek(0, SeekOrigin.Begin);
+
+            var bytes = new byte[Stream.Length];
+
+            // Should we slice the reading into smaller chunks?
+            if (Stream.Read(bytes, 0, bytes.Length) != bytes.Length)
+                throw new Exception("Unable to read all the bytes from stream");
+
+            Stream.Seek(currentPosition, SeekOrigin.Begin);
+
+            return bytes;
+        }
+
+        public void Dispose()
+        {
+            Stream?.Dispose();
+            Stream = null;
+        }
+        
+        private IndirectReferenceToken AddToken(IToken token, int reservedNumber)
         {
             var reference = new IndirectReference(reservedNumber, 0);
             var referenceToken = new IndirectReferenceToken(reference);
@@ -119,46 +155,12 @@ namespace UglyToad.PdfPig.Writer
 
             return null;
         }
-
-        public int ReserveNumber()
-        {
-            var reserved = CurrentNumber;
-            reservedNumbers.Add(reserved);
-            CurrentNumber++;
-            return reserved;
-        }
-
-        public IndirectReferenceToken ReserveNumberToken() => new IndirectReferenceToken(new IndirectReference(ReserveNumber(), 0));
-
+        
         private static void WriteString(string text, Stream stream)
         {
             var bytes = OtherEncodings.StringAsLatin1Bytes(text);
             stream.Write(bytes, 0, bytes.Length);
             stream.WriteNewLine();
         }
-
-        public byte[] ToArray()
-        {
-            if (!Stream.CanSeek)
-                throw new NotSupportedException("Stream can't seek");
-            
-            var currentPosition = Stream.Position;
-            Stream.Seek(0, SeekOrigin.Begin);
-
-            var bytes = new byte[Stream.Length];
-
-            // Should we slice the reading into smaller chunks?
-            if (Stream.Read(bytes, 0, bytes.Length) != bytes.Length)
-                throw new Exception("Unable to read all the bytes from stream");
-
-            Stream.Seek(currentPosition, SeekOrigin.Begin);
-
-            return bytes;
-        }
-
-        public void Dispose()
-        {
-            Stream.Dispose();
-        }
     }
 }