Allow source generated results and unions inside partial classes

Funzo.SourceGenerators.Test/ResultGeneratorTests.cs

@@ -66,7 +66,7 @@ static TwoUnionResult process(int value)
     [Fact]
     public void Generates_Nested_Results()
     {
-        InnerClass.InnerResult result = "test";
+        InnerClass.ReallyInnerClass.InnerResult result = "test";
 
         Assert.True(result.IsErr(out _));
     }
@@ -92,9 +92,11 @@ public partial class ClonedOk;
 [Result<ClonedOk, ClonedErr>]
 public partial class CheapClone;
 
-
 public static partial class InnerClass
 {
-    [Result<int, string>]
-    public partial class InnerResult;
+    public partial class ReallyInnerClass
+    {
+        [Result<int, string>]
+        public partial class InnerResult;
+    }
 }

Funzo.SourceGenerators.Test/UnionGeneratorTests.cs

@@ -47,7 +47,7 @@ public void Generates_Shared_Property()
     [Fact]
     public void Generates_Inner_Union()
     {
-        TestUnionInPartialClass.InnerUnion u = new TestUnionInPartialClass.A();
+        TestUnionInPartialClass.InnerClass.InnerUnion u = new TestUnionInPartialClass.A();
 
         u.Switch(a => { }, b => throw new InvalidOperationException());
     }
@@ -84,6 +84,9 @@ public static partial class TestUnionInPartialClass
     public record A;
     public record B;
 
-    [Union<A, B>]
-    public partial class InnerUnion;
+    public partial class InnerClass
+    {
+        [Union<A, B>]
+        public partial class InnerUnion;
+    }
 }

Funzo.SourceGenerators/Funzo.SourceGenerators.csproj

@@ -6,7 +6,6 @@
 		<Title>Funzo.Generators</Title>
 		<Description>Source generators for result and union types</Description>
 		<PackageTags>functional programming;option types;monad;union types;discriminated unions;result</PackageTags>
-		<GeneratePackageOnBuild>True</GeneratePackageOnBuild>
 		<LangVersion>preview</LangVersion>
 		<Nullable>enable</Nullable>
 		<GeneratePackageOnBuild>true</GeneratePackageOnBuild>

Funzo.SourceGenerators/Generators/SourceGeneratorBase.cs → Funzo.SourceGenerators/Generators/GeneratorBase.cs

@@ -2,10 +2,11 @@
 using Microsoft.CodeAnalysis;
 using System.Linq;
 
-namespace Funzo.SourceGenerators;
-internal abstract class SourceGeneratorBase
+namespace Funzo.SourceGenerators.Generators;
+
+internal abstract class GeneratorBase
 {
-    protected SourceGeneratorBase()
+    protected GeneratorBase()
     {
     }
 
@@ -16,5 +17,5 @@ protected void CreateDiagnosticError(SourceProductionContext context, Diagnostic
             DiagnosticSeverity.Error));
     }
 
-    internal abstract string? GetSource(SourceProductionContext context, SymbolWithAttribute symbolWithAttribute);
+    internal abstract string? GetSource(SourceProductionContext context, MarkedType type);
 }

Funzo.SourceGenerators/Generators/ResultSourceGenerator.cs

@@ -4,37 +4,39 @@
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using System;
+using System.Collections.Immutable;
 using System.Linq;
 using System.Text;
 
 namespace Funzo.SourceGenerators.Generators;
 
