Use new SegmentedList and SegmentedDictionary to process large heap dumps

src/FastSerialization/SegmentedDictionary/SegmentedDictionary.cs

@@ -781,7 +781,8 @@ public int EnsureCapacity(int capacity)
                 throw new ArgumentOutOfRangeException(nameof(capacity));
             }
 
-            var currentCapacity = _entries.Capacity;
+            // Normal usage of a dictionary should never ask for a capacity that exceeds int32.MaxValue.
+            var currentCapacity = (int)_entries.Capacity;
             if (currentCapacity >= capacity)
             {
                 return currentCapacity;

src/FastSerialization/SegmentedList.cs

@@ -100,7 +100,7 @@ public long Count
             }
         }
 
-        public int Capacity => this.capacity;
+        internal long Capacity => this.capacity;
 
         /// <summary>
         /// Copy to Array
@@ -177,7 +177,7 @@ public T this[long index]
             }
         }
 
-        public ref T GetElementByReference(int index) =>
+        internal ref T GetElementByReference(int index) =>
             ref this.items[index >> this.segmentShift][index & this.offsetMask];
 
         /// <summary>
@@ -190,6 +190,18 @@ public bool IsValidIndex(long index)
             return this.items[index >> this.segmentShift] != null;
         }
 
+        /// <summary>
+        /// Get slot of an element
+        /// </summary>
+        /// <param name="index"></param>
+        /// <param name="slot"></param>
+        /// <returns></returns>
+        public T[] GetSlot(int index, out int slot)
+        {
+            slot = index & this.offsetMask;
+            return this.items[index >> this.segmentShift];
+        }
+
         /// <summary>
         /// Adds new element at the end of the list.
         /// </summary>
@@ -443,12 +455,12 @@ public void CopyTo(T[] array, int arrayIndex)
                     "Destination array is not long enough to copy all the items in the collection. Check array index and length.");
             }
 
-            int remain = (int)this.count;
+            long remain = this.count;
 
-            for (int i = 0; (remain > 0) && (i < this.items.Length); i++)
+            for (long i = 0; (remain > 0) && (i < this.items.Length); i++)
             {
                 // We can safely cast to int, since that is the max value that items[i].Length can have.
-                int len = Math.Min(remain, this.items[i].Length);
+                int len = (int)Math.Min(remain, this.items[i].Length);
 
                 Array.Copy(this.items[i], 0, array, arrayIndex, len);
 

src/FastSerialization/SegmentedMemoryStreamReader.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
-using System.IO;
 using System.Text;
 
 namespace FastSerialization
@@ -13,12 +12,31 @@ public class SegmentedMemoryStreamReader
         /// <summary>
         /// Create a IStreamReader (reads binary data) from a given byte buffer
         /// </summary>
-        public SegmentedMemoryStreamReader(SegmentedList<byte> data) : this(data, 0, (int)data.Count) { }
+        public SegmentedMemoryStreamReader(SegmentedList<byte> data, SerializationConfiguration config = null) : this(data, 0, data.Count, config) { }
         /// <summary>
         /// Create a IStreamReader (reads binary data) from a given subregion of a byte buffer 
         /// </summary>
-        public SegmentedMemoryStreamReader(SegmentedList<byte> data, int start, int length)
+        public SegmentedMemoryStreamReader(SegmentedList<byte> data, long start, long length, SerializationConfiguration config = null)
         {
+            SerializationConfiguration = config ?? new SerializationConfiguration();
+
+            if (SerializationConfiguration.StreamLabelWidth == StreamLabelWidth.FourBytes)
+            {
+                readLabel = () =>
+                {
+                    return (StreamLabel)(uint)ReadInt32();
+                };
+                sizeOfSerializedStreamLabel = 4;
+            }
+            else
+            {
+                readLabel = () =>
+                {
+                    return (StreamLabel)(ulong)ReadInt64();
+                };
+                sizeOfSerializedStreamLabel = 8;
+            }
+
             bytes = new SegmentedList<byte>(65_536, length);
             bytes.AppendFrom(data, start, length);
             position = start;
@@ -130,18 +148,24 @@ public string ReadString()
         /// <summary>
         /// Implementation of IStreamReader
         /// </summary>
-        public StreamLabel ReadLabel()
-        {
-            return (StreamLabel)(uint)ReadInt32();
-        }
+        public StreamLabel ReadLabel() => readLabel();
+
         /// <summary>
         /// Implementation of IStreamReader
         /// </summary>
         public virtual void Goto(StreamLabel label)
         {
             Debug.Assert(label != StreamLabel.Invalid);
-            Debug.Assert((long)label <= int.MaxValue);
-            position = (int)label;
+
+            if (SerializationConfiguration.StreamLabelWidth == StreamLabelWidth.FourBytes)
+            {
+                Debug.Assert((long)label <= int.MaxValue);
+                position = (uint)label;
+            }
+            else
+            {
+                position = (long)label;
+            }
         }
         /// <summary>
         /// Implementation of IStreamReader
@@ -150,16 +174,22 @@ public virtual StreamLabel Current
         {
             get
             {
-                return (StreamLabel)(uint)position;
+                if (SerializationConfiguration.StreamLabelWidth == StreamLabelWidth.FourBytes)
+                {
+                    return (StreamLabel)(uint)position;
+                }
+                else
+                {
+                    return (StreamLabel)position;
+                }
             }
         }
         /// <summary>
         /// Implementation of IStreamReader
         /// </summary>
         public virtual void GotoSuffixLabel()
         {
-            const int serializedStreamLabelSize = 4;
-            Goto((StreamLabel)(Length - serializedStreamLabelSize));
+            Goto((StreamLabel)(Length - sizeOfSerializedStreamLabel));
             Goto(ReadLabel());
         }
         /// <summary>
@@ -174,6 +204,11 @@ public void Dispose()
         /// Dispose pattern
         /// </summary>
         protected virtual void Dispose(bool disposing) { }
+
+        /// <summary>
+        /// Returns the SerializationConfiguration for this stream reader.
+        /// </summary>
+        public SerializationConfiguration SerializationConfiguration { get; private set; }
         #endregion
 
         #region private 
@@ -182,9 +217,11 @@ protected virtual void Dispose(bool disposing) { }
             throw new Exception("Streamreader read past end of buffer");
         }
         internal /*protected*/  SegmentedList<byte> bytes;
-        internal /*protected*/  int position;
-        internal /*protected*/  int endPosition;
+        internal /*protected*/  long position;
+        internal /*protected*/  long endPosition;
         private StringBuilder sb;
+        private Func<StreamLabel> readLabel;
+        private readonly int sizeOfSerializedStreamLabel;
         #endregion
     }
 }

