Avoid per-element patch-path allocations in C# collection serializers

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator.ClientModel/test/ModelReaderWriterValidation/TestProjects/Sample_TypeSpec/DynamicModelTests.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Buffers;
 using System.ClientModel.Primitives;
 using System.Collections.Generic;
 using System.Linq;
@@ -369,6 +370,292 @@ public void JsonPatchRemove_NullDynamicListElementSnapshot(bool onlyNull)
             Assert.That(Encoding.UTF8.GetString(json), Is.EqualTo(onlyNull ? "[]" : """[{"bar":"present"},null]"""));
         }
 
+        [TestCase(false)]
+        [TestCase(true)]
+        public void JsonModelWrite_UnpatchedCollectionDoesNotAllocatePerElement(bool unrelatedPatch)
+        {
+            var model = new NullableDynamicModel
+            {
+                Children = new AnotherDynamicModel[256]
+            };
+#pragma warning disable SCME0001
+            if (unrelatedPatch)
+            {
+                model.Patch.Set("$.unrelated"u8, 1);
+            }
+#pragma warning restore SCME0001
+
+            var buffer = new ArrayBufferWriter<byte>();
+            using var writer = new Utf8JsonWriter(buffer);
+            var jsonModel = (IJsonModel<NullableDynamicModel>)model;
+            jsonModel.Write(writer, ModelReaderWriterOptions.Json);
+            writer.Flush();
+            buffer.Clear();
+            writer.Reset(buffer);
+
+            long before = GC.GetAllocatedBytesForCurrentThread();
+            jsonModel.Write(writer, ModelReaderWriterOptions.Json);
+            writer.Flush();
+            long allocated = GC.GetAllocatedBytesForCurrentThread() - before;
+
+            Assert.That(allocated, Is.LessThan(1024), "Indexed patch paths must not allocate for each unpatched element.");
+            using var document = JsonDocument.Parse(buffer.WrittenMemory);
+            Assert.That(document.RootElement.GetProperty("children").GetArrayLength(), Is.EqualTo(256));
+        }
+
+        [TestCase("nestedChildren", false)]
+        [TestCase("nestedChildren", true)]
+        [TestCase("nestedChildDictionary", false)]
+        [TestCase("nestedChildDictionary", true)]
+        [TestCase("dictionaryChildren", false)]
+        [TestCase("dictionaryChildren", true)]
+        [TestCase("listOfDictionaries", false)]
+        [TestCase("listOfDictionaries", true)]
+        public void JsonModelWrite_UnpatchedNestedCollectionDoesNotAllocatePerElement(string propertyName, bool unrelatedPatch)
+        {
+            const int Count = 256;
+            var model = new NullableDynamicModel();
+            switch (propertyName)
+            {
+                case "nestedChildren":
+                    model.NestedChildren = Enumerable.Range(0, Count)
+                        .Select(_ => (IList<AnotherDynamicModel>)[new AnotherDynamicModel("value")])
+                        .ToList();
+                    break;
+                case "nestedChildDictionary":
+                    model.NestedChildDictionary = Enumerable.Range(0, Count)
+                        .ToDictionary(
+                            index => index.ToString(),
+                            _ => (IDictionary<string, AnotherDynamicModel>)new Dictionary<string, AnotherDynamicModel>
+                            {
+                                ["value"] = new AnotherDynamicModel("value")
+                            });
+                    break;
+                case "dictionaryChildren":
+                    model.DictionaryChildren = Enumerable.Range(0, Count)
+                        .ToDictionary(
+                            index => index.ToString(),
+                            _ => (IList<AnotherDynamicModel>)[new AnotherDynamicModel("value")]);
+                    break;
+                case "listOfDictionaries":
+                    model.ListOfDictionaries = Enumerable.Range(0, Count)
+                        .Select(_ => (IDictionary<string, AnotherDynamicModel>)new Dictionary<string, AnotherDynamicModel>
+                        {
+                            ["value"] = new AnotherDynamicModel("value")
+                        })
+                        .ToList();
+                    break;
+                default:
+                    throw new ArgumentOutOfRangeException(nameof(propertyName), propertyName, null);
+            }
+
+#pragma warning disable SCME0001
+            if (unrelatedPatch)
+            {
+                model.Patch.Set("$.unrelated"u8, 1);
+            }
+#pragma warning restore SCME0001
+
+            var buffer = new ArrayBufferWriter<byte>();
+            using var writer = new Utf8JsonWriter(buffer);
+            var jsonModel = (IJsonModel<NullableDynamicModel>)model;
+            jsonModel.Write(writer, ModelReaderWriterOptions.Json);
+            writer.Flush();
+            buffer.Clear();
+            writer.Reset(buffer);
+
+            long before = GC.GetAllocatedBytesForCurrentThread();
+            jsonModel.Write(writer, ModelReaderWriterOptions.Json);
+            writer.Flush();
+            long allocated = GC.GetAllocatedBytesForCurrentThread() - before;
+
+            // Shapes whose innermost collection element is an IDictionary<,> (nestedChildDictionary,
+            // listOfDictionaries) enumerate that dictionary via its interface once per outer element,
+            // which boxes a struct enumerator regardless of any patch guard. That is unrelated overhead
+            // this PR does not address, so those shapes use a higher bound; all shapes must stay far
+            // below what unconditional per-element interpolated-path formatting would cost (tens of KB).
+            long maxAllocated = propertyName is "nestedChildDictionary" or "listOfDictionaries" ? 20 * 1024 : 4096;
+            Assert.That(allocated, Is.LessThan(maxAllocated), "Indexed patch paths must not allocate for each unpatched nested element.");
+            using var document = JsonDocument.Parse(buffer.WrittenMemory);
+            var collection = document.RootElement.GetProperty(propertyName);
+            Assert.That(
+                collection.ValueKind == JsonValueKind.Array ? collection.GetArrayLength() : collection.EnumerateObject().Count(),
+                Is.EqualTo(Count));
+        }
+
+        [TestCase("nestedChildren")]
+        [TestCase("nestedChildDictionary")]
+        [TestCase("dictionaryChildren")]
+        [TestCase("listOfDictionaries")]
+        public void JsonModelWrite_UnpatchedNestedCollectionSerializes(string propertyName)
+        {
+            const int Count = 256;
+            var model = new NullableDynamicModel();
+            switch (propertyName)
+            {
+                case "nestedChildren":
+                    model.NestedChildren = Enumerable.Range(0, Count)
+                        .Select(_ => (IList<AnotherDynamicModel>)[new AnotherDynamicModel("value")])
+                        .ToList();
+                    break;
+                case "nestedChildDictionary":
+                    model.NestedChildDictionary = Enumerable.Range(0, Count)
+                        .ToDictionary(
+                            index => index.ToString(),
+                            _ => (IDictionary<string, AnotherDynamicModel>)new Dictionary<string, AnotherDynamicModel>
+                            {
+                                ["value"] = new AnotherDynamicModel("value")
+                            });
+                    break;
+                case "dictionaryChildren":
+                    model.DictionaryChildren = Enumerable.Range(0, Count)
+                        .ToDictionary(
+                            index => index.ToString(),
+                            _ => (IList<AnotherDynamicModel>)[new AnotherDynamicModel("value")]);
+                    break;
+                case "listOfDictionaries":
+                    model.ListOfDictionaries = Enumerable.Range(0, Count)
+                        .Select(_ => (IDictionary<string, AnotherDynamicModel>)new Dictionary<string, AnotherDynamicModel>
+                        {
+                            ["value"] = new AnotherDynamicModel("value")
+                        })
+                        .ToList();
+                    break;
+                default:
+                    throw new ArgumentOutOfRangeException(nameof(propertyName), propertyName, null);
+            }
+
+            var buffer = new ArrayBufferWriter<byte>();
+            using var writer = new Utf8JsonWriter(buffer);
+            var jsonModel = (IJsonModel<NullableDynamicModel>)model;
+            jsonModel.Write(writer, ModelReaderWriterOptions.Json);
+            writer.Flush();
+            buffer.Clear();
+            writer.Reset(buffer);
+
+            jsonModel.Write(writer, ModelReaderWriterOptions.Json);
+            writer.Flush();
+
+            using var document = JsonDocument.Parse(buffer.WrittenMemory);
+            var collection = document.RootElement.GetProperty(propertyName);
+            Assert.That(
+                collection.ValueKind == JsonValueKind.Array ? collection.GetArrayLength() : collection.EnumerateObject().Count(),
+                Is.EqualTo(Count));
+        }
+
+        [TestCase("nestedChildren")]
+        [TestCase("nestedChildDictionary")]
+        [TestCase("dictionaryChildren")]
+        [TestCase("listOfDictionaries")]
+        public void JsonModelWrite_PatchedNestedCollectionSerializesParentPatch(string propertyName)
+        {
+            var model = new NullableDynamicModel();
+            var patchPath = propertyName switch
+            {
+                "nestedChildren" => SetNestedChildren(model),
+                "nestedChildDictionary" => SetNestedChildDictionary(model),
+                "dictionaryChildren" => SetDictionaryChildren(model),
+                "listOfDictionaries" => SetListOfDictionaries(model),
+                _ => throw new ArgumentOutOfRangeException(nameof(propertyName), propertyName, null)
+            };
+
+#pragma warning disable SCME0001
+            model.Patch.Set(Encoding.UTF8.GetBytes(patchPath), "patched");
+#pragma warning restore SCME0001
+
+            var data = ModelReaderWriter.Write(model, ModelReaderWriterOptions.Json, SampleTypeSpecContext.Default);
+            using var document = JsonDocument.Parse(data);
+            JsonElement patchedElement;
+            switch (propertyName)
+            {
+                case "nestedChildren":
+                    var nestedChildren = document.RootElement.GetProperty("nestedChildren")[0];
+                    patchedElement = nestedChildren[0];
+                    break;
+                case "nestedChildDictionary":
+                    var nestedChildDictionary = document.RootElement.GetProperty("nestedChildDictionary").GetProperty("outer");
+                    patchedElement = nestedChildDictionary.GetProperty("patched");
+                    break;
+                case "dictionaryChildren":
+                    var dictionaryChildren = document.RootElement.GetProperty("dictionaryChildren").GetProperty("outer");
+                    patchedElement = dictionaryChildren[0];
+                    break;
+                case "listOfDictionaries":
+                    var listOfDictionaries = document.RootElement.GetProperty("listOfDictionaries")[0];
+                    patchedElement = listOfDictionaries.GetProperty("patched");
+                    break;
+                default:
+                    throw new ArgumentOutOfRangeException(nameof(propertyName), propertyName, null);
+            }
+
+            Assert.That(patchedElement.GetProperty("extra").GetString(), Is.EqualTo("patched"));
+
+            static string SetNestedChildren(NullableDynamicModel model)
+            {
+                model.NestedChildren = [[null!]];
+                return "$.nestedChildren[0][0].extra";
+            }
+
+            static string SetNestedChildDictionary(NullableDynamicModel model)
+            {
+                model.NestedChildDictionary = new Dictionary<string, IDictionary<string, AnotherDynamicModel>>
+                {
+                    ["outer"] = new Dictionary<string, AnotherDynamicModel>
+                    {
+                        ["patched"] = null!
+                    }
+                };
+                return "$.nestedChildDictionary.outer.patched.extra";
+            }
+
+            static string SetDictionaryChildren(NullableDynamicModel model)
+            {
+                model.DictionaryChildren = new Dictionary<string, IList<AnotherDynamicModel>>
+                {
+                    ["outer"] = [null!]
+                };
+                return "$.dictionaryChildren.outer[0].extra";
+            }
+
+            static string SetListOfDictionaries(NullableDynamicModel model)
+            {
+                model.ListOfDictionaries =
+                [
+                    new Dictionary<string, AnotherDynamicModel>
+                    {
+                        ["patched"] = null!
+                    }
+                ];
+                return "$.listOfDictionaries[0].patched.extra";
+            }
+        }
+
+        [TestCase(false)]
+        [TestCase(true)]
+        public void JsonPatchRemove_ChildRootWithUnpatchedParentCollection(bool unrelatedPatch)
+        {
+            var removed = new AnotherDynamicModel("removed");
+            var model = new NullableDynamicModel
+            {
+                Children = [null, removed, new AnotherDynamicModel("present")]
+            };
+
+#pragma warning disable SCME0001
+            removed.Patch.Remove("$"u8);
+            if (unrelatedPatch)
+            {
+                model.Patch.Set("$.unrelated"u8, 1);
+            }
+            Assert.That(model.Patch.Contains("$"u8, "children"u8), Is.False);
+            var snapshot = model.Patch.GetJson("$.children"u8);
+#pragma warning restore SCME0001
+
+            Assert.That(Encoding.UTF8.GetString(snapshot), Is.EqualTo("""[null,{"bar":"present"}]"""));
+            var data = ModelReaderWriter.Write(model, ModelReaderWriterOptions.Json, SampleTypeSpecContext.Default);
+            using var document = JsonDocument.Parse(data);
+            Assert.That(document.RootElement.GetProperty("children").GetRawText(), Is.EqualTo("""[null,{"bar":"present"}]"""));
+        }
+
         private static NullableDynamicModel CreateModelWithRemovedDynamicListElements(string propertyName, bool onlyNull)
         {
             var items = onlyNull ? "[null]" : """[null,{"bar":"present"},null]""";

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator.ClientModel/test/Providers/MrwSerializationTypeDefinitions/TestData/DynamicModelSerializationTests/PropagateModelListPropertyHelperMethods.cs

@@ -30,9 +30,10 @@ private bool TryResolveP1Array(out global::System.ClientModel.Primitives.JsonPat
             {
                 yield break;
             }
+            bool hasPatch = Patch.Contains("$"u8, "p1"u8);
             for (int i = 0; (i < P1.Count); i++)
             {
-                if ((!Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.p1[{i}]")) && ((P1[i] == null) || !P1[i].Patch.IsRemoved("$"u8))))
+                if (((!hasPatch || !Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.p1[{i}]"))) && ((P1[i] == null) || !P1[i].Patch.IsRemoved("$"u8))))
                 {
                     yield return P1[i];
                 }

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator.ClientModel/test/Providers/MrwSerializationTypeDefinitions/TestData/DynamicModelSerializationTests/WriteArrayProperties.cs

@@ -49,9 +49,10 @@ protected virtual void JsonModelWriteCore(global::System.Text.Json.Utf8JsonWrite
             {
                 writer.WritePropertyName("cats"u8);
                 writer.WriteStartArray();
+                bool hasPatch = Patch.Contains("$"u8, "cats"u8);
                 for (int i = 0; (i < Cats.Count); i++)
                 {
-                    if (Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.cats[{i}]")))
+                    if ((hasPatch && Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.cats[{i}]"))))
                     {
                         continue;
                     }
@@ -72,9 +73,10 @@ protected virtual void JsonModelWriteCore(global::System.Text.Json.Utf8JsonWrite
             {
                 writer.WritePropertyName("names"u8);
                 writer.WriteStartArray();
+                bool hasPatch = Patch.Contains("$"u8, "names"u8);
                 for (int i = 0; (i < Names.Count); i++)
                 {
-                    if (Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.names[{i}]")))
+                    if ((hasPatch && Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.names[{i}]"))))
                     {
                         continue;
                     }
@@ -100,9 +102,10 @@ protected virtual void JsonModelWriteCore(global::System.Text.Json.Utf8JsonWrite
             {
                 writer.WritePropertyName("optionalNames"u8);
                 writer.WriteStartArray();
+                bool hasPatch = Patch.Contains("$"u8, "optionalNames"u8);
                 for (int i = 0; (i < OptionalNames.Count); i++)
                 {
-                    if (Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.optionalNames[{i}]")))
+                    if ((hasPatch && Patch.IsRemoved(global::System.Text.Encoding.UTF8.GetBytes($"$.optionalNames[{i}]"))))
                     {
                         continue;
                     }