Implement MD array support

eiriktsarpalis · Apr 6, 2024 · 6fedf2a · 6fedf2a
1 parent 9550799
commit 6fedf2a
Show file tree

Hide file tree

Showing 7 changed files with 115 additions and 58 deletions.
diff --git a/src/TypeShape.Applications/JsonSerializer/Converters/JsonEnumerableConverter.cs b/src/TypeShape.Applications/JsonSerializer/Converters/JsonEnumerableConverter.cs
@@ -1,6 +1,9 @@
-namespace TypeShape.Applications.JsonSerializer.Converters;
+using System.Diagnostics;
+
+namespace TypeShape.Applications.JsonSerializer.Converters;
 
 using System;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Text.Json;
@@ -123,68 +126,115 @@ private protected override TEnumerable Construct(PooledList<TElement> buffer)
         => spanConstructor(buffer.AsSpan());
 }
 
-internal sealed class Json2DArrayConverter<TElement>(JsonConverter<TElement> elementConverter) : JsonConverter<TElement[,]>
+internal sealed class JsonMDArrayConverter<TArray, TElement>(JsonConverter<TElement> elementConverter, int rank) : JsonConverter<TArray>
 {
-    public override TElement[,]? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+    [ThreadStatic] private static int[]? _dimensions;
+
+    [UnconditionalSuppressMessage("AOT", "IL3050", Justification = "The Array.CreateInstance method generates TArray instances.")]
+    public override TArray? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
     {
         if (reader.TokenType is JsonTokenType.Null)
         {
-            return null;
+            return default;
         }
 
-        reader.EnsureTokenType(JsonTokenType.StartArray);
-        reader.EnsureRead();
+        int[] dimensions = _dimensions ??= new int[rank];
+        dimensions.AsSpan().Fill(-1);
+        PooledList<TElement> buffer = new();
+        try
+        {
+            ReadSubArray(ref reader, ref buffer, dimensions, options);
+            dimensions.AsSpan().Replace(-1, 0);
+            Array result = Array.CreateInstance(typeof(TElement), dimensions);
+            buffer.AsSpan().CopyTo(AsSpan(result));
+            return (TArray)(object)result;
+        }
+        finally
+        {
+            buffer.Dispose();
+        }
+    }
 
-        using PooledList<TElement> buffer = new();
-        int rows = 0;
-        int? columns = null;
+    public override void Write(Utf8JsonWriter writer, TArray value, JsonSerializerOptions options)
+    {
+        var array = (Array)(object)value!;
+        Debug.Assert(rank == array.Rank);
 
+        int[] dimensions = _dimensions ??= new int[rank];
+        for (int i = 0; i < rank; i++) dimensions[i] = array.GetLength(i);
+        WriteSubArray(writer, dimensions, AsSpan(array), options);
+    }
+
+    private void ReadSubArray(
+        ref Utf8JsonReader reader,
+        ref PooledList<TElement> buffer,
+        Span<int> dimensions,
+        JsonSerializerOptions options)
+    {
+        Debug.Assert(dimensions.Length > 0);
+        reader.EnsureTokenType(JsonTokenType.StartArray);
+        reader.EnsureRead();
+
+        int dimension = 0;
         while (reader.TokenType != JsonTokenType.EndArray)
         {
-            reader.EnsureTokenType(JsonTokenType.StartArray);
-            reader.EnsureRead();
-
-            int rowLength = 0;
-            while (reader.TokenType != JsonTokenType.EndArray)
+            if (dimensions.Length > 1)
+            {
+                ReadSubArray(ref reader, ref buffer, dimensions[1..], options);
+            }
+            else
             {
                 TElement? element = elementConverter.Read(ref reader, typeof(TElement), options);
                 buffer.Add(element!);
-                reader.EnsureRead();
-                rowLength++;
             }
-
+            
             reader.EnsureRead();
-            rows++;
-
-            if ((columns ??= rowLength) != rowLength)
-            {
-                JsonHelpers.ThrowJsonException("The deserialized jagged array must be rectangular.");
-            }
+            dimension++;
+        }
+
+        if (dimensions[0] < 0)
+        {
+            dimensions[0] = dimension;
+        }
+        else if (dimensions[0] != dimension)
+        {
+            JsonHelpers.ThrowJsonException("The deserialized jagged array was not rectangular.");
         }
-
-        TElement[,] result = new TElement[rows, columns ?? 0];
-        Span<TElement> destination = MemoryMarshal.CreateSpan(ref Unsafe.As<byte, TElement>(ref MemoryMarshal.GetArrayDataReference(result)), result.Length);
-        buffer.AsSpan().CopyTo(destination);
-        return result;
     }
-
-    public override void Write(Utf8JsonWriter writer, TElement[,] value, JsonSerializerOptions options)
+
+    private void WriteSubArray(
+        Utf8JsonWriter writer,
+        ReadOnlySpan<int> dimensions,
+        ReadOnlySpan<TElement> elements,
+        JsonSerializerOptions options)
     {
-        int n = value.GetLength(0);
-        int m = value.GetLength(1);
-
+        Debug.Assert(dimensions.Length > 0);
+
         writer.WriteStartArray();
-        for (int i = 0; i < n; i++)
+
+        int outerDim = dimensions[0];
+        if (dimensions.Length > 1 && outerDim > 0)
         {
-            writer.WriteStartArray();
-            for (int j = 0; j < m; j++)
+            int subArrayLength = elements.Length / outerDim;
+            for (int i = 0; i < outerDim; i++)
             {
-                elementConverter.Write(writer, value[i, j], options);
+                WriteSubArray(writer, dimensions[1..], elements[..subArrayLength], options);
+                elements = elements[subArrayLength..];
             }
-
-            writer.WriteEndArray();
         }
-
+        else
+        {
+            for (int i = 0; i < outerDim; i++)
+            {
+                elementConverter.Write(writer, elements[i], options);
+            }
+        }
+
         writer.WriteEndArray();
     }
+
+    private static Span<TElement> AsSpan(Array array) =>
+        MemoryMarshal.CreateSpan(
+            ref Unsafe.As<byte, TElement>(ref MemoryMarshal.GetArrayDataReference(array)), 
+            array.Length);
 }
