Union visitors (#309)

* Union visitors

* Fix up stuff

* Fix UT

src/FlatSharp.Compiler/SchemaModel/UnionSchemaModel.cs

@@ -16,6 +16,7 @@
 
 using System.Linq;
 using FlatSharp.Compiler.Schema;
+using FlatSharp.CodeGen;
 
 namespace FlatSharp.Compiler.SchemaModel;
 
@@ -176,6 +177,7 @@ public class UnionSchemaModel : BaseSchemaModel
             }
 
             // Switch methods.
+            this.WriteAcceptMethod(writer, innerTypes);
             this.WriteSwitchMethod(writer, true, true, innerTypes);
             this.WriteSwitchMethod(writer, true, false, innerTypes);
             this.WriteSwitchMethod(writer, false, true, innerTypes);
@@ -183,6 +185,35 @@ public class UnionSchemaModel : BaseSchemaModel
         }
     }
 
+    private void WriteAcceptMethod(
+        CodeWriter writer,
+        List<(string resolvedType, EnumVal value)> components)
+    {
+        string visitorBaseType = $"IFlatBufferUnionVisitor<TReturn, {string.Join(", ", components.Select(x => x.resolvedType))}>";
+
+        writer.AppendSummaryComment("A convenience interface for implementing a visitor.");
+        writer.AppendLine($"public interface Visitor<TReturn> : {visitorBaseType} {{ }}");
+
+        writer.AppendSummaryComment("Accepts a visitor into this FlatBufferUnion.");
+        writer.AppendLine($"public TReturn Accept<TVisitor, TReturn>(TVisitor visitor)");
+        writer.AppendLine($"   where TVisitor : {visitorBaseType}");
+        using (writer.WithBlock())
+        {
+            writer.AppendLine("var disc = this.Discriminator;");
+            writer.AppendLine("switch (disc)");
+            using (writer.WithBlock())
+            {
+                foreach (var item in components)
+                {
+                    long index = item.value.Value;
+                    writer.AppendLine($"case {index}: return visitor.Visit(({item.resolvedType})this.value);");
+                }
+
+                writer.AppendLine($"default: throw new {typeof(InvalidOperationException).GetCompilableTypeName()}(\"Unexpected discriminator: \" + disc);");
+            }
+        }
+    }
+
     private void WriteSwitchMethod(
         CodeWriter writer,
         bool hasReturn,

src/FlatSharp.Runtime/UnionTypes.tt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 James Courtney
+ * Copyright 2022 James Courtney
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -32,6 +32,29 @@ public interface IFlatBufferUnion
     byte Discriminator { get; }
 }
 
+<#
+    for (int i = 1; i <= numGenerics; ++i)
+    {
+#>
+
+public interface IFlatBufferUnionVisitor<TReturn, <#= string.Join(", ", Enumerable.Range(1, i).Select(x => $"T{x}")) #>>
+{
+<#
+        for (int j = 1; j <= i; ++j)
+        {
+#>
+    
+    TReturn Visit(T<#= j #> item);
+
+<#
+        }
+#>
+
+}
+
+<#
+    }
+#>
 
 <#
     for (int i = 0; i < numGenerics; ++i)
@@ -47,6 +70,9 @@ public interface IFlatBufferUnion
         <# foreach (var genericType in range) { #>
             T<#= genericType #> Item<#= genericType #> { get; }
         <# } #>
+
+            //TReturn Accept<TVisitor, TReturn>(TVisitor visitor)
+            //    where TVisitor : IFlatBufferUnionVisitor<TReturn, <#= genericList #>>;
         }
 
         [ExcludeFromCodeCoverage]
@@ -110,6 +136,21 @@ public interface IFlatBufferUnion
             <# 
             } 
             #>