-internal class ResultSourceGenerator : SourceGeneratorBase
+internal class ResultSourceGenerator : GeneratorBase
 {
-    internal override string? GetSource(SourceProductionContext context, SymbolWithAttribute symbolWithAttribute)
+    internal override string? GetSource(SourceProductionContext context, MarkedType type)
     {
-        if (HasErrors(context, symbolWithAttribute))
+        if (HasErrors(context, type))
         {
             return null;
         }
 
-        var isSimpleResult = symbolWithAttribute.TypeArguments.Length == 1;
+        var is1ArityResult = type.AttributeTypeArguments.Length == 1;
 
-        Results.ResultGenerator generator = isSimpleResult
-            ? new Result1AritySourceGenerator(symbolWithAttribute)
-            : new Result2AritySourceGenerator(symbolWithAttribute);
+        ResultGenerator generator = is1ArityResult
+            ? new Result1AritySourceGenerator(type)
+            : new Result2AritySourceGenerator(type);
 
-        var containingType = symbolWithAttribute.Symbol.ContainingType?.Name;
+        var sb = new StringBuilder();
+        sb.AppendLine($@"// <auto-generated />
 
-        var sb = new StringBuilder($@"// <auto-generated />
-
-namespace {symbolWithAttribute.Symbol.ContainingNamespace.ToDisplayString()}
+namespace {type.Symbol.ContainingNamespace.ToDisplayString()}
 {{");
 
-        if (containingType is not null)
+        // Add all needed upper classes
+        foreach (var containingType in type.ContainerClasses)
         {
-            var isStatic = symbolWithAttribute.Symbol.ContainingType!.IsStatic ? "static " : "";
-            sb.AppendLine($"{isStatic}partial class {containingType} {{");
+            var isStatic = type.Symbol.ContainingType!.IsStatic ? "static " : "";
+
+            sb.AppendLine($"{isStatic}partial class {containingType.Name} {{");
         }
 
         sb.AppendLine($@"
@@ -49,44 +51,58 @@ namespace {symbolWithAttribute.Symbol.ContainingNamespace.ToDisplayString()}
     }}
 }}");
 
-        if (containingType is not null)
-        {
-            sb.AppendLine("}");
-        }
+        // Close all the braces opened for parent classes
+        sb.AppendLine(new string('}', type.ContainerClasses.Count));
 
         var src = sb.ToString();
 
         return src;
     }
 
-    private bool HasErrors(SourceProductionContext context, SymbolWithAttribute symbolWithAttribute)
+    private bool HasErrors(SourceProductionContext context, MarkedType type)
     {
-        var (ClassSymbol, _) = symbolWithAttribute;
+        var symbol = type.Symbol;
+        var typeArguments = type.AttributeTypeArguments;
 
-        if (!ClassSymbol.ContainingSymbol.Equals(ClassSymbol.ContainingNamespace, SymbolEqualityComparer.Default) && ClassSymbol.ContainingType is { } containerType
-           && containerType.DeclaringSyntaxReferences.Any(syntax =>
-               syntax.GetSyntax() is BaseTypeDeclarationSyntax declaration
-               && !declaration.Modifiers.Any(mod => mod.IsKind(SyntaxKind.PartialKeyword))))
+        if (!IsSymbolTopLevelOrInsidePartialClasses(symbol))
         {
-            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Result.TopLevelError, ClassSymbol);
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Result.TopLevelError, symbol);
             return true;
         }
 
-        if (ClassSymbol.BaseType is not null && ClassSymbol.BaseType.Name != "Object")
+        if (IsSymbolInheritingFromOtherClasses(symbol))
         {
-            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Result.WrongBaseType, ClassSymbol);
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Result.WrongBaseType, symbol);
             return true;
         }
 
-        foreach (var typeSymbol in symbolWithAttribute.TypeArguments)
+        if (IsAnyTypeArgumentObject(typeArguments))
         {
-            if (typeSymbol.Name == nameof(Object))
-            {
-                CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Result.ObjectNotValidType, ClassSymbol);
-                return true;
-            }
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Result.ObjectNotValidType, symbol);
+            return true;
+        }
+
+        if (AreTypeArgumentsTheSame(typeArguments))
+        {
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Result.RepeatedTypeSymbols, symbol);
+            return true;
         }
 
         return false;
     }
