CSHARP-5348: Avoid allocations for Bson*Context

src/MongoDB.Bson/IO/BsonBinaryReader.cs

@@ -14,8 +14,10 @@
 */
 
 using System;
+using System.Collections.Generic;
 using System.Globalization;
 using System.IO;
+using System.Linq;
 
 namespace MongoDB.Bson.IO
 {
@@ -30,6 +32,7 @@ public class BsonBinaryReader : BsonReader
 #pragma warning restore CA2213 // Disposable never disposed
         private readonly BsonStream _bsonStream;
         private BsonBinaryReaderContext _context;
+        private readonly Lazy<Stack<BsonBinaryReaderContext>> _contexts = new(() => new());
 
         // constructors
         /// <summary>
@@ -61,7 +64,7 @@ public BsonBinaryReader(Stream stream, BsonBinaryReaderSettings settings)
             _baseStream = stream;
             _bsonStream = (stream as BsonStream) ?? new BsonStreamAdapter(stream);
 
-            _context = new BsonBinaryReaderContext(null, ContextType.TopLevel, 0, 0);
+            _context = new BsonBinaryReaderContext(ContextType.TopLevel, 0, 0);
         }
 
         // public properties
@@ -109,10 +112,8 @@ public override void Close()
         /// Gets a bookmark to the reader's current position and state.
         /// </summary>
         /// <returns>A bookmark.</returns>
-        public override BsonReaderBookmark GetBookmark()
-        {
-            return new BsonBinaryReaderBookmark(State, CurrentBsonType, CurrentName, _context, _bsonStream.Position);
-        }
+        public override BsonReaderBookmark GetBookmark() =>
+            new BsonBinaryReaderBookmark(State, CurrentBsonType, CurrentName, _context, _contexts.Value.ToArray(), _bsonStream.Position);
 
         /// <summary>
         /// Determines whether this reader is at end of file.
@@ -201,12 +202,11 @@ public override BsonType ReadBsonType()
 
             if (_context.ContextType == ContextType.Array)
             {
-                _context.CurrentArrayIndex++;
+                _context.ArrayIndex++;
             }
 
             try
             {
-
                 CurrentBsonType = _bsonStream.ReadBsonType();
             }
             catch (FormatException ex)
@@ -342,7 +342,8 @@ public override void ReadEndArray()
                 ThrowInvalidState("ReadEndArray", BsonReaderState.EndOfArray);
             }
 
-            _context = _context.PopContext(_bsonStream.Position);
+            _context = PopContext();
+
             switch (_context.ContextType)
             {
                 case ContextType.Array: State = BsonReaderState.Type; break;
@@ -371,10 +372,10 @@ public override void ReadEndDocument()
                 ThrowInvalidState("ReadEndDocument", BsonReaderState.EndOfDocument);
             }
 
-            _context = _context.PopContext(_bsonStream.Position);
+            _context = PopContext();
             if (_context.ContextType == ContextType.JavaScriptWithScope)
             {
-                _context = _context.PopContext(_bsonStream.Position); // JavaScriptWithScope
+                _context = PopContext(); // JavaScriptWithScope
             }
             switch (_context.ContextType)
             {
@@ -432,7 +433,10 @@ public override string ReadJavaScriptWithScope()
 
             var startPosition = _bsonStream.Position; // position of size field
             var size = ReadSize();
-            _context = new BsonBinaryReaderContext(_context, ContextType.JavaScriptWithScope, startPosition, size);
+
+            PushContext(_context);
+            _context = new BsonBinaryReaderContext(ContextType.JavaScriptWithScope, startPosition, size);
+
             var code = _bsonStream.ReadString(Settings.Encoding);
 
             State = BsonReaderState.ScopeDocument;
@@ -486,7 +490,7 @@ public override string ReadName(INameDecoder nameDecoder)
 
             if (_context.ContextType == ContextType.Document)
             {
-                _context.CurrentElementName = CurrentName;
+                _context.ElementName = CurrentName;
             }
 
             return CurrentName;
@@ -553,7 +557,7 @@ public override IByteBuffer ReadRawBsonDocument()
 
             if (_context.ContextType == ContextType.JavaScriptWithScope)
             {
-                _context = _context.PopContext(_bsonStream.Position); // JavaScriptWithScope
+                _context = PopContext(); // JavaScriptWithScope
             }
             switch (_context.ContextType)
             {
@@ -590,7 +594,10 @@ public override void ReadStartArray()
 
             var startPosition = _bsonStream.Position; // position of size field
             var size = ReadSize();
-            _context = new BsonBinaryReaderContext(_context, ContextType.Array, startPosition, size);
+
+            PushContext(_context);
+            _context = new(ContextType.Array, startPosition, size);
+
             State = BsonReaderState.Type;
         }
 
@@ -605,7 +612,8 @@ public override void ReadStartDocument()
             var contextType = (State == BsonReaderState.ScopeDocument) ? ContextType.ScopeDocument : ContextType.Document;
             var startPosition = _bsonStream.Position; // position of size field
             var size = ReadSize();
-            _context = new BsonBinaryReaderContext(_context, contextType, startPosition, size);
+            PushContext(_context);
+            _context = new(contextType, startPosition, size);
             State = BsonReaderState.Type;
         }
 
