LightweightIocContainer/LightweightIocContainer.FactoryGenerator/FactoryGenerator.cs

// Author: Simon.Gockner
// Created: 2025-12-01
// Copyright(c) 2025 SimonG. All Rights Reserved.

using System.Collections.Immutable;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace LightweightIocContainer.FactoryGenerator;

[Generator]
public class FactoryGenerator : IIncrementalGenerator
{
    private const string EXTENSION_CLASS_NAME = "FactoryExtensions";
    private const string BUILDER_CLASS_NAME = "GeneratedFactoryBuilder";
    
    private const string GENERATED_FILE_HEADER = "//---GENERATED by FactoryGenerator! DO NOT EDIT!---";
    private const string INDENT = "    ";
    
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        string? classNamespace = typeof(FactoryGenerator).Namespace;
        context.RegisterPostInitializationOutput(c => c.AddSource($"{EXTENSION_CLASS_NAME}.g.cs", GenerateFactoryExtensionsClass(classNamespace, EXTENSION_CLASS_NAME)));

        IncrementalValuesProvider<ITypeSymbol?> syntaxProvider = context.SyntaxProvider.CreateSyntaxProvider(IsCallToGenerateFactory, GetTypeArgument);
        context.RegisterSourceOutput(syntaxProvider.Collect(), GenerateTypeDependentClasses);
    }

    private string GenerateFactoryExtensionsClass(string? classNamespace, string className)
    {
        StringBuilder stringBuilder = new();
        
        stringBuilder.AppendLine(GENERATED_FILE_HEADER);
        stringBuilder.AppendLine();

        stringBuilder.AppendLine("using LightweightIocContainer.Interfaces.Registrations;");
        stringBuilder.AppendLine("using LightweightIocContainer.Interfaces.Registrations.Fluent;");
        stringBuilder.AppendLine("using LightweightIocContainer.Registrations;");
        stringBuilder.AppendLine();
        
        if (!string.IsNullOrEmpty(classNamespace))
        {
            stringBuilder.AppendLine($"namespace {classNamespace};");
            stringBuilder.AppendLine();
        }
        
        stringBuilder.AppendLine($"public static class {className}");;
        stringBuilder.AppendLine("{");
        
        stringBuilder.AppendLine($"{INDENT}public static IRegistrationBase WithGeneratedFactory<TFactory>(this IRegistrationBase registration)");
        stringBuilder.AppendLine($"{INDENT}{{");
        stringBuilder.AppendLine($"{INDENT}{INDENT}TFactory factory = Builder.Create<TFactory>();");
        stringBuilder.AppendLine();
        stringBuilder.AppendLine($"{INDENT}{INDENT}if (registration is not RegistrationBase registrationBase)");
        stringBuilder.AppendLine($"{INDENT}{INDENT}{INDENT}throw new InvalidOperationException(\"The registration must be of type RegistrationBase to add a generated factory.\");");
        stringBuilder.AppendLine();
        stringBuilder.AppendLine($"{INDENT}{INDENT}registrationBase.AddGeneratedFactory<TFactory>(factory);");
        stringBuilder.AppendLine();
        stringBuilder.AppendLine($"{INDENT}{INDENT}return registrationBase;");
        stringBuilder.AppendLine($"{INDENT}}}");
        
        stringBuilder.AppendLine("}");
        
        return stringBuilder.ToString();
    }

    private bool IsCallToGenerateFactory(SyntaxNode node, CancellationToken cancellationToken)
    {
        if (!node.IsKind(SyntaxKind.GenericName) || node is not GenericNameSyntax genericNameSyntax)
            return false;

        if (!genericNameSyntax.ToString().StartsWith("WithGeneratedFactory"))
            return false;

        if (genericNameSyntax.TypeArgumentList.Arguments[0] is not IdentifierNameSyntax)
            return false;

        return true;
    }
    
    private ITypeSymbol? GetTypeArgument(GeneratorSyntaxContext syntaxContext, CancellationToken cancellationToken)
    {
        if (syntaxContext.Node is not GenericNameSyntax genericNameSyntax)
            return null;
        
        if (genericNameSyntax.TypeArgumentList.Arguments[0] is not IdentifierNameSyntax identifierNameSyntax)
            return null;

        if (syntaxContext.SemanticModel.GetSymbolInfo(identifierNameSyntax).Symbol is not ITypeSymbol typeSymbol)
            return null;

        return typeSymbol;
    }
    
    private void GenerateTypeDependentClasses(SourceProductionContext context, ImmutableArray<ITypeSymbol?> types)
    {
        string? classNamespace = typeof(FactoryGenerator).Namespace;
        context.AddSource($"{BUILDER_CLASS_NAME}.g.cs", GenerateBuilderClassSourceCode(classNamespace, types));
    }

    private string GenerateBuilderClassSourceCode(string? classNamespace, ImmutableArray<ITypeSymbol?> types)
    {
        StringBuilder stringBuilder = new();

        stringBuilder.AppendLine(GENERATED_FILE_HEADER);
        stringBuilder.AppendLine();

        foreach (string typeNamespace in GetNamespacesOfTypes(types)) 
            stringBuilder.AppendLine($"using {typeNamespace};");
        
        stringBuilder.AppendLine();

        if (classNamespace is not null)
        {
            stringBuilder.AppendLine($"namespace {classNamespace};");
            stringBuilder.AppendLine();
        }
        
        stringBuilder.AppendLine("public static class Builder");
        stringBuilder.AppendLine("{");
        stringBuilder.AppendLine($"{INDENT}public static TFactory Create<TFactory>()");
        stringBuilder.AppendLine($"{INDENT}{{");

        foreach (ITypeSymbol? type in types)
        {
            if (type is null)
                continue;
            
            stringBuilder.AppendLine($"{INDENT}{INDENT}if (typeof(TFactory) == typeof({type.Name}))");
            stringBuilder.AppendLine($"{INDENT}{INDENT}{{");
            stringBuilder.AppendLine($"{INDENT}{INDENT}{INDENT}return (TFactory) (object) new Generated{type.Name}();");
            stringBuilder.AppendLine($"{INDENT}{INDENT}}}");
            stringBuilder.AppendLine();
        }
        
        stringBuilder.AppendLine($"{INDENT}{INDENT}throw new Exception(\"Invalid type.\");");
        
        stringBuilder.AppendLine($"{INDENT}}}");
        stringBuilder.AppendLine("}");

        return stringBuilder.ToString();
    }

    private IEnumerable<string> GetNamespacesOfTypes(ImmutableArray<ITypeSymbol?> types) =>
        types.OfType<ITypeSymbol>()
            .Select(s => s.ContainingNamespace.IsGlobalNamespace ? null : s.ContainingNamespace.ToString())
            .OfType<string>()
            .Distinct();
}