+
+    private static bool IsSymbolTopLevelOrInsidePartialClasses(INamedTypeSymbol symbol)
+        => !(!symbol.ContainingSymbol.Equals(symbol.ContainingNamespace, SymbolEqualityComparer.Default) && symbol.ContainingType is { } containerType
+           && containerType.DeclaringSyntaxReferences.Any(syntax =>
+               syntax.GetSyntax() is BaseTypeDeclarationSyntax declaration
+               && !declaration.Modifiers.Any(mod => mod.IsKind(SyntaxKind.PartialKeyword))));
+
+    private static bool IsSymbolInheritingFromOtherClasses(INamedTypeSymbol symbol)
+        => symbol.BaseType is not null && symbol.BaseType.Name != "Object";
+
+    private static bool IsAnyTypeArgumentObject(ImmutableArray<ITypeSymbol> typeArguments)
+        => typeArguments.Any(a => a.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "object");
+
+    private static bool AreTypeArgumentsTheSame(ImmutableArray<ITypeSymbol> typeArguments)
+        => typeArguments.Count() == 2 && SymbolEqualityComparer.Default.Equals(typeArguments.First(), typeArguments.Last());
 }

Funzo.SourceGenerators/Generators/Results/Result1AritySourceGenerator.cs

@@ -5,7 +5,7 @@
 namespace Funzo.SourceGenerators.Generators.Results;
 internal class Result1AritySourceGenerator : ResultGenerator
 {
-    internal Result1AritySourceGenerator(SymbolWithAttribute symbolWithAttribute) : base(symbolWithAttribute)
+    internal Result1AritySourceGenerator(MarkedType symbolWithAttribute) : base(symbolWithAttribute)
     {
     }
 

Funzo.SourceGenerators/Generators/Results/Result2AritySourceGenerator.cs

@@ -7,7 +7,7 @@
 namespace Funzo.SourceGenerators.Generators.Results;
 internal class Result2AritySourceGenerator : ResultGenerator
 {
-    internal Result2AritySourceGenerator(SymbolWithAttribute symbolWithAttribute) : base(symbolWithAttribute)
+    internal Result2AritySourceGenerator(MarkedType symbolWithAttribute) : base(symbolWithAttribute)
     {
     }
 

Funzo.SourceGenerators/Generators/Results/ResultGenerator.cs

@@ -5,13 +5,14 @@
 using System.Text;
 
 namespace Funzo.SourceGenerators.Generators.Results;
+
 internal abstract class ResultGenerator
 {
-    private readonly SymbolWithAttribute _symbolWithAttribute;
+    private readonly MarkedType _symbolWithAttribute;
     protected INamedTypeSymbol ClassSymbol => _symbolWithAttribute.Symbol;
-    protected ImmutableArray<ITypeSymbol> TypeArguments => _symbolWithAttribute.AttributeData.AttributeClass!.TypeArguments;
+    protected ImmutableArray<ITypeSymbol> TypeArguments => _symbolWithAttribute.AttributeTypeArguments;
 
-    protected ResultGenerator(SymbolWithAttribute symbolWithAttribute)
+    protected ResultGenerator(MarkedType symbolWithAttribute)
     {
         _symbolWithAttribute = symbolWithAttribute;
     }

Funzo.SourceGenerators/Generators/Unions/UnionGenerator.cs

@@ -0,0 +1,134 @@
+using Funzo.SourceGenerators.Helpers;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+using System;
+using System.Collections.Immutable;
+using System.Linq;
+using System.Text;
+
+namespace Funzo.SourceGenerators.Generators.Unions;
+
+internal class UnionGenerator : GeneratorBase
+{
+    internal override string? GetSource(SourceProductionContext context, MarkedType type)
+    {
+        if (HasErrors(context, type))
+        {
+            return null;
+        }
+
+        var commonProperties = GetCommonProperties(type.AttributeTypeArguments);
+
+        var classSymbol = type.Symbol;
+        var className = $"{classSymbol.Name}";
+        var typeArguments = type.AttributeTypeArguments;
+
+        StringBuilder source = new($@"// <auto-generated />
+
+namespace {classSymbol.ContainingNamespace.ToDisplayString()}
+{{");
+
+        // Add all needed upper classes
+        foreach (var containingType in type.ContainerClasses)
+        {
+            var isStatic = type.Symbol.ContainingType!.IsStatic ? "static " : "";
+
+            source.AppendLine($"{isStatic}partial class {containingType.Name} {{");
+        }
+
+        source.AppendLine($@"         partial class {className} : global::{FunzoAttributeSources.AttributeNamespace}.Union{typeArguments.OpenGenericPart()}
+    {{");
+
+        foreach (var typeArgument in typeArguments)
+        {
+            source.Append($@"
+        public {classSymbol.Name}({typeArgument.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)} _) : base(_) {{}}
+        public static implicit operator {className}({typeArgument.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)} _) => new {className}(_);
+");
+        }
+
+        foreach (var prop in commonProperties)
+        {
+            source.Append($@"
+        public {prop.Type.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)} {prop.Name} => Match({string.Join(",", Enumerable.Range(0, typeArguments.Length).Select(_ => $"x => x.{prop.Name}"))});
+");
+        }
+
+        // Close all the braces opened for parent classes
+        source.AppendLine(new string('}', type.ContainerClasses.Count));
+
+        source.Append(@"    }
+}");
+
+        var src = source.ToString();
+        return src;
+    }
+
+    private bool HasErrors(SourceProductionContext context, MarkedType type)
+    {
+        var symbol = type.Symbol;
+
+        if (!IsSymbolTopLevelOrInsidePartialClasses(symbol))
+        {
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Union.TopLevelError, symbol);
+            return true;
+        }
+
+        if (IsSymbolInheritingFromOtherClasses(symbol))
+        {
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Union.WrongBaseType, symbol);
+            return true;
+        }
+
+
+        if (IsAnyTypeArgumentObject(type.AttributeTypeArguments))
+        {
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Union.ObjectNotValidType, symbol);
+            return true;
+        }
+
+        if (IsAnyTypeArgumentAnInterface(type.AttributeTypeArguments))
+        {
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Union.InterfaceNotValidType, symbol);
+            return true;
+        }
+
+        if (IsAnyTypeRepeatedInUnion(type.AttributeTypeArguments))
+        {
+            CreateDiagnosticError(context, FunzoDiagnosticDescriptors.Union.RepeatedTypeSymbols, symbol);
+            return true;
+        }
+
+        return false;
+    }
+
+    private static ImmutableArray<IPropertySymbol> GetCommonProperties(ImmutableArray<ITypeSymbol> typeArguments)
+    {
+        var typeProperties = typeArguments.SelectMany(t => t.GetPublicProperties()).ToLookup(x => x.Name);
+
+        var commonProperties = typeProperties.Where(g => g.Count() == typeArguments.Length && g.All(s => s.Type.Equals(g.First().Type, SymbolEqualityComparer.IncludeNullability)))
+            .Select(x => x.First())
+            .ToImmutableArray();
+
+        return commonProperties;
+    }
+
+    private static bool IsSymbolTopLevelOrInsidePartialClasses(INamedTypeSymbol type)
+        => !(type.ContainingSymbol.Equals(type.ContainingNamespace, SymbolEqualityComparer.Default) && type.ContainingType is { } containerType
+            && containerType.DeclaringSyntaxReferences.Any(syntax =>
+                syntax.GetSyntax() is BaseTypeDeclarationSyntax declaration
+                && !declaration.Modifiers.Any(mod => mod.IsKind(SyntaxKind.PartialKeyword))));
+
+    private static bool IsSymbolInheritingFromOtherClasses(INamedTypeSymbol symbol)
+        => symbol.BaseType is not null && symbol.BaseType.Name != "Object";
+
+    private static bool IsAnyTypeArgumentObject(ImmutableArray<ITypeSymbol> typeArguments)
+        => typeArguments.Any(a => a.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "object");
+
+    private static bool IsAnyTypeArgumentAnInterface(ImmutableArray<ITypeSymbol> typeArguments)
+        => typeArguments.Any(t => t.TypeKind == TypeKind.Interface);
+
+    private static bool IsAnyTypeRepeatedInUnion(ImmutableArray<ITypeSymbol> typeArguments)
+        => typeArguments.ToLookup(x => x, SymbolEqualityComparer.Default).Any(g => g.Count() > 1);
+}