Poolable / Resettable CBOR Readers (#88104)

Co-authored-by: Adam Sitnik <adam.sitnik@gmail.com> Co-authored-by: Jeremy Barton <jbarton@microsoft.com>
dotnet · Jul 17, 2023 · d4db2d9 · d4db2d9
1 parent 0f56e16
commit d4db2d9
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 2 deletions.
diff --git a/src/libraries/System.Formats.Cbor/ref/System.Formats.Cbor.cs b/src/libraries/System.Formats.Cbor/ref/System.Formats.Cbor.cs
@@ -65,6 +65,7 @@ public void ReadStartIndefiniteLengthTextString() { }
         [System.CLSCompliantAttribute(false)]
         public ulong ReadUInt64() { throw null; }
         public System.DateTimeOffset ReadUnixTimeSeconds() { throw null; }
+        public void Reset(System.ReadOnlyMemory<byte> data) { }
         public void SkipToParent(bool disableConformanceModeChecks = false) { }
         public void SkipValue(bool disableConformanceModeChecks = false) { }
         public bool TryReadByteString(System.Span<byte> destination, out int bytesWritten) { throw null; }

diff --git a/src/libraries/System.Formats.Cbor/src/System/Formats/Cbor/Reader/CborReader.cs b/src/libraries/System.Formats.Cbor/src/System/Formats/Cbor/Reader/CborReader.cs
@@ -9,7 +9,7 @@ namespace System.Formats.Cbor
     /// <summary>A stateful, forward-only reader for Concise Binary Object Representation (CBOR) encoded data.</summary>
     public partial class CborReader
     {
-        private readonly ReadOnlyMemory<byte> _data;
+        private ReadOnlyMemory<byte> _data;
         private int _offset;
 
         private Stack<StackFrame>? _nestedDataItems;
@@ -60,7 +60,7 @@ public CborReader(ReadOnlyMemory<byte> data, CborConformanceMode conformanceMode
             _data = data;
             ConformanceMode = conformanceMode;
             AllowMultipleRootLevelValues = allowMultipleRootLevelValues;
-            _definiteLength = allowMultipleRootLevelValues ? null : (int?)1;
+            _definiteLength = allowMultipleRootLevelValues ? null : 1;
         }
 
         /// <summary>Reads the next CBOR data item, returning a <see cref="ReadOnlyMemory{T}" /> view of the encoded value. For indefinite length encodings this includes the break byte.</summary>
@@ -81,6 +81,34 @@ public ReadOnlyMemory<byte> ReadEncodedValue(bool disableConformanceModeChecks =
             return _data.Slice(initialOffset, _offset - initialOffset);
         }
 
+        /// <summary>
+        /// Resets the <see cref="CborReader"/> instance over the specified <paramref name="data"/> with unchanged configuration.
+        /// <see cref="ConformanceMode"/> and <see cref="AllowMultipleRootLevelValues"/> are unchanged.
+        /// </summary>
+        /// <param name="data">The CBOR-encoded data to read.</param>
+        public void Reset(ReadOnlyMemory<byte> data)
+        {
+            // ConformanceMode and AllowMultipleRootLevelValues are set in ctor, they remain unchanged.
+
+            _data = data;
+            _offset = 0;
+
+            _nestedDataItems?.Clear();
+            _currentMajorType = default;
+            _definiteLength = AllowMultipleRootLevelValues ? null : 1;
+            _itemsRead = default;
+            _frameOffset = default;
+            _isTagContext = default;
+            _currentKeyOffset = default;
+            _previousKeyEncodingRange = default;
+            _keyEncodingRanges?.Clear();
+            _isConformanceModeCheckEnabled = true;
+            _cachedState = CborReaderState.Undefined;
+
+            // We don't need to clear the reusable instances in _pooledKeyEncodingRangeAllocations
+            // or _indefiniteLengthStringRangeAllocation.
+        }
+
         private CborInitialByte PeekInitialByte()
         {
             if (_definiteLength - _itemsRead == 0)

diff --git a/src/libraries/System.Formats.Cbor/tests/Reader/CborReaderTests.cs b/src/libraries/System.Formats.Cbor/tests/Reader/CborReaderTests.cs
@@ -54,6 +54,80 @@ public static void Depth_ShouldReturnExpectedValue(int depth)
             Assert.Equal(CborReaderState.Finished, reader.PeekState());
         }
 
+        [Fact]
+        public static void Reset_DoesNotAffect_ConformanceMode()
+        {
+            var reader = new CborReader(Array.Empty<byte>());
+            Assert.Equal(CborConformanceMode.Strict, reader.ConformanceMode);
+
+            reader.Reset(Array.Empty<byte>());
+            Assert.Equal(CborConformanceMode.Strict, reader.ConformanceMode);
+
+            foreach(var mode in new[] { CborConformanceMode.Canonical, CborConformanceMode.Ctap2Canonical, CborConformanceMode.Lax, CborConformanceMode.Strict })
+            {
+                reader = new CborReader(Array.Empty<byte>(), mode);
+                Assert.Equal(mode, reader.ConformanceMode);
+
+                reader.Reset(Array.Empty<byte>());
+                Assert.Equal(mode, reader.ConformanceMode);
+            }
+        }
+
+        [Fact]
+        public static void Reset_DoesNotAffect_AllowMultipleRootLevelValues()
+        {
+            var reader = new CborReader(Array.Empty<byte>(),  allowMultipleRootLevelValues: false);
+            Assert.False(reader.AllowMultipleRootLevelValues);
+
+            reader.Reset(Array.Empty<byte>());
+            Assert.False(reader.AllowMultipleRootLevelValues);
+
+            reader = new CborReader(Array.Empty<byte>(), allowMultipleRootLevelValues: true);
+            Assert.True(reader.AllowMultipleRootLevelValues);
+
+            reader.Reset(Array.Empty<byte>());
+            Assert.True(reader.AllowMultipleRootLevelValues);
+        }
+
+        [Theory]
+        [InlineData(0)]
+        [InlineData(1)]
+        [InlineData(3)]
+        [InlineData(7)]
+        public static void ResetInDepth_ShouldReturnExpectedValue(int depth)
+        {
+            byte[] encoding = Enumerable.Repeat<byte>(0x81, depth).Append<byte>(0x01).ToArray();
+
+            var reader = new CborReader(encoding);
+            for (int i = 0; i < depth; i++)
+            {
+                Assert.Equal(i, reader.CurrentDepth);
+                reader.ReadStartArray();
+            }
+
+            reader.Reset(encoding);
+            Assert.Equal(0, reader.CurrentDepth);
+            Assert.Equal(encoding.Length, reader.BytesRemaining);
+
+            for (int i = 0; i < depth; i++)
+            {
+                Assert.Equal(i, reader.CurrentDepth);
+                reader.ReadStartArray();
+            }
+
+            Assert.Equal(depth, reader.CurrentDepth);
+            reader.ReadInt32();
+            Assert.Equal(depth, reader.CurrentDepth);
+
+            for (int i = depth - 1; i >= 0; i--)
+            {
+                reader.ReadEndArray();
+                Assert.Equal(i, reader.CurrentDepth);
+            }
+
+            Assert.Equal(CborReaderState.Finished, reader.PeekState());
+        }
+
         [Fact]
         public static void BytesRemaining_NoReads_ShouldReturnTotalLength()
         {