diff --git a/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.GeneratesExpectedSourceCode_sourceFileName=FlattenedNestedUnionWithPartition.00.verified.cs b/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.GeneratesExpectedSourceCode_sourceFileName=FlattenedNestedUnionWithPartition.00.verified.cs
new file mode 100644
index 0000000..653d24f
--- /dev/null
+++ b/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.GeneratesExpectedSourceCode_sourceFileName=FlattenedNestedUnionWithPartition.00.verified.cs
@@ -0,0 +1,23 @@
+//HintName: DiscriminatedUnionAttribute.g.cs
+//
+#nullable enable
+
+namespace Funcky
+{
+ [global::System.Diagnostics.Conditional("Funcky_DiscriminatedUnion")]
+ [global::System.AttributeUsage(global::System.AttributeTargets.Class)]
+ internal sealed class DiscriminatedUnionAttribute : global::System.Attribute
+ {
+ /// Allow only consumers in the same assembly to use the exhaustive Match and Switch methods.
+ public bool NonExhaustive { get; set; }
+
+ /// Generates exhaustive Match and Switch methods for the entire type hierarchy.
+ public bool Flatten { get; set; }
+
+ /// If a specialized partition extension method for IEnumerable should be generated. Defaults to .
+ public bool GeneratePartitionExtension { get; set; }
+
+ /// Customized the generic type name used for the result in the generated Match methods. Defaults to TResult.
+ public string? MatchResultTypeName { get; set; }
+ }
+}
diff --git a/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.GeneratesExpectedSourceCode_sourceFileName=FlattenedNestedUnionWithPartition.01.verified.cs b/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.GeneratesExpectedSourceCode_sourceFileName=FlattenedNestedUnionWithPartition.01.verified.cs
new file mode 100644
index 0000000..958c82b
--- /dev/null
+++ b/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.GeneratesExpectedSourceCode_sourceFileName=FlattenedNestedUnionWithPartition.01.verified.cs
@@ -0,0 +1,65 @@
+//HintName: DiscriminatedUnionGenerator.g.cs
+//
+#nullable enable
+
+namespace Funcky.DiscriminatedUnion.Test
+{
+ partial record FlattenedNestedUnionWithPartition
+ {
+ [global::System.CodeDom.Compiler.GeneratedCode("Funcky.DiscriminatedUnion.SourceGeneration", "1.2.0.0")]
+ public abstract TResult Match(global::System.Func keyword, global::System.Func integer);
+
+ [global::System.CodeDom.Compiler.GeneratedCode("Funcky.DiscriminatedUnion.SourceGeneration", "1.2.0.0")]
+ public abstract void Switch(global::System.Action keyword, global::System.Action integer);
+
+ partial record Keyword
+ {
+ [global::System.CodeDom.Compiler.GeneratedCode("Funcky.DiscriminatedUnion.SourceGeneration", "1.2.0.0")]
+ public override TResult Match(global::System.Func keyword, global::System.Func integer) => keyword(this);
+
+ [global::System.CodeDom.Compiler.GeneratedCode("Funcky.DiscriminatedUnion.SourceGeneration", "1.2.0.0")]
+ public override void Switch(global::System.Action keyword, global::System.Action integer) => keyword(this);
+ }
+
+ partial record Literal
+ {
+ partial record Number
+ {
+ partial record Integer
+ {
+ [global::System.CodeDom.Compiler.GeneratedCode("Funcky.DiscriminatedUnion.SourceGeneration", "1.2.0.0")]
+ public override TResult Match(global::System.Func keyword, global::System.Func integer) => integer(this);
+
+ [global::System.CodeDom.Compiler.GeneratedCode("Funcky.DiscriminatedUnion.SourceGeneration", "1.2.0.0")]
+ public override void Switch(global::System.Action keyword, global::System.Action integer) => integer(this);
+ }
+ }
+ }
+ }
+
+ [global::System.CodeDom.Compiler.GeneratedCode("Funcky.DiscriminatedUnion.SourceGeneration", "1.2.0.0")]
+ public static partial class FlattenedNestedUnionWithPartitionEnumerableExtensions
+ {
+ public static (global::System.Collections.Generic.IReadOnlyList keyword, global::System.Collections.Generic.IReadOnlyList integer) Partition(this global::System.Collections.Generic.IEnumerable source)
+ {
+ var keywordItems = new global::System.Collections.Generic.List();
+ var integerItems = new global::System.Collections.Generic.List();
+ foreach (var item in source)
+ {
+ item.Switch(keyword: keywordItems.Add, integer: integerItems.Add);
+ }
+ return (keywordItems.AsReadOnly(), integerItems.AsReadOnly());
+ }
+
+ public static TResult Partition(this global::System.Collections.Generic.IEnumerable source, global::System.Func, global::System.Collections.Generic.IReadOnlyList, TResult> resultSelector)
+ {
+ var keywordItems = new global::System.Collections.Generic.List();
+ var integerItems = new global::System.Collections.Generic.List();
+ foreach (var item in source)
+ {
+ item.Switch(keyword: keywordItems.Add, integer: integerItems.Add);
+ }
+ return resultSelector(keywordItems.AsReadOnly(), integerItems.AsReadOnly());
+ }
+ }
+}
diff --git a/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.cs b/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.cs
index ca18156..c90b294 100644
--- a/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.cs
+++ b/Funcky.DiscriminatedUnion.Test/SourceGeneratorTest.cs
@@ -22,6 +22,7 @@ public sealed class SourceGeneratorTest
[InlineData("JsonPolymorphic")]
[InlineData("UnionWithPartitionUsage")]
[InlineData("FlattenedUnionWithPartition")]
+ [InlineData("FlattenedNestedUnionWithPartition")]
public async Task GeneratesExpectedSourceCode(string sourceFileName) => await Verify(sourceFileName);
[Fact]
diff --git a/Funcky.DiscriminatedUnion.Test/Sources/FlattenedNestedUnionWithPartition.cs b/Funcky.DiscriminatedUnion.Test/Sources/FlattenedNestedUnionWithPartition.cs
new file mode 100644
index 0000000..08da131
--- /dev/null
+++ b/Funcky.DiscriminatedUnion.Test/Sources/FlattenedNestedUnionWithPartition.cs
@@ -0,0 +1,23 @@
+namespace Funcky.DiscriminatedUnion.Test;
+
+[DiscriminatedUnion(Flatten = true, GeneratePartitionExtension = true)]
+public abstract partial record FlattenedNestedUnionWithPartition
+{
+ public sealed partial record Keyword(string Value) : FlattenedNestedUnionWithPartition;
+
+ public abstract partial record Literal : FlattenedNestedUnionWithPartition
+ {
+ public abstract partial record Number : Literal
+ {
+ public sealed partial record Integer(int Value) : Number;
+ }
+ }
+}
+
+public static class FlattenedNestedUnionWithPartitionTest
+{
+ public static void Test(FlattenedNestedUnionWithPartition[] items)
+ {
+ var (keywords, integers) = items.Partition();
+ }
+}