Bug fix and small perf improvement to MJD WriteTo() (#35744)

* Bug fix and small perf improvement to MJD WriteTo() * more perf * pr fb
Azure · Apr 24, 2023 · b6c723c · b6c723c
1 parent 42e0234
commit b6c723c
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 5 deletions.
diff --git a/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.net461.cs b/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.net461.cs
@@ -164,8 +164,10 @@ internal MutableJsonDocument() { }
         public Azure.Core.Json.MutableJsonElement RootElement { get { throw null; } }
         public void Dispose() { }
         public static Azure.Core.Json.MutableJsonDocument Parse(System.BinaryData utf8Json) { throw null; }
+        public static Azure.Core.Json.MutableJsonDocument Parse(System.ReadOnlyMemory<byte> utf8Json) { throw null; }
         public static Azure.Core.Json.MutableJsonDocument Parse(string json) { throw null; }
         public void WriteTo(System.IO.Stream stream, System.Buffers.StandardFormat format = default(System.Buffers.StandardFormat)) { }
+        public void WriteTo(System.Text.Json.Utf8JsonWriter writer) { }
     }
     [System.Diagnostics.DebuggerDisplayAttribute("{DebuggerDisplay,nq}")]
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]

diff --git a/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.net6.0.cs b/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.net6.0.cs
@@ -164,8 +164,10 @@ internal MutableJsonDocument() { }
         public Azure.Core.Json.MutableJsonElement RootElement { get { throw null; } }
         public void Dispose() { }
         public static Azure.Core.Json.MutableJsonDocument Parse(System.BinaryData utf8Json) { throw null; }
+        public static Azure.Core.Json.MutableJsonDocument Parse(System.ReadOnlyMemory<byte> utf8Json) { throw null; }
         public static Azure.Core.Json.MutableJsonDocument Parse(string json) { throw null; }
         public void WriteTo(System.IO.Stream stream, System.Buffers.StandardFormat format = default(System.Buffers.StandardFormat)) { }
+        public void WriteTo(System.Text.Json.Utf8JsonWriter writer) { }
     }
     [System.Diagnostics.DebuggerDisplayAttribute("{DebuggerDisplay,nq}")]
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]

diff --git a/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.netstandard2.0.cs b/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.netstandard2.0.cs
@@ -164,8 +164,10 @@ internal MutableJsonDocument() { }
         public Azure.Core.Json.MutableJsonElement RootElement { get { throw null; } }
         public void Dispose() { }
         public static Azure.Core.Json.MutableJsonDocument Parse(System.BinaryData utf8Json) { throw null; }
+        public static Azure.Core.Json.MutableJsonDocument Parse(System.ReadOnlyMemory<byte> utf8Json) { throw null; }
         public static Azure.Core.Json.MutableJsonDocument Parse(string json) { throw null; }
         public void WriteTo(System.IO.Stream stream, System.Buffers.StandardFormat format = default(System.Buffers.StandardFormat)) { }
+        public void WriteTo(System.Text.Json.Utf8JsonWriter writer) { }
     }
     [System.Diagnostics.DebuggerDisplayAttribute("{DebuggerDisplay,nq}")]
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]

