Improve serialization speed of enums without data in C# (#2762)

Author: kazimuth

crates/bindings-csharp/BSATN.Codegen/Type.cs

@@ -73,7 +73,11 @@ public static TypeUse Parse(ISymbol member, ITypeSymbol typeSymbol, DiagReporter
                     Parse(member, named.TypeArguments[0], diag)
                 ),
                 _ => named.IsValueType
-                    ? new ValueUse(type, typeInfo)
+                    ? (
+                        named.TypeKind == Microsoft.CodeAnalysis.TypeKind.Enum
+                            ? new EnumUse(type, typeInfo)
+                            : new ValueUse(type, typeInfo)
+                    )
                     : new ReferenceUse(type, typeInfo),
             },
             _ => throw new InvalidOperationException($"Unsupported type {type}"),
@@ -113,7 +117,32 @@ public abstract string EqualsStatement(
 }
 
 /// <summary>
-/// A use of a value type.
+/// A use of an enum type.
+/// (This is a C# enum, not one of our tagged enums.)
+/// </summary>
+/// <param name="Type"></param>
+/// <param name="TypeInfo"></param>
+public record EnumUse(string Type, string TypeInfo) : TypeUse(Type, TypeInfo)
+{
+    // We just use `==` here, rather than `.Equals`, because
+    // C# enums don't provide a `bool Equals(Self other)`, and
+    // using `.Equals(object other)` allocates, which we want to avoid.
+    //
+    // We could instead generate custom .Equals for enums -- except that requires
+    // partial enums, and I'm not sure such things exist.
+    public override string EqualsStatement(
+        string inVar1,
+        string inVar2,
+        string outVar,
+        int level = 0
+    ) => $"var {outVar} = {inVar1} == {inVar2};";
+
+    public override string GetHashCodeStatement(string inVar, string outVar, int level = 0) =>
+        $"var {outVar} = {inVar}.GetHashCode();";
+}
+
+/// <summary>
+/// A use of a value type (that is not an enum).
 /// </summary>
 /// <param name="Type"></param>
 /// <param name="TypeInfo"></param>

crates/bindings-csharp/BSATN.Runtime.Tests/Tests.cs

@@ -656,6 +656,51 @@ public override int GetHashCode([DisallowNull] IEnumerable<T> obj)
         }
     }
 
+    [Type]
+    enum Banana
+    {
+        Cavendish,
+        LadyFinger,
+        RedBanana,
+        Manzano,
+        BlueJava,
+        GreenPlantain,
+        YellowPlantain,
+        PisangRaja,
+    }
+
+    [Fact]
+    public static void EnumSerializationWorks()
+    {
+        var serializer = new Enum<Banana>();
+        var bananas = new Banana[]
+        {
+            Banana.Cavendish,
+            Banana.LadyFinger,
+            Banana.RedBanana,
+            Banana.Manzano,
+            Banana.BlueJava,
+            Banana.GreenPlantain,
+            Banana.YellowPlantain,
+            Banana.PisangRaja,
+        };
+        for (var i = 0; i < bananas.Length; i++)
+        {
+            var stream = new MemoryStream();
+            var writer = new BinaryWriter(stream);
+            var banana = bananas[i];
+            serializer.Write(writer, banana);
+
+            stream.Seek(0, SeekOrigin.Begin);
+            var tag = new BinaryReader(stream).ReadByte();
+            Assert.Equal(tag, i);
+
+            stream.Seek(0, SeekOrigin.Begin);
+            var newBanana = serializer.Read(new BinaryReader(stream));
+            Assert.Equal(banana, newBanana);
+        }
+    }
+
     [Fact]
     public static void GeneratedNestedListEqualsWorks()
     {
@@ -775,4 +820,45 @@ [new BasicEnum.X(1), null],
         );
 #pragma warning restore CS8625 // Cannot convert null literal to non-nullable reference type.
     }
+
+    [Type]
+    partial struct ContainsEnum
+    {
+        public Banana TheBanana;
+        public int BananaCount;
+    }
+
+    static readonly Gen<(Banana, int)> GenContainsEnum = Gen.Select(
+        Gen.Enum<Banana>(),
+        Gen.Int[0, 3]
+    );
+    static readonly Gen<((Banana, int), (Banana, int))> GenTwoContainsEnum = Gen.Select(
+        GenContainsEnum,
+        GenContainsEnum
+    );
+
+    [Fact]
+    public static void GeneratedEnumEqualsWorks()
+    {
+        GenTwoContainsEnum.Sample(
+            example =>
+            {
+                var ((b1, c1), (b2, c2)) = example;
+                var struct1 = new ContainsEnum { TheBanana = b1, BananaCount = c1 };
+                var struct2 = new ContainsEnum { TheBanana = b2, BananaCount = c2 };
+
+                if ((b1, c1) == (b2, c2))
+                {
+                    Assert.True(struct1.Equals(struct2));
+                    Assert.Equal(struct1, struct2);
+                }
+                else
+                {
+                    Assert.False(struct1.Equals(struct2));
+                    Assert.NotEqual(struct1, struct2);
+                }
+            },
+            iter: 10_000
+        );
+    }
 }

crates/bindings-csharp/BSATN.Runtime/BSATN/Runtime.cs

