From 5bf8b7a4c12e50c706b3956fd956e9374450de9d Mon Sep 17 00:00:00 2001
From: James Duley <jagduley@gmail.com>
Date: Tue, 17 Oct 2023 13:05:08 +1300
Subject: [PATCH] Add Match method to Union types

---
 .../SchemaModel/ReferenceUnionSchemaModel.cs  | 28 ++++++++++++++++++-
 .../SchemaModel/ValueUnionSchemaModel.cs      | 27 ++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/src/FlatSharp.Compiler/SchemaModel/ReferenceUnionSchemaModel.cs b/src/FlatSharp.Compiler/SchemaModel/ReferenceUnionSchemaModel.cs
index acb3971e..77d0f318 100644
--- a/src/FlatSharp.Compiler/SchemaModel/ReferenceUnionSchemaModel.cs
+++ b/src/FlatSharp.Compiler/SchemaModel/ReferenceUnionSchemaModel.cs
@@ -107,6 +107,7 @@ protected override void OnWriteCode(CodeWriter writer, CompileContext context)
             this.WriteDefaultConstructor(writer);
             this.WriteReturnToPool(writer);
             this.WriteAcceptMethod(writer, innerTypes);
+            this.WriteMatchMethod(writer, innerTypes);
         }
     }
 
@@ -187,7 +188,32 @@ private void WriteAcceptMethod(
                 foreach (var item in components)
                 {
                     long index = item.value.Value;
-                    writer.AppendLine($"case {index}: return visitor.Visit(this.value_{item.value.Value});");
+                    writer.AppendLine($"case {index}: return visitor.Visit(this.value_{index});");
+                }
+
+                writer.AppendLine($"default: throw new {typeof(InvalidOperationException).GetCompilableTypeName()}(\"Unexpected discriminator: \" + disc);");
+            }
+        }
+    }
+
+    private void WriteMatchMethod(
+        CodeWriter writer,
+        List<(string resolvedType, EnumVal value, Type? propertyType)> components)
+    {
+        string parameters = string.Join(", ", components.Select((x, i) => $"Func<{x.resolvedType}, TReturn> f{i + 1}"));
+
+        writer.AppendSummaryComment("Matches this FlatBufferUnion with a function for each value.");
+        writer.AppendLine($"public TReturn Match<TReturn>({parameters})");
+        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 f{index}(this.value_{index});");
                 }
 
                 writer.AppendLine($"default: throw new {typeof(InvalidOperationException).GetCompilableTypeName()}(\"Unexpected discriminator: \" + disc);");
diff --git a/src/FlatSharp.Compiler/SchemaModel/ValueUnionSchemaModel.cs b/src/FlatSharp.Compiler/SchemaModel/ValueUnionSchemaModel.cs
index fa2c6722..358a7679 100644
--- a/src/FlatSharp.Compiler/SchemaModel/ValueUnionSchemaModel.cs
+++ b/src/FlatSharp.Compiler/SchemaModel/ValueUnionSchemaModel.cs
@@ -230,6 +230,7 @@ protected override void OnWriteCode(CodeWriter writer, CompileContext context)
             }
 
             this.WriteAcceptMethod(writer, innerTypes);
+            this.WriteMatchMethod(writer, innerTypes);
         }
     }
 
@@ -262,6 +263,32 @@ private void WriteAcceptMethod(
         }
     }
 
+    private void WriteMatchMethod(
+        CodeWriter writer,
+        List<(string resolvedType, EnumVal value, int? size)> components)
+    {
+        string parameters = string.Join(", ", components.Select(x => $"Func<{x.resolvedType}, TReturn> f{x.value.Value}"));
+
+        writer.AppendSummaryComment("Matches this FlatBufferUnion with a function for each value.");
+        writer.AppendLine($"public TReturn Match<TReturn>({parameters})");
+        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 f{index}(this.UncheckedGetItem{index}());");
+                }
+
+                writer.AppendLine($"default: throw new {typeof(InvalidOperationException).GetCompilableTypeName()}(\"Unexpected discriminator: \" + disc);");
+            }
+        }
+    }
+
+
     private void WriteUncheckedGetItemMethod(CodeWriter writer, string resolvedType, EnumVal unionValue, Type? propertyType, bool generateUnsafeItems)
     {
         if (propertyType?.IsValueType == true && generateUnsafeItems)