@@ -665,7 +673,15 @@ public override void ReturnToBookmark(BsonReaderBookmark bookmark)
             State = binaryReaderBookmark.State;
             CurrentBsonType = binaryReaderBookmark.CurrentBsonType;
             CurrentName = binaryReaderBookmark.CurrentName;
-            _context = binaryReaderBookmark.CloneContext();
+
+            _context = binaryReaderBookmark.CurrentContext;
+
+            _contexts.Value.Clear();
+            foreach (var context in binaryReaderBookmark.ContextsStack.Reverse())
+            {
+                _contexts.Value.Push(context);
+            }
+
             _bsonStream.Position = binaryReaderBookmark.Position;
         }
 
@@ -686,7 +702,7 @@ public override void SkipName()
 
             if (_context.ContextType == ContextType.Document)
             {
-                _context.CurrentElementName = CurrentName;
+                _context.ElementName = CurrentName;
             }
         }
 
@@ -767,35 +783,41 @@ private string GenerateDottedElementName()
             }
             else if (_context.ContextType == ContextType.Array)
             {
-                elementName = _context.CurrentArrayIndex.ToString(NumberFormatInfo.InvariantInfo);
+                elementName = _context.ArrayIndex.ToString(NumberFormatInfo.InvariantInfo);
             }
             else
             {
                 elementName = "?";
             }
 
-            return GenerateDottedElementName(_context.ParentContext, elementName);
+            return GenerateDottedElementName(_contexts.Value.ToArray(), 0, elementName);
         }
 
-        private string GenerateDottedElementName(BsonBinaryReaderContext context, string elementName)
+        private string GenerateDottedElementName(BsonBinaryReaderContext[] contexts, int currentContextIndex, string elementName)
         {
+            if (currentContextIndex >= contexts.Length)
+                return elementName;
+
+            var context = contexts[currentContextIndex];
+            var nextIndex = currentContextIndex + 1;
+
             if (context.ContextType == ContextType.Document)
             {
-                return GenerateDottedElementName(context.ParentContext, (context.CurrentElementName ?? "?") + "." + elementName);
+                return GenerateDottedElementName(contexts,  nextIndex, (context.ElementName ?? "?") + "." + elementName);
             }
-            else if (context.ContextType == ContextType.Array)
-            {
-                var indexElementName = context.CurrentArrayIndex.ToString(NumberFormatInfo.InvariantInfo);
-                return GenerateDottedElementName(context.ParentContext, indexElementName + "." + elementName);
-            }
-            else if (context.ParentContext != null)
+
+            if (context.ContextType == ContextType.Array)
             {
-                return GenerateDottedElementName(context.ParentContext, "?." + elementName);
+                var indexElementName = context.ArrayIndex.ToString(NumberFormatInfo.InvariantInfo);
+                return GenerateDottedElementName(contexts,  nextIndex, indexElementName + "." + elementName);
             }
-            else
+
+            if (nextIndex < contexts.Length)
             {
-                return elementName;
+                return GenerateDottedElementName(contexts, nextIndex, "?." + elementName);
             }
+
+            return elementName;
         }
 
         private BsonReaderState GetNextState()
@@ -828,5 +850,21 @@ private int ReadSize()
             }
             return size;
         }
+
+        private BsonBinaryReaderContext PopContext()
+        {
+            var actualSize = _bsonStream.Position - _context.StartPosition;
+            if (actualSize != _context.Size)
+            {
+                throw new FormatException($"Expected size to be {_context.Size}, not {actualSize}.");
+            }
+
+            return _contexts.Value.Pop();
+        }
+
+        private void PushContext(BsonBinaryReaderContext context)
+        {
+            _contexts.Value.Push(context);
+        }
     }
 }

