Improve JSON validation perf (#111332)

Author: PranavSenthilnathan

src/libraries/System.Text.Json/src/System/Text/Json/BitStack.cs

@@ -14,6 +14,7 @@ internal struct BitStack
 
         private const int DefaultInitialArraySize = 2;
 
+        // The backing array for the stack used when the depth exceeds AllocationFreeMaxDepth.
         private int[]? _array;
 
         // This ulong container represents a tiny stack to track the state during nested transitions.
@@ -26,8 +27,14 @@ internal struct BitStack
 
         private int _currentDepth;
 
-        public int CurrentDepth => _currentDepth;
+        /// <summary>
+        /// Gets the number of elements in the stack.
+        /// </summary>
+        public readonly int CurrentDepth => _currentDepth;
 
+        /// <summary>
+        /// Pushes <see langword="true"/> onto the stack.
+        /// </summary>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void PushTrue()
         {
@@ -42,6 +49,9 @@ public void PushTrue()
             _currentDepth++;
         }
 
+        /// <summary>
+        /// Pushes <see langword="false"/> onto the stack.
+        /// </summary>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void PushFalse()
         {
@@ -56,7 +66,10 @@ public void PushFalse()
             _currentDepth++;
         }
 
-        // Allocate the bit array lazily only when it is absolutely necessary
+        /// <summary>
+        /// Pushes a bit onto the stack. Allocate the bit array lazily only when it is absolutely necessary.
+        /// </summary>
+        /// <param name="value">The bit to push onto the stack.</param>
         [MethodImpl(MethodImplOptions.NoInlining)]
         private void PushToArray(bool value)
         {
@@ -94,6 +107,10 @@ private void PushToArray(bool value)
             _array[elementIndex] = newValue;
         }
 
+        /// <summary>
+        /// Pops the bit at the top of the stack and returns its value.
+        /// </summary>
+        /// <returns>The bit that was popped.</returns>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public bool Pop()
         {
@@ -110,13 +127,18 @@ public bool Pop()
             }
             else
             {
-                inObject = PopFromArray();
+                // Decrementing depth above effectively pops the last element in the array-backed case.
+                inObject = PeekInArray();
             }
             return inObject;
         }
 
+        /// <summary>
+        /// If the stack has a backing array allocated, this method will find the topmost bit in the array and return its value.
+        /// This should only be called if the depth is greater than AllocationFreeMaxDepth and an array has been allocated.
+        /// </summary>
         [MethodImpl(MethodImplOptions.NoInlining)]
-        private bool PopFromArray()
+        private readonly bool PeekInArray()
         {
             int index = _currentDepth - AllocationFreeMaxDepth - 1;
             Debug.Assert(_array != null);
@@ -129,6 +151,14 @@ private bool PopFromArray()
             return (_array[elementIndex] & (1 << extraBits)) != 0;
         }
 
+        /// <summary>
+        /// Peeks at the bit at the top of the stack.
+        /// </summary>
+        /// <returns>The bit at the top of the stack.</returns>
+        public readonly bool Peek()
+            // If the stack is small enough, we can use the allocation-free container, otherwise check the allocated array.
+            => _currentDepth <= AllocationFreeMaxDepth ? (_allocationFreeContainer & 1) != 0 : PeekInArray();
+
         private void DoubleArray(int minSize)
         {
             Debug.Assert(_array != null);
@@ -141,13 +171,19 @@ private void DoubleArray(int minSize)
             Array.Resize(ref _array, nextDouble);
         }
 
+        /// <summary>
+        /// Optimization to push <see langword="true"/> as the first bit when the stack is empty.
+        /// </summary>
         public void SetFirstBit()
         {
             Debug.Assert(_currentDepth == 0, "Only call SetFirstBit when depth is 0");
             _currentDepth++;
             _allocationFreeContainer = 1;
         }
 
+        /// <summary>
+        /// Optimization to push <see langword="false"/> as the first bit when the stack is empty.
+        /// </summary>
         public void ResetFirstBit()
         {
             Debug.Assert(_currentDepth == 0, "Only call ResetFirstBit when depth is 0");

src/libraries/System.Text.Json/src/System/Text/Json/JsonTokenType.cs

@@ -12,70 +12,70 @@ namespace System.Text.Json
     /// </summary>
     public enum JsonTokenType : byte
     {
-        // Do not re-order.
-        // We rely on the ordering to quickly check things like IsTokenTypePrimitive
+        // Do not re-number.
+        // We rely on the underlying values to quickly check things like JsonReaderHelper.IsTokenTypePrimitive and Utf8JsonWriter.CanWriteValue
 
         /// <summary>
         ///   Indicates that there is no value (as distinct from <see cref="Null"/>).
         /// </summary>
         /// <remarks>
         ///   This is the default token type if no data has been read by the <see cref="Utf8JsonReader"/>.
         /// </remarks>
-        None,
+        None = 0,
 
         /// <summary>
         ///   Indicates that the token type is the start of a JSON object.
         /// </summary>
-        StartObject,
+        StartObject = 1,
 
         /// <summary>
         ///   Indicates that the token type is the end of a JSON object.
         /// </summary>
-        EndObject,
+        EndObject = 2,
 
         /// <summary>
         ///   Indicates that the token type is the start of a JSON array.
         /// </summary>
-        StartArray,
+        StartArray = 3,
 
         /// <summary>
         ///   Indicates that the token type is the end of a JSON array.
         /// </summary>
-        EndArray,
+        EndArray = 4,
 
         /// <summary>
         ///   Indicates that the token type is a JSON property name.
         /// </summary>
-        PropertyName,
+        PropertyName = 5,
 
         /// <summary>
         ///   Indicates that the token type is the comment string.
         /// </summary>
-        Comment,
+        Comment = 6,
 
         /// <summary>
         ///   Indicates that the token type is a JSON string.
         /// </summary>
-        String,
+        String = 7,
 
         /// <summary>
         ///   Indicates that the token type is a JSON number.
         /// </summary>
-        Number,
+        Number = 8,
 
         /// <summary>
         ///   Indicates that the token type is the JSON literal <c>true</c>.
         /// </summary>
-        True,
+        True = 9,
 
         /// <summary>
         ///   Indicates that the token type is the JSON literal <c>false</c>.
         /// </summary>
-        False,
+        False = 10,
 
         /// <summary>
         ///   Indicates that the token type is the JSON literal <c>null</c>.
         /// </summary>
-        Null,
+        Null = 11,
     }
 }