src/FastSerialization/SegmentedMemoryStreamWriter.cs

@@ -1,16 +1,31 @@
 using System;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.IO;
-using System.Text;
 
 namespace FastSerialization
 {
     public class SegmentedMemoryStreamWriter
     {
-        public SegmentedMemoryStreamWriter() : this(64) { }
-        public SegmentedMemoryStreamWriter(int initialSize)
+        public SegmentedMemoryStreamWriter(SerializationConfiguration config = null) : this(64, config) { }
+        public SegmentedMemoryStreamWriter(long initialSize, SerializationConfiguration config = null)
         {
+            SerializationConfiguration = config ?? new SerializationConfiguration();
+
+            if (SerializationConfiguration.StreamLabelWidth == StreamLabelWidth.FourBytes)
+            {
+                writeLabel = (value) =>
+                {
+                    Write((int)value);
+                };
+            }
+            else
+            {
+                writeLabel = (value) =>
+                {
+                    Write((long)value);
+                };
+            }
+
             bytes = new SegmentedList<byte>(65_536, initialSize);
         }
 
@@ -45,10 +60,8 @@ public void Write(long value)
             bytes.Add((byte)value); value = value >> 8;
             bytes.Add((byte)value); value = value >> 8;
         }
-        public void Write(StreamLabel value)
-        {
-            Write((int)value);
-        }
+        public void Write(StreamLabel value) => writeLabel(value);
+
         public void Write(string value)
         {
             if (value == null)
@@ -85,10 +98,15 @@ public void WriteToStream(Stream outputStream)
         public SegmentedMemoryStreamReader GetReader()
         {
             var readerBytes = bytes;
-            return new SegmentedMemoryStreamReader(readerBytes, 0, (int)readerBytes.Count);
+            return new SegmentedMemoryStreamReader(readerBytes, 0, readerBytes.Count, SerializationConfiguration);
         }
         public void Dispose() { }
 
+        /// <summary>
+        /// Returns the SerializationConfiguration for this stream writer.
+        /// </summary>
+        public SerializationConfiguration SerializationConfiguration { get; private set; }
+
         #region private
         protected virtual void MakeSpace()
         {
@@ -101,6 +119,7 @@ public byte[] GetBytes()
         }
 
         protected SegmentedList<byte> bytes;
+        private Action<StreamLabel> writeLabel;
         #endregion
     }
 }

src/HeapDump/GCHeapDump.cs

