Parallel codec visiting

src/EdgeDB.Net.Driver/Binary/Builders/ObjectBuilder.cs

@@ -1,5 +1,6 @@
 using EdgeDB.Binary.Codecs;
 using EdgeDB.DataTypes;
+using EdgeDB.Utils;
 using Microsoft.Extensions.Logging;
 using System.Collections;
 using System.Collections.Concurrent;
@@ -12,7 +13,8 @@ internal sealed class ObjectBuilder
     private static readonly ConcurrentDictionary<Type, (int Version, ICodec Codec)> _codecVisitorStateTable = new();
     private static readonly object _visitorLock = new();
 
-    public static PreheatedCodec PreheatCodec<T>(EdgeDBBinaryClient client, ICodec codec)
+    public static async Task<PreheatedCodec> PreheatCodecAsync<T>(EdgeDBBinaryClient client, ICodec codec,
+        CancellationToken token)
     {
         // if the codec has been visited before and we have the most up-to-date version, return it.
         if (
@@ -28,17 +30,18 @@ public static PreheatedCodec PreheatCodec<T>(EdgeDBBinaryClient client, ICodec c
 
         var visitor = new TypeVisitor(client);
         visitor.SetTargetType(typeof(T));
-        visitor.Visit(ref codec);
+        var reference = new Ref<ICodec>(codec);
+        await visitor.VisitAsync(reference, token);
 
         if (typeof(T) != typeof(object))
-            _codecVisitorStateTable[typeof(T)] = (version, codec);
+            _codecVisitorStateTable[typeof(T)] = (version, reference.Value);
 
         if (client.Logger.IsEnabled(LogLevel.Debug))
         {
-            client.Logger.ObjectDeserializationPrep(CodecFormatter.Format(codec).ToString());
+            client.Logger.ObjectDeserializationPrep(CodecFormatter.Format(reference.Value).ToString());
         }
 
-        return new PreheatedCodec(codec);
+        return new PreheatedCodec(reference.Value);
     }
 
     public static T? BuildResult<T>(EdgeDBBinaryClient client, in PreheatedCodec preheated,
@@ -54,8 +57,9 @@ public static PreheatedCodec PreheatCodec<T>(EdgeDBBinaryClient client, ICodec c
         return (T?)ConvertTo(typeof(T), value);
     }
 
-    public static T? BuildResult<T>(EdgeDBBinaryClient client, ICodec codec, in ReadOnlyMemory<byte> data)
-        => BuildResult<T>(client, PreheatCodec<T>(client, codec), data);
+    public static async Task<T?> BuildResultAsync<T>(
+        EdgeDBBinaryClient client, ICodec codec, ReadOnlyMemory<byte> data, CancellationToken token)
+        => BuildResult<T>(client, await PreheatCodecAsync<T>(client, codec, token), data);
 
     public static object? ConvertTo(Type type, object? value)
     {

src/EdgeDB.Net.Driver/Binary/Codecs/ArgumentCodecContext.cs

@@ -0,0 +1,12 @@
+namespace EdgeDB.Binary.Codecs;
+
+internal sealed class ArgumentCodecContext : CodecContext
+{
+    public readonly ICodec[] Codecs;
+
+    public ArgumentCodecContext(ICodec[] codecs, EdgeDBBinaryClient client)
+        : base(client)
+    {
+        Codecs = codecs;
+    }
+}

src/EdgeDB.Net.Driver/Binary/Codecs/ArrayCodec.cs

@@ -3,15 +3,25 @@
 namespace EdgeDB.Binary.Codecs;
 
 internal sealed class ArrayCodec<T>
-    : BaseCodec<T?[]>, IWrappingCodec, ICacheableCodec
+    : BaseCodec<T?[]>, IWrappingCodec, ICacheableCodec, ICompiledCodec
 {
     public static readonly byte[] EMPTY_ARRAY = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1};
 
+    public Type CompiledFrom { get; }
+    public CompilableWrappingCodec Template { get; }
+
     internal ICodec<T> InnerCodec;
 
-    public ArrayCodec(in Guid id, ICodec<T> innerCodec, CodecMetadata? metadata = null)
+    public ArrayCodec(
+        in Guid id,
+        Type compiledFrom,
+        CompilableWrappingCodec template,
+        ICodec<T> innerCodec,
+        CodecMetadata? metadata = null)
         : base(in id, metadata)
     {
+        CompiledFrom = compiledFrom;
+        Template = template;
         InnerCodec = innerCodec;
     }
 

src/EdgeDB.Net.Driver/Binary/Codecs/CodecContext.cs

@@ -2,7 +2,7 @@
 
 namespace EdgeDB.Binary.Codecs;
 
-internal sealed class CodecContext
+internal class CodecContext
 {
     public CodecContext(EdgeDBBinaryClient client)
     {

src/EdgeDB.Net.Driver/Binary/Codecs/CodecFormatter.cs

@@ -5,6 +5,45 @@ namespace EdgeDB.Binary.Codecs;
 
 internal static class CodecFormatter
 {
+    public static StringBuilder FormatCodecAsTree(ICodec codec)
+    {
+        var sb = new StringBuilder();
+
+        AppendCodecToTree(sb, codec, 0);
+
+        return sb;
+    }
+
+
+    private const int CODEC_TREE_SPACING = 2;
+    private static void AppendCodecToTree(StringBuilder tree, ICodec codec, int depth, string? prefix = null)
+    {
+        tree.AppendLine("".PadLeft(depth) + $"{prefix} {codec}");
+
+        if (codec is IMultiWrappingCodec multiwrap)
+        {
+            for (int i = 0; i != multiwrap.InnerCodecs.Length; i++)
+            {
+                var innerCodec = multiwrap.InnerCodecs[i];
+                AppendCodecToTree(
+                    tree,
+                    innerCodec,
+                    depth + CODEC_TREE_SPACING,
+                    i == multiwrap.InnerCodecs.Length - 1 ? "\u2514" : "\u251c"
+                );
+            }
+        }
+        else if (codec is IWrappingCodec wrapping)
+        {
+            AppendCodecToTree(
+                tree,
+                wrapping.InnerCodec,
+                depth + CODEC_TREE_SPACING,
+                "\u2514"
+            );
+        }
+    }
+
     public static StringBuilder Format(ICodec codec, int spacing = 2)
     {
         StringBuilder sb = new();

src/EdgeDB.Net.Driver/Binary/Codecs/CompilableCodec.cs

@@ -31,7 +31,7 @@ void ICodec.Serialize(ref PacketWriter writer, object? value, CodecContext conte
     object? ICodec.Deserialize(ref PacketReader reader, CodecContext context) => throw new NotSupportedException();
 
     // to avoid state changes to this compilable, pass in the inner codec post-walk.
-    public ICodec Compile(IProtocolProvider provider, Type type, ICodec? innerCodec = null)
+    public ICompiledCodec Compile(IProtocolProvider provider, Type type, ICodec? innerCodec = null)
     {
         innerCodec ??= InnerCodec;
 
@@ -41,15 +41,19 @@ public ICodec Compile(IProtocolProvider provider, Type type, ICodec? innerCodec
 
         return CodecBuilder.GetProviderCache(provider).CompiledCodecCache.GetOrAdd(cacheKey, k =>
         {
-            var codec = (ICodec)Activator.CreateInstance(genType, Id, innerCodec, Metadata)!;
+            var codec = (ICompiledCodec)Activator.CreateInstance(
+                genType,
+                Id, type, this, innerCodec, Metadata
+            )!;
 
             if (codec is IComplexCodec complex)
             {
                 complex.BuildRuntimeCodecs(provider);
             }
 
             return codec;
-        });
+        }) as ICompiledCodec
+            ?? throw new InvalidOperationException("Codec that was returned from the cache was not a compiled codec, this shouldn't happen");
     }
 
     public Type GetInnerType()
@@ -58,5 +62,5 @@ public Type GetInnerType()
             : InnerCodec.ConverterType;
 
     public override string ToString()
-        => $"compilable({_rootCodecType.Name})";
+        => $"compilable({InnerCodec})";
 }

src/EdgeDB.Net.Driver/Binary/Codecs/Impl/BaseArgumentCodec.cs

@@ -9,8 +9,8 @@ protected BaseArgumentCodec(in Guid id, CodecMetadata? metadata)
     {
     }
 
-    public abstract void SerializeArguments(ref PacketWriter writer, T? value, CodecContext context);
+    public abstract void SerializeArguments(ref PacketWriter writer, T? value, ArgumentCodecContext context);
 
-    void IArgumentCodec.SerializeArguments(ref PacketWriter writer, object? value, CodecContext context)
+    void IArgumentCodec.SerializeArguments(ref PacketWriter writer, object? value, ArgumentCodecContext context)
         => SerializeArguments(ref writer, (T?)value, context);
 }

src/EdgeDB.Net.Driver/Binary/Codecs/Impl/IArgumentCodec.cs

@@ -2,10 +2,10 @@ namespace EdgeDB.Binary.Codecs;
 
 internal interface IArgumentCodec<T> : IArgumentCodec, ICodec<T>
 {
-    void SerializeArguments(ref PacketWriter writer, T? value, CodecContext context);
+    void SerializeArguments(ref PacketWriter writer, T? value, ArgumentCodecContext context);
 }
 
 internal interface IArgumentCodec
 {
-    void SerializeArguments(ref PacketWriter writer, object? value, CodecContext context);
+    void SerializeArguments(ref PacketWriter writer, object? value, ArgumentCodecContext context);
 }

src/EdgeDB.Net.Driver/Binary/Codecs/Impl/ICompiledCodec.cs

@@ -0,0 +1,7 @@
+namespace EdgeDB.Binary.Codecs;
+
+internal interface ICompiledCodec : ICodec, IWrappingCodec
+{
+    Type CompiledFrom { get; }
+    CompilableWrappingCodec Template { get; }
+}

src/EdgeDB.Net.Driver/Binary/Codecs/MultiRangeCodec.cs

@@ -4,14 +4,32 @@
 
 namespace EdgeDB.Binary.Codecs;
 
-internal sealed class MultiRangeCodec<T> : BaseCodec<MultiRange<T>>, IWrappingCodec, ICacheableCodec
+internal sealed class MultiRangeCodec<T>
+    : BaseCodec<MultiRange<T>>, IWrappingCodec, ICacheableCodec, ICompiledCodec
     where T : struct
 {
+    public Type CompiledFrom { get; }
+    public CompilableWrappingCodec Template { get; }
+
     private RangeCodec<T> _rangeCodec;
 
-    public MultiRangeCodec(in Guid id, ICodec<T> rangeInnerCodec, CodecMetadata? metadata) : base(in id, metadata)
+    public MultiRangeCodec(
+        in Guid id,
+        Type compiledFrom,
+        CompilableWrappingCodec template,
+        ICodec<T> rangeInnerCodec,
+        CodecMetadata? metadata
+        ) : base(in id, metadata)
     {
-        _rangeCodec = new RangeCodec<T>(in id, rangeInnerCodec, metadata);
+        CompiledFrom = compiledFrom;
+        Template = template;
+        _rangeCodec = new RangeCodec<T>(
+            in id,
+            compiledFrom,
+            template,
+            rangeInnerCodec,
+            metadata
+        );
     }
 
     public override void Serialize(ref PacketWriter writer, MultiRange<T> value, CodecContext context)

src/EdgeDB.Net.Driver/Binary/Codecs/NullCodec.cs

@@ -9,7 +9,7 @@ public NullCodec() { }
 
     public NullCodec(CodecMetadata? metadata = null) { } // used in generic codec construction
 
-    public void SerializeArguments(ref PacketWriter writer, object? value, CodecContext context) { }
+    public void SerializeArguments(ref PacketWriter writer, object? value, ArgumentCodecContext context) { }
 
     public Guid Id
         => Guid.Empty;

src/EdgeDB.Net.Driver/Binary/Codecs/ObjectCodec.cs

@@ -53,6 +53,9 @@ public EdgeDBTypeDeserializeInfo Deserializer
             reader.Position = enumerator.Reader.Position;
         }
     }
+
+    public override string ToString()
+        => $"object ({TargetType.Name})";
 }
 
 internal class ObjectCodec
@@ -109,10 +112,13 @@ public TypeInitializedObjectCodec GetOrCreateTypeCodec(Type type)
         }
     }
 
-    public override void SerializeArguments(ref PacketWriter writer, object? value, CodecContext context)
-        => Serialize(ref writer, value, context);
+    public override void SerializeArguments(ref PacketWriter writer, object? value, ArgumentCodecContext context)
+        => Serialize(ref writer, value, context.Codecs, context);
 
     public override void Serialize(ref PacketWriter writer, object? value, CodecContext context)
+        => Serialize(ref writer, value, InnerCodecs, context);
+
+    private void Serialize(ref PacketWriter writer, object? value, ICodec[] codecs, CodecContext context)
     {
         object?[]? values = null;
 
@@ -128,9 +134,6 @@ public override void Serialize(ref PacketWriter writer, object? value, CodecCont
 
         writer.Write(values.Length);
 
-        // TODO: maybe cache the visited codecs based on the 'value'.
-        var visitor = context.CreateTypeVisitor();
-
         for (var i = 0; i != values.Length; i++)
         {
             var element = values[i];
@@ -156,17 +159,11 @@ public override void Serialize(ref PacketWriter writer, object? value, CodecCont
             }
             else
             {
-                var innerCodec = InnerCodecs[i];
+                var innerCodec = codecs[i];
 
                 // special case for enums
                 if (element.GetType().IsEnum && innerCodec is TextCodec)
                     element = element.ToString();
-                else
-                {
-                    visitor.SetTargetType(element.GetType());
-                    visitor.Visit(ref innerCodec);
-                    visitor.Reset();
-                }
 
                 writer.WriteToWithInt32Length((ref PacketWriter innerWriter) =>
                     innerCodec.Serialize(ref innerWriter, element, context));

src/EdgeDB.Net.Driver/Binary/Codecs/RangeCodec.cs

@@ -5,26 +5,37 @@
 namespace EdgeDB.Binary.Codecs;
 
 internal sealed class RangeCodec<T>
-    : BaseComplexCodec<Range<T>>, IWrappingCodec, ICacheableCodec
+    : BaseComplexCodec<Range<T>>, IWrappingCodec, ICacheableCodec, ICompiledCodec
     where T : struct
 {
     [Flags]
     public enum RangeFlags : byte
     {
         Empty = 1 << 0,
-        IncudeLowerBound = 1 << 1,
+        IncludeLowerBound = 1 << 1,
         IncludeUpperBound = 1 << 2,
         InfiniteLowerBound = 1 << 3,
         InfiniteUpperBound = 1 << 4
     }
 
-    public ICodec<T> _innerCodec;
+    public Type CompiledFrom { get; }
+    public CompilableWrappingCodec Template { get; }
 
-    public RangeCodec(in Guid id, ICodec<T> innerCodec, CodecMetadata? metadata = null)
+    private ICodec<T> _innerCodec;
+
+    public RangeCodec(
+        in Guid id,
+        Type compiledFrom,
+        CompilableWrappingCodec template,
+        ICodec<T> innerCodec,
+        CodecMetadata? metadata = null)
         : base(in id, metadata)
     {
         _innerCodec = innerCodec;
 
+        CompiledFrom = compiledFrom;
+        Template = template;
+
         AddConverter(From, To);
     }
 
@@ -108,15 +119,15 @@ public override Range<T> Deserialize(ref PacketReader reader, CodecContext conte
             upperBound = _innerCodec.Deserialize(ref reader, context);
         }
 
-        return new Range<T>(lowerBound, upperBound, (flags & RangeFlags.IncudeLowerBound) != 0,
+        return new Range<T>(lowerBound, upperBound, (flags & RangeFlags.IncludeLowerBound) != 0,
             (flags & RangeFlags.IncludeUpperBound) != 0);
     }
 
     public override void Serialize(ref PacketWriter writer, Range<T> value, CodecContext context)
     {
         var flags = value.IsEmpty
             ? RangeFlags.Empty
-            : (value.IncludeLower ? RangeFlags.IncudeLowerBound : 0) |
+            : (value.IncludeLower ? RangeFlags.IncludeLowerBound : 0) |
               (value.IncludeUpper ? RangeFlags.IncludeUpperBound : 0) |
               (!value.Lower.HasValue ? RangeFlags.InfiniteLowerBound : 0) |
               (!value.Upper.HasValue ? RangeFlags.InfiniteUpperBound : 0);