src/MongoDB.Bson/IO/BsonBinaryReaderBookmark.cs

@@ -13,40 +13,28 @@
 * limitations under the License.
 */
 
-namespace MongoDB.Bson.IO
+namespace MongoDB.Bson.IO;
+
+/// <summary>
+/// Represents a bookmark that can be used to return a reader to the current position and state.
+/// </summary>
+public class BsonBinaryReaderBookmark : BsonReaderBookmark
 {
-    /// <summary>
-    /// Represents a bookmark that can be used to return a reader to the current position and state.
-    /// </summary>
-    public class BsonBinaryReaderBookmark : BsonReaderBookmark
+    internal BsonBinaryReaderBookmark(
+        BsonReaderState state,
+        BsonType currentBsonType,
+        string currentName,
+        BsonBinaryReaderContext currentContext,
+        BsonBinaryReaderContext[] contextsStack,
+        long position)
+        : base(state, currentBsonType, currentName)
     {
-        // private fields
-        private BsonBinaryReaderContext _context;
-        private long _position;
-
-        // constructors
-        internal BsonBinaryReaderBookmark(
-            BsonReaderState state,
-            BsonType currentBsonType,
-            string currentName,
-            BsonBinaryReaderContext context,
-            long position)
-            : base(state, currentBsonType, currentName)
-        {
-            _context = context.Clone();
-            _position = position;
-        }
-
-        // internal properties
-        internal long Position
-        {
-            get { return _position; }
-        }
-
-        // internal methods
-        internal BsonBinaryReaderContext CloneContext()
-        {
-            return _context.Clone();
-        }
+        CurrentContext = currentContext;
+        ContextsStack = contextsStack;
+        Position = position;
     }
+
+    internal BsonBinaryReaderContext CurrentContext { get; }
+    internal BsonBinaryReaderContext[] ContextsStack { get; }
+    internal long Position { get; }
 }

src/MongoDB.Bson/IO/BsonBinaryReaderContext.cs

@@ -13,75 +13,17 @@
 * limitations under the License.
 */
 
-using System;
+namespace MongoDB.Bson.IO;
 
-namespace MongoDB.Bson.IO
+internal struct BsonBinaryReaderContext(
+    ContextType contextType,
+    long startPosition,
+    long size)
 {
-    internal class BsonBinaryReaderContext
-    {
-        // private fields
-        private readonly BsonBinaryReaderContext _parentContext;
-        private readonly ContextType _contextType;
-        private readonly long _startPosition;
-        private readonly long _size;
-        private string _currentElementName;
-        private int _currentArrayIndex = -1;
+    public ContextType ContextType { get; } = contextType;
+    public long StartPosition { get; } = startPosition;
+    public long Size { get; } = size;
 
-        // constructors
-        internal BsonBinaryReaderContext(
-            BsonBinaryReaderContext parentContext,
-            ContextType contextType,
-            long startPosition,
-            long size)
-        {
-            _parentContext = parentContext;
-            _contextType = contextType;
-            _startPosition = startPosition;
-            _size = size;
-        }
-
-        // public properties
-        public ContextType ContextType
-        {
-            get { return _contextType; }
-        }
-
-        public int CurrentArrayIndex
-        {
-            get { return _currentArrayIndex; }
-            set { _currentArrayIndex = value; }
-        }
-
-        public string CurrentElementName
-        {
-            get { return _currentElementName; }
-            set { _currentElementName = value; }
-        }
-
-        public BsonBinaryReaderContext ParentContext
-        {
-            get { return _parentContext; }
-        }
-
-        // public methods
-        /// <summary>
-        /// Creates a clone of the context.
-        /// </summary>
-        /// <returns>A clone of the context.</returns>
-        public BsonBinaryReaderContext Clone()
-        {
-            return new BsonBinaryReaderContext(_parentContext, _contextType, _startPosition, _size);
-        }
-
-        public BsonBinaryReaderContext PopContext(long position)
-        {
-            var actualSize = position - _startPosition;
-            if (actualSize != _size)
-            {
-                var message = string.Format("Expected size to be {0}, not {1}.", _size, actualSize);
-                throw new FormatException(message);
-            }
-            return _parentContext;
-        }
-    }
+    public int ArrayIndex { get; set; } = -1;
+    public string ElementName { get; set; }
 }