@@ -709,7 +709,7 @@ void IFastSerializable.FromStream(Deserializer deserializer)
 /// </graph>
 /// 
 /// </summary>
-public class XmlGcHeapDump
+internal class XmlGcHeapDump
 {
     public static GCHeapDump ReadGCHeapDumpFromXml(string fileName)
     {
@@ -840,7 +840,7 @@ public static MemoryGraph ReadMemoryGraphFromXml(XmlReader reader)
         return graph;
     }
 
-    public static void WriteGCDumpToXml(GCHeapDump gcDump, StreamWriter writer)
+    internal static void WriteGCDumpToXml(GCHeapDump gcDump, StreamWriter writer)
     {
         writer.WriteLine("<GCHeapDump>");
 

src/MemoryGraph/MemoryGraph.cs

@@ -12,7 +12,7 @@ public MemoryGraph(int expectedSize, bool isVeryLargeGraph = false)
         {
             // If we have too many addresses we will reach the Dictionary's internal array's size limit and throw.
             // Therefore use a new implementation of it that is similar in performance but that can handle the extra load.
-            if (expectedSize > 200_000)
+            if (isVeryLargeGraph)
             {
                 m_addressToNodeIndex = new SegmentedDictionary<Address, NodeIndex>(expectedSize);
             }

src/MemoryGraph/graph.cs

@@ -5,7 +5,6 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
-using System.Text;
 using System.Text.RegularExpressions;
 using Address = System.UInt64;
 
@@ -120,9 +119,9 @@ public virtual NodeType AllocTypeNodeStorage()
         /// </summary>
         public NodeIndex NodeIndexLimit { get { return (NodeIndex)m_nodes.Count; } }
         /// <summary>
-        /// Same as NodeIndexLimit, just cast to an integer.  
+        /// Same as NodeIndexLimit.  
         /// </summary>
-        public int NodeCount { get { return (int)m_nodes.Count; } }
+        public long NodeCount { get { return m_nodes.Count; } }
         /// <summary>
         /// It is expected that users will want additional information associated with TYPES of the nodes of the graph.  They can
         /// do this by allocating an array of code:NodeTypeIndexLimit and then indexing this by code:NodeTypeIndex
@@ -161,8 +160,11 @@ public virtual NodeType AllocTypeNodeStorage()
         /// 
         /// TODO I can eliminate the need for AllowReading.  
         /// </summary>
+        /// <remarks>if isVeryLargeGraph argument is true, then StreamLabels will be serialized as longs
+        /// too acommodate for the extra size of the graph's stream representation.</remarks>
         public Graph(int expectedNodeCount, bool isVeryLargeGraph = false)
         {
+            m_isVeryLargeGraph = isVeryLargeGraph;
             m_expectedNodeCount = expectedNodeCount;
             m_types = new GrowableArray<TypeInfo>(Math.Max(expectedNodeCount / 100, 2000));
             m_nodes = new SegmentedList<StreamLabel>(SegmentSize, m_expectedNodeCount);
@@ -462,7 +464,8 @@ private void ClearWorker()
             RootIndex = NodeIndex.Invalid;
             if (m_writer == null)
             {
-                m_writer = new SegmentedMemoryStreamWriter(m_expectedNodeCount * 8);
+                m_writer = new SegmentedMemoryStreamWriter(m_expectedNodeCount * 8,
+                    m_isVeryLargeGraph ? new SerializationConfiguration() { StreamLabelWidth = StreamLabelWidth.EightBytes } : null);
             }
 
             m_totalSize = 0;
@@ -590,7 +593,9 @@ public void FromStream(Deserializer deserializer)
             // Read in the Blob stream.  
             // TODO be lazy about reading in the blobs.  
             int blobCount = deserializer.ReadInt();
-            SegmentedMemoryStreamWriter writer = new SegmentedMemoryStreamWriter(blobCount);
+            SegmentedMemoryStreamWriter writer = new SegmentedMemoryStreamWriter(blobCount,
+                m_isVeryLargeGraph ? new SerializationConfiguration() { StreamLabelWidth = StreamLabelWidth.EightBytes } : null);
+
             while (8 <= blobCount)
             {
                 writer.Write(deserializer.ReadInt64());
@@ -649,7 +654,7 @@ public void FromStream(Deserializer deserializer)
             }
         }
 
-        private int m_expectedNodeCount;                // Initial guess at graph Size. 
+        private long m_expectedNodeCount;                // Initial guess at graph Size.
         private long m_totalSize;                       // Total Size of all the nodes in the graph.  
         internal int m_totalRefs;                       // Total Number of references in the graph
         internal GrowableArray<TypeInfo> m_types;       // We expect only thousands of these