+            
+
+            public TReturn Accept<TVisitor, TReturn>(TVisitor visitor)
+                where TVisitor : IFlatBufferUnionVisitor<TReturn, <#= genericList #>>
+            {
+                switch (this.discriminator)
+                {
+                <# foreach (var genericType in range) { #>
+                    case <#= genericType #>:
+                        return visitor.Visit((T<#= genericType #>)this.value);
+                <# } #>
+                    default:
+                        throw new InvalidOperationException("Unexpected discriminator: " + this.discriminator);
+                }
+            }
 
             public void Switch(
                 System.Action defaultCase,

src/Tests/FlatSharpCompilerTests/UnionTests.cs

@@ -49,18 +49,20 @@ public class UnionTests
         string[] expectedAliases = new[] { "First", "B", "Foobar_C" };
 
         // Validate nested enum
-        Type nestedEnum = unionType.GetNestedTypes().Single();
+        Type nestedEnum = unionType.GetNestedTypes().Where(x => x.IsEnum).Single();
         Assert.True(nestedEnum.IsEnum);
         Assert.Equal("ItemKind", nestedEnum.Name);
         Assert.Equal(typeof(byte), Enum.GetUnderlyingType(nestedEnum));
         Assert.Equal(types.Length + 1, Enum.GetValues(nestedEnum).Length);
-
         Assert.Equal("NONE", Enum.GetName(nestedEnum, (byte)0));
+
         for (int i = 0; i < types.Length; ++i)
         {
             Assert.Equal(expectedAliases[i], Enum.GetName(nestedEnum, (byte)(i + 1)));
         }
 
+        Type nestedVistior = unionType.GetNestedTypes().Where(x => x.IsInterface).Single();
+
         // Custom union defines ctors for all input types.
         foreach (var type in types)
         {

src/Tests/FlatSharpEndToEndTests/FlatSharpEndToEndTests.csproj

@@ -40,6 +40,8 @@
     <Compile Include="ValueStructs\ValueStructs.fbs.generated.cs" />
     <Compile Remove="Documentation\Documentation.fbs.generated.cs" />
     <Compile Include="Documentation\Documentation.fbs.generated.cs" />
+    <Compile Remove="Unions\Unions.fbs.generated.cs" />
+    <Compile Include="Unions\Unions.fbs.generated.cs" />
   </ItemGroup>
 
   <Target Name="FBS" BeforeTargets="CoreCompile">
@@ -50,6 +52,7 @@
       <InputFile Include="Grpc\Grpc.fbs" />
       <InputFile Include="ValueStructs\ValueStructs.fbs" />
       <InputFile Include="Documentation\Documentation.fbs" />
+      <InputFile Include="Unions\Unions.fbs" />
     </ItemGroup>
     <Exec Command="dotnet $(FlatSharpCompilerDll) --nullable-warnings true -i %(InputFile.Identity) -o %(InputFile.RelativeDir)" />
   </Target>

src/Tests/FlatSharpEndToEndTests/Unions/Unions.fbs

@@ -0,0 +1,17 @@
+
+attribute "fs_serializer";
+attribute "fs_sharedString";
+attribute "fs_rpcInterface";
+
+namespace FlatSharpEndToEndTests.Unions;
+
+struct A { V : uint; }
+struct B { V : uint; }
+struct C { V : uint; }
+struct D { V : uint; }
+
+union MyUnion { A, B, C, D }
+
+table Container (fs_serializer) {
+    Value : [MyUnion];
+}

src/Tests/FlatSharpEndToEndTests/Unions/UnionsTestCases.cs

@@ -0,0 +1,163 @@
+/*
+ * Copyright 2022 James Courtney
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+namespace FlatSharpEndToEndTests.Unions;
+
+public class UnionsTestCases
+{
+    [Fact]
+    public void Union_Accept_Works()
+    {
+        var c = this.Setup();
+
+        UnionVisitor visitor = new();
+
+        Assert.Equal(typeof(A), c.Value[0].Accept<UnionVisitor, Type>(visitor));
+        Assert.Equal(typeof(B), c.Value[1].Accept<UnionVisitor, Type>(visitor));
+        Assert.Equal(typeof(C), c.Value[2].Accept<UnionVisitor, Type>(visitor));
+        Assert.Equal(typeof(D), c.Value[3].Accept<UnionVisitor, Type>(visitor));
+    }
+
+    [Fact]
+    public void Union_Switch_Func_Works()
+    {
+        static Type? CallSwitch(MyUnion union)
+        {
+            return union.Switch(
+                caseDefault: () => null,
+                caseA: a => typeof(A),
+                caseB: b => typeof(B),
+                caseC: c => typeof(C),
+                caseD: d => typeof(D));
+        }
+
+        var c = this.Setup();
+
+        Assert.Equal(typeof(A), CallSwitch(c.Value[0]));
+        Assert.Equal(typeof(B), CallSwitch(c.Value[1]));
+        Assert.Equal(typeof(C), CallSwitch(c.Value[2]));
+        Assert.Equal(typeof(D), CallSwitch(c.Value[3]));
+    }
+
+    [Fact]
+    public void Union_Switch_Func_WithState_Works()
+    {
+        static Type? CallSwitch(MyUnion union)
+        {
+            string? state = null;
+            Type? type = union.Switch(
+                "foobar",
+                caseDefault: (s) => { state = s; return (Type)null; },
+                caseA: (s, a) => { state = s; return typeof(A); },
+                caseB: (s, b) => { state = s; return typeof(B); },
+                caseC: (s, c) => { state = s; return typeof(C); },
+                caseD: (s, d) => { state = s; return typeof(D); });
+
+            Assert.Equal("foobar", state);
+            return type;
+        }
+
+        var c = this.Setup();
+
+        Assert.Equal(typeof(A), CallSwitch(c.Value[0]));
+        Assert.Equal(typeof(B), CallSwitch(c.Value[1]));
+        Assert.Equal(typeof(C), CallSwitch(c.Value[2]));
+        Assert.Equal(typeof(D), CallSwitch(c.Value[3]));
+    }
+
+    [Fact]
+    public void Union_Switch_Action_Works()
+    {
+        static Type? CallSwitch(MyUnion union)
+        {
+            Type? value = null;
+
+            union.Switch(
+                caseDefault: () => { },
+                caseA: a => { value = typeof(A); },
+                caseB: b => { value = typeof(B); },
+                caseC: c => { value = typeof(C); },
+                caseD: d => { value = typeof(D); });
+
+            return value;
+        }
+
+        var c = this.Setup();
+
+        Assert.Equal(typeof(A), CallSwitch(c.Value[0]));
+        Assert.Equal(typeof(B), CallSwitch(c.Value[1]));
+        Assert.Equal(typeof(C), CallSwitch(c.Value[2]));
+        Assert.Equal(typeof(D), CallSwitch(c.Value[3]));
+    }
+
+    [Fact]
+    public void Union_Switch_Action_WithState_Works()
+    {
+        static Type? CallSwitch(MyUnion union)
+        {
+            string? state = null;
+            Type? type = null;
+
+            union.Switch(
+                "foobar",
+                caseDefault: (s) => { state = s; },
+                caseA: (s, a) => { state = s; type = typeof(A); },
+                caseB: (s, b) => { state = s; type = typeof(B); },
+                caseC: (s, c) => { state = s; type = typeof(C); },
+                caseD: (s, d) => { state = s; type = typeof(D); });
+
+            Assert.Equal("foobar", state);
+            return type;
+        }
+
+        var c = this.Setup();
+
+        Assert.Equal(typeof(A), CallSwitch(c.Value[0]));
+        Assert.Equal(typeof(B), CallSwitch(c.Value[1]));
+        Assert.Equal(typeof(C), CallSwitch(c.Value[2]));
+        Assert.Equal(typeof(D), CallSwitch(c.Value[3]));
+    }
+
+    private Container Setup()
+    {
+        Container c = new Container
+        {
+            Value = new MyUnion[]
+            {
+                new MyUnion(new A()),
+                new MyUnion(new B()),
+                new MyUnion(new C()),
+                new MyUnion(new D()),
+            }
+        };
+
+        byte[] buffer = new byte[Container.Serializer.GetMaxSize(c)];
+        Container.Serializer.Write(buffer, c);
+
+        return Container.Serializer.Parse(buffer);
+    }
+
+    private struct UnionVisitor : MyUnion.Visitor<Type>
+    {
+        public Type Visit(A item) => typeof(A);
+
+        public Type Visit(B item) => typeof(B);
+
+        public Type Visit(C item) => typeof(C);
+
+        public Type Visit(D item) => typeof(D);
+    }
+}

src/common.props

@@ -10,7 +10,7 @@
   <PropertyGroup>
     <GeneratePackageOnBuild>true</GeneratePackageOnBuild>
     <Version>6.3.2</Version>
-    <PackageVersion>6.3.2</PackageVersion>
+    <PackageVersion>6.3.3</PackageVersion>
     <AssemblyVersion>$(Version)</AssemblyVersion>
     <Authors>James Courtney</Authors>
     <Description>FlatSharp is a fast, idiomatic implementation of the FlatBuffer binary format.</Description>