@@ -118,24 +118,37 @@ public interface IReadWrite<T>
 }
 
 /// <summary>
-/// Present for backwards-compatibility reasons, but no longer used.
-/// The auto-generated serialization code for enum now reads/writes
-/// the tag byte directly to the wire. This avoids calling into reflective code.
+/// Serializer for enums.
 /// </summary>
 /// <typeparam name="T"></typeparam>
 public readonly struct Enum<T> : IReadWrite<T>
     where T : struct, Enum
 {
-    private static readonly ulong NumVariants;
+    /// <summary>
+    /// Map from tag -> value, implemented as an array.
+    /// Note: the [Type] macro rejects enums with explicitly set values (see Codegen.Tests),
+    /// so this array is guaranteed to be continuous and indexed starting from 0.
+    /// </summary>
+    private static readonly T[] TagToValue = Enum.GetValues(typeof(T)).Cast<T>().ToArray();
 
-    static Enum()
+    public T Read(BinaryReader reader)
     {
-        NumVariants = (ulong)Enum.GetValues(typeof(T)).Length;
+        var tag = reader.ReadByte();
+        try
+        {
+            return TagToValue[tag];
+        }
+        catch
+        {
+            throw new ArgumentOutOfRangeException(
+                $"Tag {tag} is out of range of enum {typeof(T).Name}"
+            );
+        }
     }
 
-    private static T Validate(T value)
+    public void Write(BinaryWriter writer, T value)
     {
-        // Previously this was: `if (!Enum.IsDefined(typeof(T), value))`.
+        // Previously this was: `if (Enum.IsDefined(typeof(T), value))`.
         // This was quite expensive because:
         //   1. It uses reflection
         //   2. It allocates
@@ -144,24 +157,22 @@ private static T Validate(T value)
         // However, enum values are guaranteed to be sequential and zero based.
         // Hence we only ever need to do an upper bound check.
         // See `SpacetimeDB.Type.ParseEnum` for the syntax analysis.
-
-        // Later note: this STILL uses reflection. We've deprecated this class entirely
-        // because of this.
-        if (Convert.ToUInt64(value) >= NumVariants)
+        //
+        // Note: this actually still uses reflection and allocates.
+        // It's hard to figure out how to avoid this without custom-generating a writer for each enum type.
+        var tag = Convert.ToByte(value);
+        if (tag < TagToValue.Length)
+        {
+            writer.Write(tag);
+        }
+        else
         {
             throw new ArgumentOutOfRangeException(
-                nameof(value),
-                $"Value {value} is out of range of enum {typeof(T).Name}"
+                $"Value {value} is out of range for enum {typeof(T).Name}"
             );
         }
-        return value;
     }
 
-    public T Read(BinaryReader reader) => Validate((T)Enum.ToObject(typeof(T), reader.ReadByte()));
-
-    public void Write(BinaryWriter writer, T value) =>
-        writer.Write(Convert.ToByte(Validate(value)));
-
     public AlgebraicType GetAlgebraicType(ITypeRegistrar registrar) =>
         registrar.RegisterType<T>(
             (_) =>

crates/bindings-csharp/Codegen.Tests/fixtures/client/snapshots/Type#PublicTable.verified.cs

@@ -255,7 +255,7 @@ public bool Equals(PublicTable that)
         var ___eqConnectionIdField = this.ConnectionIdField.Equals(that.ConnectionIdField);
         var ___eqCustomStructField = this.CustomStructField.Equals(that.CustomStructField);
         var ___eqCustomClassField = this.CustomClassField.Equals(that.CustomClassField);
-        var ___eqCustomEnumField = this.CustomEnumField.Equals(that.CustomEnumField);
+        var ___eqCustomEnumField = this.CustomEnumField == that.CustomEnumField;
         var ___eqCustomTaggedEnumField =
             this.CustomTaggedEnumField == null
                 ? that.CustomTaggedEnumField == null

crates/bindings-csharp/Codegen.Tests/fixtures/diag/snapshots/Module#TestUniqueNotEquatable.verified.cs

@@ -279,7 +279,7 @@ public bool Equals(PublicTable that)
             this.CustomClassField == null
                 ? that.CustomClassField == null
                 : this.CustomClassField.Equals(that.CustomClassField);
-        var ___eqCustomEnumField = this.CustomEnumField.Equals(that.CustomEnumField);
+        var ___eqCustomEnumField = this.CustomEnumField == that.CustomEnumField;
         var ___eqCustomTaggedEnumField =
             this.CustomTaggedEnumField == null
                 ? that.CustomTaggedEnumField == null

crates/bindings-csharp/Codegen.Tests/fixtures/diag/snapshots/Type#TestUnsupportedType.verified.cs

@@ -152,7 +152,7 @@ public bool Equals(CustomNestedClass? that)
             this.NestedNullableClass == null
                 ? that.NestedNullableClass == null
                 : this.NestedNullableClass.Equals(that.NestedNullableClass);
-        var ___eqNestedEnum = this.NestedEnum.Equals(that.NestedEnum);
+        var ___eqNestedEnum = this.NestedEnum == that.NestedEnum;
         var ___eqNestedNullableEnum = this.NestedNullableEnum.Equals(that.NestedNullableEnum);
         var ___eqNestedTaggedEnum =
             this.NestedTaggedEnum == null