Improve JSON validation perf

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,8 +12,8 @@ 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"/>).

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

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

@@ -12,54 +12,116 @@ namespace System.Text.Json
 {
     public sealed partial class Utf8JsonWriter
     {
-        private bool HasPartialStringData => PartialStringDataLength != 0;
+        /// <summary>
+        /// Returns whether a JSON value can be written at the current position based on the current <see cref="_enclosingContainer"/>:
+        /// <list type="bullet">
+        /// <item>
+        /// <see cref="EnclosingContainerType.Array"/>: Writing a value is always allowed.
+        /// </item>
+        /// <item>
+        /// <see cref="EnclosingContainerType.Object"/>: Writing a value is allowed only if <see cref="_tokenType"/> is a property name.
+        /// Because we designed <see cref="EnclosingContainerType.Object"/> == <see cref="JsonTokenType.PropertyName"/>, we can just check for equality.
+        /// </item>
+        /// <item>
+        /// <see cref="EnclosingContainerType.None"/>: Writing a value is allowed only if <see cref="_tokenType"/> is None (only one value may be written at the root).
+        /// This case is identical to the previous case.
+        /// </item>
+        /// <item>
+        /// <see cref="EnclosingContainerType.Utf8StringSequence"/>, <see cref="EnclosingContainerType.Utf16StringSequence"/>, <see cref="EnclosingContainerType.Base64StringSequence"/>:
+        /// Writing a value is never valid and <see cref="_enclosingContainer"/> does not equal any <see cref="JsonTokenType"/> by construction.
+        /// </item>
+        /// </list>
+        /// This method performs better without short circuiting (this often gets inlined so using simple branch free code seems to have some benefits).
+        /// </summary>
+        private bool CanWriteValue => _enclosingContainer == EnclosingContainerType.Array | (byte)_enclosingContainer == (byte)_tokenType;
 
-        private void ClearPartialStringData() => PartialStringDataLength = 0;
+        private bool HasPartialStringData => _partialStringDataLength != 0;
 
-        private void ValidateEncodingDidNotChange(SegmentEncoding currentSegmentEncoding)
+        private void ClearPartialStringData() => _partialStringDataLength = 0;
+
+        private void ValidateWritingValue()
         {
-            if (PreviousSegmentEncoding != currentSegmentEncoding)
+            if (!CanWriteValue)
             {
-                ThrowHelper.ThrowInvalidOperationException_CannotMixEncodings(PreviousSegmentEncoding, currentSegmentEncoding);
+                OnValidateWritingValueFailed();
             }
         }
 
-        private void ValidateNotWithinUnfinalizedString()
+        [DoesNotReturn]
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private void OnValidateWritingValueFailed()
         {
-            if (_tokenType == StringSegmentSentinel)
+            Debug.Assert(!_options.SkipValidation);
+
+            if (IsWritingPartialString)
             {
-                ThrowHelper.ThrowInvalidOperationException(ExceptionResource.CannotWriteWithinString, currentDepth: default, maxDepth: _options.MaxDepth, token: default, _tokenType);
+                ThrowInvalidOperationException(ExceptionResource.CannotWriteWithinString);
             }
 
-            Debug.Assert(PreviousSegmentEncoding == SegmentEncoding.None);
             Debug.Assert(!HasPartialStringData);
+
+            if (_enclosingContainer == EnclosingContainerType.Object)
+            {
+                Debug.Assert(_tokenType != JsonTokenType.PropertyName);
+                Debug.Assert(_tokenType != JsonTokenType.None && _tokenType != JsonTokenType.StartArray);
+                ThrowInvalidOperationException(ExceptionResource.CannotWriteValueWithinObject);
+            }
+            else
+            {
+                Debug.Assert(_tokenType != JsonTokenType.PropertyName);
+                Debug.Assert(CurrentDepth == 0 && _tokenType != JsonTokenType.None);
+                ThrowInvalidOperationException(ExceptionResource.CannotWriteValueAfterPrimitiveOrClose);
+            }
         }
 
-        private void ValidateWritingValue()
+        private void ValidateWritingSegment(EnclosingContainerType currentSegmentEncoding)
         {
-            Debug.Assert(!_options.SkipValidation);
+            Debug.Assert(currentSegmentEncoding is EnclosingContainerType.Utf8StringSequence or EnclosingContainerType.Utf16StringSequence or EnclosingContainerType.Base64StringSequence);
+
+            // A string segment can be written if either:
+            // 1) The writer is currently in a partial string of the same type. In this case the new segment
+            // will continue the partial string.
+            // - or -
+            // 2) The writer can write a value at the current position, in which case a new string can be started.
+            if (_enclosingContainer != currentSegmentEncoding && !CanWriteValue)
+            {
+                OnValidateWritingSegmentFailed(currentSegmentEncoding);
+            }
+        }
+
+        [DoesNotReturn]
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private void OnValidateWritingSegmentFailed(EnclosingContainerType currentSegmentEncoding)
+        {
+            if (IsWritingPartialString)
+            {
+                ThrowHelper.ThrowInvalidOperationException_CannotMixEncodings(_enclosingContainer, currentSegmentEncoding);
+            }
 
-            // Make sure a new value is not attempted within an unfinalized string.
-            ValidateNotWithinUnfinalizedString();
+            Debug.Assert(!HasPartialStringData);
 
-            if (_inObject)
+            if (_enclosingContainer == EnclosingContainerType.Object)
             {
-                if (_tokenType != JsonTokenType.PropertyName)
-                {
-                    Debug.Assert(_tokenType != JsonTokenType.None && _tokenType != JsonTokenType.StartArray);
-                    ThrowHelper.ThrowInvalidOperationException(ExceptionResource.CannotWriteValueWithinObject, currentDepth: default, maxDepth: _options.MaxDepth, token: default, _tokenType);
-                }
+                Debug.Assert(_tokenType != JsonTokenType.PropertyName);
+                Debug.Assert(_tokenType != JsonTokenType.None && _tokenType != JsonTokenType.StartArray);
+                ThrowInvalidOperationException(ExceptionResource.CannotWriteValueWithinObject);
             }
             else
             {
                 Debug.Assert(_tokenType != JsonTokenType.PropertyName);
+                Debug.Assert(CurrentDepth == 0 && _tokenType != JsonTokenType.None);
+                ThrowInvalidOperationException(ExceptionResource.CannotWriteValueAfterPrimitiveOrClose);
+            }
+        }
 
-                // It is more likely for CurrentDepth to not equal 0 when writing valid JSON, so check that first to rely on short-circuiting and return quickly.
-                if (CurrentDepth == 0 && _tokenType != JsonTokenType.None)
-                {
-                    ThrowHelper.ThrowInvalidOperationException(ExceptionResource.CannotWriteValueAfterPrimitiveOrClose, currentDepth: default, maxDepth: _options.MaxDepth, token: default, _tokenType);
-                }
+        private void ValidateNotWithinUnfinalizedString()
+        {
+            if (IsWritingPartialString)
+            {
+                ThrowInvalidOperationException(ExceptionResource.CannotWriteWithinString);
             }
+
+            Debug.Assert(!HasPartialStringData);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]