CSHARP-5666: Remove GetBitArray allocations in BsonClassMapSerializer.DeserializeClass

Author: BorisDog

src/MongoDB.Bson/Serialization/Serializers/BsonClassMapSerializer.cs

@@ -14,9 +14,11 @@
 */
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.ComponentModel;
 using System.Reflection;
+using System.Runtime.CompilerServices;
 using MongoDB.Bson.IO;
 using MongoDB.Bson.Serialization.Conventions;
 using MongoDB.Bson.Serialization.Serializers;
@@ -82,7 +84,7 @@ public override TClass Deserialize(BsonDeserializationContext context, BsonDeser
         {
             var bsonReader = context.Reader;
 
-            if (bsonReader.GetCurrentBsonType() == Bson.BsonType.Null)
+            if (bsonReader.GetCurrentBsonType() == BsonType.Null)
             {
                 bsonReader.ReadNull();
                 return default(TClass);
@@ -149,7 +151,9 @@ public TClass DeserializeClass(BsonDeserializationContext context)
             var discriminatorConvention = _classMap.GetDiscriminatorConvention();
             var allMemberMaps = _classMap.AllMemberMaps;
             var extraElementsMemberMapIndex = _classMap.ExtraElementsMemberMapIndex;
-            var memberMapBitArray = FastMemberMapHelper.GetBitArray(allMemberMaps.Count);
+
+            var (bitArrayLength, useStackAlloc) = FastMemberMapHelper.GetBitArrayLength(_classMap.AllMemberMaps.Count);
+            using var bitArray = useStackAlloc ? FastMemberMapHelper.GetBitArray(stackalloc uint[bitArrayLength]) : FastMemberMapHelper.GetBitArray(bitArrayLength);
 
             bsonReader.ReadStartDocument();
             var elementTrie = _classMap.ElementTrie;
@@ -193,7 +197,8 @@ public TClass DeserializeClass(BsonDeserializationContext context)
                             DeserializeExtraElementValue(context, values, elementName, memberMap);
                         }
                     }
-                    memberMapBitArray[memberMapIndex >> 5] |= 1U << (memberMapIndex & 31);
+
+                    bitArray.SetMemberIndex(memberMapIndex);
                 }
                 else
                 {
@@ -221,7 +226,7 @@ public TClass DeserializeClass(BsonDeserializationContext context)
                         {
                             DeserializeExtraElementValue(context, values, elementName, extraElementsMemberMap);
                         }
-                        memberMapBitArray[extraElementsMemberMapIndex >> 5] |= 1U << (extraElementsMemberMapIndex & 31);
+                        bitArray.SetMemberIndex(extraElementsMemberMapIndex);
                     }
                     else if (_classMap.IgnoreExtraElements)
                     {
@@ -239,51 +244,38 @@ public TClass DeserializeClass(BsonDeserializationContext context)
             bsonReader.ReadEndDocument();
 
             // check any members left over that we didn't have elements for (in blocks of 32 elements at a time)
-            for (var bitArrayIndex = 0; bitArrayIndex < memberMapBitArray.Length; ++bitArrayIndex)
+            var bitArraySpan = bitArray.Span;
+            for (var bitArrayIndex = 0; bitArrayIndex < bitArraySpan.Length; bitArrayIndex++)
             {
                 var memberMapIndex = bitArrayIndex << 5;
-                var memberMapBlock = ~memberMapBitArray[bitArrayIndex]; // notice that bits are flipped so 1's are now the missing elements
+                var memberMapBlock = ~bitArraySpan[bitArrayIndex]; // notice that bits are flipped so 1's are now the missing elements
 
                 // work through this memberMapBlock of 32 elements
-                while (true)
+                for (; memberMapBlock != 0 && memberMapIndex < allMemberMaps.Count; memberMapIndex++, memberMapBlock >>= 1)
                 {
-                    // examine missing elements (memberMapBlock is shifted right as we work through the block)
-                    for (; (memberMapBlock & 1) != 0; ++memberMapIndex, memberMapBlock >>= 1)
-                    {
-                        var memberMap = allMemberMaps[memberMapIndex];
-                        if (memberMap.IsReadOnly)
-                        {
-                            continue;
-                        }
-
-                        if (memberMap.IsRequired)
-                        {
-                            var fieldOrProperty = (memberMap.MemberInfo is FieldInfo) ? "field" : "property";
-                            var message = string.Format(
-                                "Required element '{0}' for {1} '{2}' of class {3} is missing.",
-                                memberMap.ElementName, fieldOrProperty, memberMap.MemberName, _classMap.ClassType.FullName);
-                            throw new FormatException(message);
-                        }
+                    if ((memberMapBlock & 1) == 0)
+                        continue;
 
-                        if (document != null)
-                        {
-                            memberMap.ApplyDefaultValue(document);
-                        }
-                        else if (memberMap.IsDefaultValueSpecified && !memberMap.IsReadOnly)
-                        {
-                            values[memberMap.ElementName] = memberMap.DefaultValue;
-                        }
+                    var memberMap = allMemberMaps[memberMapIndex];
+                    if (memberMap.IsReadOnly)
+                    {
+                        continue;
                     }
 
-                    if (memberMapBlock == 0)
+                    if (memberMap.IsRequired)
                     {
-                        break;
+                        var fieldOrProperty = (memberMap.MemberInfo is FieldInfo) ? "field" : "property";
+                        throw new FormatException($"Required element 'memberMap.ElementName' for {fieldOrProperty} '{memberMap.MemberName}' of class {_classMap.ClassType.FullName} is missing.");
                     }
 
-                    // skip ahead to the next missing element
-                    var leastSignificantBit = FastMemberMapHelper.GetLeastSignificantBit(memberMapBlock);
-                    memberMapIndex += leastSignificantBit;
-                    memberMapBlock >>= leastSignificantBit;
+                    if (document != null)
+                    {
+                        memberMap.ApplyDefaultValue(document);
+                    }
+                    else if (memberMap.IsDefaultValueSpecified && !memberMap.IsReadOnly)
+                    {
+                        values[memberMap.ElementName] = memberMap.DefaultValue;
+                    }
                 }
             }
 
@@ -335,13 +327,11 @@ public bool GetDocumentId(
                 idGenerator = idMemberMap.IdGenerator;
                 return true;
             }
-            else
-            {
-                id = null;
-                idNominalType = null;
-                idGenerator = null;
-                return false;
-            }
+
+            id = null;
+            idNominalType = null;
+            idGenerator = null;
+            return false;
         }
 
         /// <summary>
@@ -694,47 +684,73 @@ private bool ShouldSerializeDiscriminator(Type nominalType)
 
         // nested classes
         // helper class that implements member map bit array helper functions
-        private static class FastMemberMapHelper
+        internal static class FastMemberMapHelper
         {
-            public static uint[] GetBitArray(int memberCount)
+            internal ref struct BitArray()
             {
-                var bitArrayOffset = memberCount & 31;
-                var bitArrayLength = memberCount >> 5;
-                if (bitArrayOffset == 0)
-                {
-                    return new uint[bitArrayLength];
-                }
-                var bitArray = new uint[bitArrayLength + 1];
-                bitArray[bitArrayLength] = ~0U << bitArrayOffset; // set unused bits to 1
-                return bitArray;
-            }
+                private readonly ArrayPool<uint> _arrayPool;
+                private readonly Span<uint> _bitArray;
+                private readonly uint[] _rentedBuffer;
+                private bool _isDisposed = false;
 
-            // see http://graphics.stanford.edu/~seander/bithacks.html#ZerosOnRightBinSearch
-            // also returns 31 if no bits are set; caller must check this case
-            public static int GetLeastSignificantBit(uint bitBlock)
-            {
-                var leastSignificantBit = 1;
-                if ((bitBlock & 65535) == 0)
+                public BitArray(Span<uint> bitArray) : this()
                 {
-                    bitBlock >>= 16;
-                    leastSignificantBit |= 16;
+                    _arrayPool = null;
+                    _bitArray = bitArray;
+                    _rentedBuffer = null;
                 }
-                if ((bitBlock & 255) == 0)
+
+                public BitArray(int spanLength, uint[] rentedBuffer, ArrayPool<uint> arrayPool) : this()
                 {
-                    bitBlock >>= 8;
-                    leastSignificantBit |= 8;
+                    _arrayPool = arrayPool;
+                    _bitArray = rentedBuffer.AsSpan(0, spanLength);
+                    _rentedBuffer = rentedBuffer;
                 }
-                if ((bitBlock & 15) == 0)
+
+                public Span<uint> Span => _bitArray;
+                public ArrayPool<uint> ArrayPool => _arrayPool;
+
+                public void SetMemberIndex(int memberMapIndex) =>
+                    _bitArray[memberMapIndex >> 5] |= 1U << (memberMapIndex & 31);
+
+                public void Dispose()
                 {
-                    bitBlock >>= 4;
-                    leastSignificantBit |= 4;
+                    if (_isDisposed)
+                        return;
+
+                    if (_rentedBuffer != null)
+                    {
+                        _arrayPool.Return(_rentedBuffer);
+                    }
+                    _isDisposed = true;
                 }
-                if ((bitBlock & 3) == 0)
+            }
+
+            public static (int BitArrayLength, bool UseStackAlloc) GetBitArrayLength(int membersCount)
+            {
+                var length = (membersCount + 31) >> 5;
+                return (length, length <= 8); // Use stackalloc for up to 256 members
+            }
+
+            public static BitArray GetBitArray(Span<uint> span) =>
+                new(ResetSpan(span));
+
+            public static BitArray GetBitArray(int length)
+            {
+                var rentedBuffer = ArrayPool<uint>.Shared.Rent(length);
+                ResetSpan(rentedBuffer);
+
+                return new(length, rentedBuffer, ArrayPool<uint>.Shared);
+            }
+
+            private static Span<uint> ResetSpan(Span<uint> span)
+            {
+                for (var i = 0; i < span.Length; i++)
                 {
-                    bitBlock >>= 2;
-                    leastSignificantBit |= 2;
+                    span[i] = 0;
                 }
-                return leastSignificantBit - (int)(bitBlock & 1);
+
+                return span;
             }
         }
     }