CSHARP-5348: Avoid allocations for Bson*Context (#1791)

Author: BorisDog

src/MongoDB.Bson/IO/BsonBinaryReader.cs

@@ -14,6 +14,7 @@
 */
 
 using System;
+using System.Collections.Generic;
 using System.Globalization;
 using System.IO;
 
@@ -30,6 +31,7 @@ public class BsonBinaryReader : BsonReader
 #pragma warning restore CA2213 // Disposable never disposed
         private readonly BsonStream _bsonStream;
         private BsonBinaryReaderContext _context;
+        private readonly Stack<BsonBinaryReaderContext> _contextStack = new(4);
 
         // constructors
         /// <summary>
@@ -61,7 +63,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 +111,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, _contextStack, _bsonStream.Position);
 
         /// <summary>
         /// Determines whether this reader is at end of file.
@@ -201,12 +201,11 @@ public override BsonType ReadBsonType()
 
             if (_context.ContextType == ContextType.Array)
             {
-                _context.CurrentArrayIndex++;
+                _context.ArrayIndex++;
             }
 
             try
             {
-
                 CurrentBsonType = _bsonStream.ReadBsonType();
             }
             catch (FormatException ex)
@@ -342,7 +341,8 @@ public override void ReadEndArray()
                 ThrowInvalidState("ReadEndArray", BsonReaderState.EndOfArray);
             }
 
-            _context = _context.PopContext(_bsonStream.Position);
+            PopContext();
+
             switch (_context.ContextType)
             {
                 case ContextType.Array: State = BsonReaderState.Type; break;
@@ -371,10 +371,10 @@ public override void ReadEndDocument()
                 ThrowInvalidState("ReadEndDocument", BsonReaderState.EndOfDocument);
             }
 
-            _context = _context.PopContext(_bsonStream.Position);
+            PopContext();
             if (_context.ContextType == ContextType.JavaScriptWithScope)
             {
-                _context = _context.PopContext(_bsonStream.Position); // JavaScriptWithScope
+                PopContext(); // JavaScriptWithScope
             }
             switch (_context.ContextType)
             {
@@ -432,7 +432,9 @@ public override string ReadJavaScriptWithScope()
 
             var startPosition = _bsonStream.Position; // position of size field
             var size = ReadSize();
-            _context = new BsonBinaryReaderContext(_context, ContextType.JavaScriptWithScope, startPosition, size);
+
+            PushContext(new(ContextType.JavaScriptWithScope, startPosition, size));
+
             var code = _bsonStream.ReadString(Settings.Encoding);
 
             State = BsonReaderState.ScopeDocument;
@@ -486,7 +488,7 @@ public override string ReadName(INameDecoder nameDecoder)
 
             if (_context.ContextType == ContextType.Document)
             {
-                _context.CurrentElementName = CurrentName;
+                _context.ElementName = CurrentName;
             }
 
             return CurrentName;
@@ -553,7 +555,7 @@ public override IByteBuffer ReadRawBsonDocument()
 
             if (_context.ContextType == ContextType.JavaScriptWithScope)
             {
-                _context = _context.PopContext(_bsonStream.Position); // JavaScriptWithScope
+                PopContext(); // JavaScriptWithScope
             }
             switch (_context.ContextType)
             {
@@ -590,7 +592,9 @@ public override void ReadStartArray()
 
             var startPosition = _bsonStream.Position; // position of size field
             var size = ReadSize();
-            _context = new BsonBinaryReaderContext(_context, ContextType.Array, startPosition, size);
+
+            PushContext(new(ContextType.Array, startPosition, size));
+
             State = BsonReaderState.Type;
         }
 
@@ -605,7 +609,9 @@ 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(new(contextType, startPosition, size));
+
             State = BsonReaderState.Type;
         }
 