src/libraries/System.Text.Json/src/System/Text/Json/ThrowHelper.cs

@@ -4,6 +4,7 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
+using static System.Text.Json.Utf8JsonWriter;
 
 namespace System.Text.Json
 {
@@ -312,9 +313,23 @@ public static void ThrowInvalidOperationException_CannotSkipOnPartial()
         }
 
         [DoesNotReturn]
-        public static void ThrowInvalidOperationException_CannotMixEncodings(Utf8JsonWriter.SegmentEncoding previousEncoding, Utf8JsonWriter.SegmentEncoding currentEncoding)
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        public static void ThrowInvalidOperationException_CannotMixEncodings(EnclosingContainerType previousEncoding, EnclosingContainerType currentEncoding)
         {
-            throw GetInvalidOperationException(SR.Format(SR.CannotMixEncodings, previousEncoding, currentEncoding));
+            throw GetInvalidOperationException(SR.Format(SR.CannotMixEncodings, GetEncodingName(previousEncoding), GetEncodingName(currentEncoding)));
+
+            static string GetEncodingName(EnclosingContainerType encoding)
+            {
+                switch (encoding)
+                {
+                    case EnclosingContainerType.Utf8StringSequence: return "UTF-8";
+                    case EnclosingContainerType.Utf16StringSequence: return "UTF-16";
+                    case EnclosingContainerType.Base64StringSequence: return "Base64";
+                    default:
+                        Debug.Fail("Unknown encoding.");
+                        return "Unknown";
+                };
+            }
         }
 
         private static InvalidOperationException GetInvalidOperationException(string message, JsonTokenType tokenType)

src/libraries/System.Text.Json/src/System/Text/Json/Writer/Utf8JsonWriter.WriteProperties.Helpers.cs

@@ -3,6 +3,7 @@
 
 using System.Buffers;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 
@@ -36,13 +37,10 @@ private void ValidateWritingProperty()
         {
             if (!_options.SkipValidation)
             {
-                // Make sure a new property is not attempted within an unfinalized string.
-                ValidateNotWithinUnfinalizedString();
-
-                if (!_inObject || _tokenType == JsonTokenType.PropertyName)
+                if (_enclosingContainer != EnclosingContainerType.Object || _tokenType == JsonTokenType.PropertyName)
                 {
                     Debug.Assert(_tokenType != JsonTokenType.StartObject);
-                    ThrowHelper.ThrowInvalidOperationException(ExceptionResource.CannotWritePropertyWithinArray, currentDepth: default, maxDepth: _options.MaxDepth, token: default, _tokenType);
+                    OnValidateWritingPropertyFailed();
                 }
             }
         }
@@ -52,18 +50,28 @@ private void ValidateWritingProperty(byte token)
         {
             if (!_options.SkipValidation)
             {
-                // Make sure a new property is not attempted within an unfinalized string.
-                ValidateNotWithinUnfinalizedString();
-
-                if (!_inObject || _tokenType == JsonTokenType.PropertyName)
+                if (_enclosingContainer != EnclosingContainerType.Object || _tokenType == JsonTokenType.PropertyName)
                 {
                     Debug.Assert(_tokenType != JsonTokenType.StartObject);
-                    ThrowHelper.ThrowInvalidOperationException(ExceptionResource.CannotWritePropertyWithinArray, currentDepth: default, maxDepth: _options.MaxDepth, token: default, _tokenType);
+                    OnValidateWritingPropertyFailed();
                 }
                 UpdateBitStackOnStart(token);
             }
         }
 
+        [DoesNotReturn]
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private void OnValidateWritingPropertyFailed()
+        {
+            if (IsWritingPartialString)
+            {
+                ThrowInvalidOperationException(ExceptionResource.CannotWriteWithinString);
+            }
+
+            Debug.Assert(_enclosingContainer != EnclosingContainerType.Object || _tokenType == JsonTokenType.PropertyName);
+            ThrowInvalidOperationException(ExceptionResource.CannotWritePropertyWithinArray);
+        }
+
         private void WritePropertyNameMinimized(ReadOnlySpan<byte> escapedPropertyName, byte token)
         {
             Debug.Assert(escapedPropertyName.Length < int.MaxValue - 5);