diff --git a/sdk/core/Azure.Core.Experimental/src/MutableJsonDocument.cs b/sdk/core/Azure.Core.Experimental/src/MutableJsonDocument.cs
@@ -18,7 +18,7 @@ public sealed partial class MutableJsonDocument : IDisposable
     {
         internal static readonly JsonSerializerOptions DefaultJsonSerializerOptions = new JsonSerializerOptions();
 
-        private readonly Memory<byte> _original;
+        private readonly ReadOnlyMemory<byte> _original;
         private readonly JsonDocument _originalDocument;
 
         internal ChangeTracker Changes { get; } = new();
@@ -47,11 +47,14 @@ public MutableJsonElement RootElement
         /// </summary>
         /// <param name="stream">The stream to which to write the document.</param>
         /// <param name="format">A format string indicating the format to use when writing the document.</param>
+        /// <exception cref="ArgumentNullException">The <paramref name="stream"/> parameter is <see langword="null"/>.</exception>
         /// <exception cref="FormatException">Thrown if an unsupported value is passed for format.</exception>
         /// <remarks>The value of <paramref name="format"/> can be default or 'J' to write the document as JSON.</remarks>
         public void WriteTo(Stream stream, StandardFormat format = default)
         {
-            if (format != default || format.Symbol != 'J')
+            Argument.AssertNotNull(stream, nameof(stream));
+
+            if (format != default && format.Symbol != 'J')
             {
                 throw new FormatException($"Unsupported format {format.Symbol}. Supported formats are: 'J' - JSON.");
             }
@@ -66,8 +69,17 @@ public void WriteTo(Stream stream, StandardFormat format = default)
             RootElement.WriteTo(writer);
         }
 
-        internal void WriteTo(Utf8JsonWriter writer)
+        /// <summary>
+        /// Writes the document to the provided stream as a JSON value.
+        /// </summary>
+        /// <param name="writer">The writer to which to write the document.</param>
+        /// <exception cref="ArgumentNullException">The <paramref name="writer"/> parameter is <see langword="null"/>.</exception>
+#pragma warning disable AZC0014 // Avoid using banned types in public API
+        public void WriteTo(Utf8JsonWriter writer)
+#pragma warning restore AZC0014 // Avoid using banned types in public API
         {
+            Argument.AssertNotNull(writer, nameof(writer));
+
             if (!Changes.HasChanges)
             {
                 _originalDocument.RootElement.WriteTo(writer);
@@ -96,17 +108,31 @@ private static void Write(Stream stream, ReadOnlySpan<byte> buffer)
         /// </summary>
         /// <param name="utf8Json">A UTF-8 encoded string representing a JSON value.</param>
         /// <returns>A <see cref="MutableJsonDocument"/> representation of the value.</returns>
+        /// <exception cref="JsonException"><paramref name="utf8Json"/> does not represent a valid single JSON value.</exception>
+        public static MutableJsonDocument Parse(ReadOnlyMemory<byte> utf8Json)
+        {
+            var doc = JsonDocument.Parse(utf8Json);
+            return new MutableJsonDocument(doc, utf8Json);
+        }
+
+        /// <summary>
+        /// Parses a UTF-8 encoded string representing a single JSON value into a <see cref="MutableJsonDocument"/>.
+        /// </summary>
+        /// <param name="utf8Json">A UTF-8 encoded string representing a JSON value.</param>
+        /// <returns>A <see cref="MutableJsonDocument"/> representation of the value.</returns>
+        /// <exception cref="JsonException"><paramref name="utf8Json"/> does not represent a valid single JSON value.</exception>
         public static MutableJsonDocument Parse(BinaryData utf8Json)
         {
             var doc = JsonDocument.Parse(utf8Json);
-            return new MutableJsonDocument(doc, utf8Json.ToArray().AsMemory());
+            return new MutableJsonDocument(doc, utf8Json.ToMemory());
         }
 
         /// <summary>
         /// Parses test representing a single JSON value into a <see cref="MutableJsonDocument"/>.
         /// </summary>
         /// <param name="json">The JSON string.</param>
         /// <returns>A <see cref="MutableJsonDocument"/> representation of the value.</returns>
+        /// <exception cref="JsonException"><paramref name="json"/> does not represent a valid single JSON value.</exception>
         public static MutableJsonDocument Parse(string json)
         {
             byte[] utf8 = Encoding.UTF8.GetBytes(json);
@@ -120,7 +146,7 @@ public void Dispose()
             _originalDocument.Dispose();
         }
 
-        internal MutableJsonDocument(JsonDocument jsonDocument, Memory<byte> utf8Json)
+        internal MutableJsonDocument(JsonDocument jsonDocument, ReadOnlyMemory<byte> utf8Json)
         {
             _original = utf8Json;
             _originalDocument = jsonDocument;

diff --git a/sdk/core/Azure.Core.Experimental/src/MutableJsonElement.WriteTo.cs b/sdk/core/Azure.Core.Experimental/src/MutableJsonElement.WriteTo.cs
@@ -17,6 +17,32 @@ private void WriteElement(string path, int highWaterMark, JsonElement element, U
         {
             if (Changes.TryGetChange(path, highWaterMark, out MutableJsonChange change))
             {
+                switch (change.Value)
+                {
+                    case int i:
+                        writer.WriteNumberValue(i);
+                        return;
+                    case long l:
+                        writer.WriteNumberValue(l);
+                        return;
+                    case double d:
+                        writer.WriteNumberValue(d);
+                        return;
+                    case float f:
+                        writer.WriteNumberValue(f);
+                        return;
+                    case bool b:
+                        writer.WriteBooleanValue(b);
+                        return;
+                    case null:
+                        writer.WriteNullValue();
+                        return;
+                    default:
+                        break;
+
+                        // Note: string is not included to let JsonElement handle escaping.
+                }
+
                 element = change.AsJsonElement();
                 highWaterMark = change.Index;
             }

diff --git a/sdk/core/Azure.Core.Experimental/tests/MutableJsonDocumentWriteToTests.cs b/sdk/core/Azure.Core.Experimental/tests/MutableJsonDocumentWriteToTests.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Linq;
 using System.Text.Json;
 using Azure.Core.Json;
@@ -464,6 +465,48 @@ public void CanWriteArrayWithChangedElements()
             MutableJsonDocumentTests.ValidateWriteTo(expected, mdoc);
         }
 
+        [Test]
+        public void CanWriteToStream()
+        {
+            string json = """{ "foo" : 1 }""";
+            MutableJsonDocument mdoc = MutableJsonDocument.Parse(json);
+
+            using Stream stream = new MemoryStream();
+            mdoc.WriteTo(stream);
+            stream.Flush();
+            stream.Position = 0;
+
+            string actual = BinaryData.FromStream(stream).ToString();
+
+            Assert.AreEqual(json, actual);
+        }
+
+        [Test]
+        public void CanWriteToStreamWithJsonFormat()
+        {
+            string json = """{ "foo" : 1 }""";
+            MutableJsonDocument mdoc = MutableJsonDocument.Parse(json);
+
+            using Stream stream = new MemoryStream();
+            mdoc.WriteTo(stream, 'J');
+            stream.Flush();
+            stream.Position = 0;
+
+            string actual = BinaryData.FromStream(stream).ToString();
+
+            Assert.AreEqual(json, actual);
+        }
+
+        [Test]
+        public void CannotWriteToStreamWithPatchFormat()
+        {
+            string json = """{ "foo" : 1 }""";
+            MutableJsonDocument mdoc = MutableJsonDocument.Parse(json);
+
+            Stream stream = new MemoryStream();
+            Assert.Throws<FormatException>(() => mdoc.WriteTo(stream, 'P'));
+        }
+
         [TestCaseSource(nameof(TestCases))]
         public void WriteToBehaviorMatchesJsonDocument(dynamic json)
         {