@@ -662,10 +668,11 @@ public override void ReadUndefined()
         public override void ReturnToBookmark(BsonReaderBookmark bookmark)
         {
             var binaryReaderBookmark = (BsonBinaryReaderBookmark)bookmark;
+
             State = binaryReaderBookmark.State;
             CurrentBsonType = binaryReaderBookmark.CurrentBsonType;
             CurrentName = binaryReaderBookmark.CurrentName;
-            _context = binaryReaderBookmark.CloneContext();
+            _context = binaryReaderBookmark.RestoreContext(_contextStack);
             _bsonStream.Position = binaryReaderBookmark.Position;
         }
 
@@ -686,7 +693,7 @@ public override void SkipName()
 
             if (_context.ContextType == ContextType.Document)
             {
-                _context.CurrentElementName = CurrentName;
+                _context.ElementName = CurrentName;
             }
         }
 
@@ -745,7 +752,7 @@ protected override void Dispose(bool disposing)
                 {
                     Close();
                 }
-                catch { } // ignore exceptions
+                catch { /* ignore exceptions */ }
             }
             base.Dispose(disposing);
         }
@@ -767,35 +774,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(_contextStack.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);
-            }
-            else if (context.ContextType == ContextType.Array)
-            {
-                var indexElementName = context.CurrentArrayIndex.ToString(NumberFormatInfo.InvariantInfo);
-                return GenerateDottedElementName(context.ParentContext, indexElementName + "." + elementName);
+                return GenerateDottedElementName(contexts,  nextIndex, (context.ElementName ?? "?") + "." + 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()
@@ -813,6 +826,23 @@ private BsonReaderState GetNextState()
             }
         }
 
+        private void PopContext()
+        {
+            var actualSize = _bsonStream.Position - _context.StartPosition;
+            if (actualSize != _context.Size)
+            {
+                throw new FormatException($"Expected size to be {_context.Size}, not {actualSize}.");
+            }
+
+            _context =_contextStack.Pop();
+        }
+
+        private void PushContext(BsonBinaryReaderContext newContext)
+        {
+            _contextStack.Push(_context);
+            _context = newContext;
+        }
+
         private int ReadSize()
         {
             int size = _bsonStream.ReadInt32();

src/MongoDB.Bson/IO/BsonBinaryReaderBookmark.cs

@@ -13,40 +13,43 @@
 * limitations under the License.
 */
 
-namespace MongoDB.Bson.IO
+using System.Collections.Generic;
+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
+    private readonly BsonBinaryReaderContext _context;
+    private readonly BsonBinaryReaderContext[] _contextArray;
+    private readonly long _position;
+
+    internal BsonBinaryReaderBookmark(
+        BsonReaderState state,
+        BsonType currentBsonType,
+        string currentName,
+        BsonBinaryReaderContext currentContext,
+        Stack<BsonBinaryReaderContext> contextStack,
+        long position)
+        : base(state, currentBsonType, currentName)
     {
-        // private fields
-        private BsonBinaryReaderContext _context;
-        private long _position;
+        _context = currentContext;
+        _contextArray = contextStack.ToArray();
+        _position = position;
+    }
 
-        // constructors
-        internal BsonBinaryReaderBookmark(
-            BsonReaderState state,
-            BsonType currentBsonType,
-            string currentName,
-            BsonBinaryReaderContext context,
-            long position)
-            : base(state, currentBsonType, currentName)
-        {
-            _context = context.Clone();
-            _position = position;
-        }
+    internal long Position => _position;
 
-        // internal properties
-        internal long Position
-        {
-            get { return _position; }
-        }
+    internal BsonBinaryReaderContext RestoreContext(Stack<BsonBinaryReaderContext> contextStack)
+    {
+        contextStack.Clear();
 
-        // internal methods
-        internal BsonBinaryReaderContext CloneContext()
+        for (var i = _contextArray.Length - 1; i >= 0; i--)
         {
-            return _context.Clone();
+            contextStack.Push(_contextArray[i]);
         }
+
+        return _context;
     }
 }

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; }
 }