diff --git a/src/TypeShape.Applications/JsonSerializer/Converters/JsonValueConverters.cs b/src/TypeShape.Applications/JsonSerializer/Converters/JsonValueConverters.cs
@@ -140,7 +140,7 @@ public override void WriteAsPropertyName(Utf8JsonWriter writer, object? value, J
 
         static ITypeShapeJsonConverter ResolveDerivedConverter(Type derivedType, ITypeShapeProvider provider)
         {
-            if (provider.GetShape(derivedType) is not ITypeShape derivedShape)
+            if (provider.GetShape(derivedType) is not { } derivedShape)
             {
                 throw new NotSupportedException($"Unsupported derived type '{derivedType}'.");
             }

diff --git a/src/TypeShape.Applications/JsonSerializer/TypeShapeJsonSerializer.Builder.cs b/src/TypeShape.Applications/JsonSerializer/TypeShapeJsonSerializer.Builder.cs
@@ -36,8 +36,7 @@ private JsonConverter<T> BuildConverter<T>(ITypeShape<T> typeShape)
                 return new JsonObjectConverter(type.Provider);
             }
 
-            JsonPropertyConverter<T>[] properties = type
-                .GetProperties()
+            JsonPropertyConverter<T>[] properties = type.GetProperties()
                 .Select(prop => (JsonPropertyConverter<T>)prop.Accept(this)!)
                 .ToArray();
 
@@ -91,11 +90,7 @@ private JsonConverter<T> BuildConverter<T>(ITypeShape<T> typeShape)
             if (enumerableTypeShape.Rank > 1)
             {
                 Debug.Assert(typeof(TEnumerable).IsArray);
-                return enumerableTypeShape.Rank switch
-                {
-                    2 => new Json2DArrayConverter<TElement>(elementConverter),
-                    _ => throw new NotImplementedException("Array rank > 2 not implemented."),
-                };
+                return new JsonMDArrayConverter<TEnumerable, TElement>(elementConverter, enumerableTypeShape.Rank);
             }
 
             return enumerableTypeShape.ConstructionStrategy switch

diff --git a/src/TypeShape.Applications/JsonSerializer/TypeShapeJsonSerializer.cs b/src/TypeShape.Applications/JsonSerializer/TypeShapeJsonSerializer.cs
@@ -1,8 +1,4 @@
-using System.Text.Json;
-using System.Text.Json.Serialization;
-using System.Text.Json.Serialization.Metadata;
-
-namespace TypeShape.Applications.JsonSerializer;
+namespace TypeShape.Applications.JsonSerializer;
 
 public static partial class TypeShapeJsonSerializer
 {

diff --git a/tests/TypeShape.Tests/JsonSchemaTests.cs b/tests/TypeShape.Tests/JsonSchemaTests.cs
@@ -95,11 +95,6 @@ void AssertType(string type)
     [MemberData(nameof(TestTypes.GetTestCases), MemberType = typeof(TestTypes))]
     public void SchemaMatchesJsonSerializer<T>(TestCase<T> testCase)
     {
-        if (typeof(T).IsArray && typeof(T).GetArrayRank() > 2)
-        {
-            return; // Serialization not implemented for arrays with rank > 2
-        }
-
         if (typeof(T) == typeof(Int128) || typeof(T) == typeof(UInt128))
         {
             return; // Not supported by JsonSchema.NET

diff --git a/tests/TypeShape.Tests/JsonTests.cs b/tests/TypeShape.Tests/JsonTests.cs
@@ -314,6 +314,8 @@ public void MultiDimensionalArrays_SerializedAsJaggedArray<TArray>(TArray array,
     public static IEnumerable<object?[]> GetMultiDimensionalArraysAndExpectedJson()
     {
         yield return Wrap(new int[,] { }, """[]""");
+        yield return Wrap(new int[,,] { }, """[]""");
+        yield return Wrap(new int[,,,,,] { }, """[]""");
 
         yield return Wrap(
             new int[,] { { 1, 0, }, { 0, 1 } },
@@ -326,6 +328,24 @@ public void MultiDimensionalArrays_SerializedAsJaggedArray<TArray>(TArray array,
         yield return Wrap(
             new int[,] { { 1, 2, 3 }, { 4, 5, 6 } },
             """[[1,2,3],[4,5,6]]""");
+
+        yield return Wrap(
+            new int[,,] // 3 x 2 x 2
+            {
+                { { 1, 0 }, { 0, 1 } }, 
+                { { 1, 2 }, { 3, 4 } }, 
+                { { 1, 1 }, { 1, 1 } }
+            },
+            """[[[1,0],[0,1]],[[1,2],[3,4]],[[1,1],[1,1]]]""");
+
+        yield return Wrap(
+            new int[,,] // 3 x 2 x 5
+            {
+                { { 1, 0, 0, 0, 0 }, { 0, 1, 0, 0, 0 } }, 
+                { { 1, 2, 3, 4, 5 }, { 6, 7, 8, 9, 10 } }, 
+                { { 1, 1, 1, 1, 1 }, { 1, 1, 1, 1, 1 } }
+            },
+            """[[[1,0,0,0,0],[0,1,0,0,0]],[[1,2,3,4,5],[6,7,8,9,10]],[[1,1,1,1,1],[1,1,1,1,1]]]""");
 
         static object?[] Wrap<TArray>(TArray tuple, string expectedJson) where TArray : IEnumerable
             => [tuple, expectedJson];

diff --git a/tests/TypeShape.Tests/TestTypes.cs b/tests/TypeShape.Tests/TestTypes.cs
@@ -1329,6 +1329,7 @@ public record DerivedClassWithShadowingMember : BaseClassWithShadowingMembers
 [GenerateShape<int[][]>]
 [GenerateShape<int[,]>]
 [GenerateShape<int[,,]>]
+[GenerateShape<int[,,,,,]>]
 [GenerateShape<Memory<int>>]
 [GenerateShape<ReadOnlyMemory<int>>]
 [GenerateShape<List<string>>]