AnnotationReachabilityTrimmer.cs

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace Light.GuardClauses.SourceCodeTransformation;

/// <summary>
/// Removes the bundled support types that the generated file does not reference any longer. This step must
/// run after <see cref="CleanupStep.RemoveAnnotations" /> because only the merged tree shows which annotation
/// usages actually survive cleanup - the original source files still contain every removable usage.
/// </summary>
internal static class AnnotationReachabilityTrimmer
{
    private const string JetBrainsAnnotationsNamespace = "JetBrains.Annotations";
    private const string NotNullAttributeName = "NotNullAttribute";
    private const string ValidatedNotNullAttributeName = "ValidatedNotNullAttribute";

    private static readonly HashSet<string> BundledNamespaceNames = new (StringComparer.Ordinal)
    {
        JetBrainsAnnotationsNamespace,
        "System.Diagnostics.CodeAnalysis",
        "System.Runtime.CompilerServices",
    };

    public static CompilationUnitSyntax Trim(CompilationUnitSyntax compilationUnit, SourceFileMergeOptions options)
    {
        ArgumentNullException.ThrowIfNull(compilationUnit);
        ArgumentNullException.ThrowIfNull(options);

        // At this point in the pipeline the root does not belong to a syntax tree that reflects it, thus the
        // tree must be recreated with the export's parse options before anything can be bound. Creating the
        // tree clones the root, so all further work happens on the tree's root.
        var syntaxTree = CSharpSyntaxTree.Create(
            compilationUnit,
            SourceReachabilityAnalyzer.CreateParseOptions(options.TargetFramework)
        );
        var root = (CompilationUnitSyntax) syntaxTree.GetRoot();

        var catalog = BundledTypeCatalog.Create(root, options);
        var reachableTypes = catalog.Types.Count == 0 ?
            new () :
            DetermineReachableTypes(syntaxTree, root, catalog);

        root = RemoveUnreachableTypes(root, catalog, reachableTypes);
        return AdjustJetBrainsUsingDirectives(root, options, reachableTypes);
    }

    private static HashSet<BundledType> DetermineReachableTypes(
        SyntaxTree syntaxTree,
        CompilationUnitSyntax root,
        BundledTypeCatalog catalog
    )
    {
        var compilation = CSharpCompilation.Create(
            "Light.GuardClauses.SourceExportAnnotationReachability",
            [syntaxTree],
            SourceReachabilityAnalyzer.GetMetadataReferences(),
            new (
                OutputKind.DynamicallyLinkedLibrary,
                nullableContextOptions: NullableContextOptions.Enable,
                allowUnsafe: true
            )
        );
        var semanticModel = compilation.GetSemanticModel(syntaxTree, true);

        var typesBySymbol = new Dictionary<ISymbol, BundledType>(SymbolEqualityComparer.Default);
        foreach (var bundledType in catalog.Types)
        {
            if (GetDeclaredSymbol(semanticModel, bundledType.Declaration) is { } typeSymbol)
            {
                typesBySymbol[typeSymbol] = bundledType;
            }
        }

        var reachableTypes = new HashSet<BundledType>();
        var typesToScan = new Queue<BundledType>();

        // Reachability is established from attribute usages and using aliases. Any future bundled type must
        // be reachable through one of these two - a type that is referenced in another position (a base type,
        // a member signature, a typeof expression) would be trimmed although it is still used. The only cref
        // references among the bundled types are self-references that disappear with their own declaration.
        foreach (var attribute in root.DescendantNodes().OfType<AttributeSyntax>())
        {
            if (catalog.IsDeclaredInsideBundledType(attribute))
            {
                continue;
            }

            MarkAttributeType(attribute);
        }

        // A using alias names its target type, thus an alias that outlives its target does not compile even
        // when nothing else refers to the alias.
        foreach (var usingDirective in root.Usings)
        {
            if (usingDirective.Alias != null)
            {
                MarkAliasedType(usingDirective);
            }
        }

        while (typesToScan.TryDequeue(out var scannedType))
        {
            foreach (var attribute in scannedType.Declaration.DescendantNodes().OfType<AttributeSyntax>())
            {
                MarkAttributeType(attribute);
            }
        }

        return reachableTypes;

        void MarkAttributeType(AttributeSyntax attribute)
        {
            var attributeType = ResolveAttributeType(semanticModel, attribute);
            if (attributeType == null)
            {
                // A degraded binding must produce a slightly larger file, never one that fails to compile.
                foreach (var matchingType in catalog.FindByUnqualifiedName(attribute.Name))
                {
                    Mark(matchingType);
                }

                return;
            }

            if (typesBySymbol.TryGetValue(attributeType, out var bundledType))
            {
                Mark(bundledType);
            }
        }

        void MarkAliasedType(UsingDirectiveSyntax usingDirective)
        {
            var aliasedType = GetAttributeType(semanticModel.GetSymbolInfo(usingDirective.NamespaceOrType).Symbol);
            if (aliasedType is not null and not IErrorTypeSymbol)
            {
                if (typesBySymbol.TryGetValue(aliasedType, out var bundledType))
                {
                    Mark(bundledType);
                }

                return;
            }

            if (usingDirective.NamespaceOrType is NameSyntax aliasedName)
            {
                foreach (var matchingType in catalog.FindByUnqualifiedName(aliasedName))
                {
                    Mark(matchingType);
                }
            }
        }

        void Mark(BundledType bundledType)
        {
            if (reachableTypes.Add(bundledType))
            {
                typesToScan.Enqueue(bundledType);
            }
        }
    }

    private static INamedTypeSymbol? ResolveAttributeType(SemanticModel semanticModel, AttributeSyntax attribute)
    {
        var symbolInfo = semanticModel.GetSymbolInfo(attribute);
        var attributeType = GetAttributeType(symbolInfo.Symbol);
        if (attributeType == null)
        {
            foreach (var candidateSymbol in symbolInfo.CandidateSymbols)
            {
                attributeType = GetAttributeType(candidateSymbol);
                if (attributeType != null)
                {
                    break;
                }
            }
        }

        attributeType ??= GetAttributeType(semanticModel.GetSymbolInfo(attribute.Name).Symbol);
        return attributeType is null or IErrorTypeSymbol ? null : attributeType;
    }

    private static INamedTypeSymbol? GetAttributeType(ISymbol? symbol) =>
        symbol switch
        {
            IMethodSymbol methodSymbol => methodSymbol.ContainingType,
            INamedTypeSymbol namedTypeSymbol => namedTypeSymbol,
            _ => null,
        };

    private static CompilationUnitSyntax RemoveUnreachableTypes(
        CompilationUnitSyntax root,
        BundledTypeCatalog catalog,
        HashSet<BundledType> reachableTypes
    )
    {
        var nodesToRemove = new List<SyntaxNode>();
        foreach (var namespaceDeclaration in catalog.BundledNamespaces)
        {
            var unreachableDeclarations = new List<SyntaxNode>();
            var retainsMember = false;
            foreach (var member in namespaceDeclaration.Members)
            {
                // Anything that is not a bundled type counts as retained: only the catalog decides what may
                // be removed for lack of references.
                if (catalog.TryGetBundledType(member, out var bundledType) && !reachableTypes.Contains(bundledType))
                {
                    unreachableDeclarations.Add(member);
                }
                else
                {
                    retainsMember = true;
                }
            }

            if (retainsMember)
            {
                nodesToRemove.AddRange(unreachableDeclarations);
            }
            else
            {
                // Removing the namespace node also removes its license banner: the banner is leading trivia
                // of the namespace token.
                nodesToRemove.Add(namespaceDeclaration);
            }
        }

        foreach (var bundledType in catalog.TypesOutsideBundledNamespaces)
        {
            if (!reachableTypes.Contains(bundledType))
            {
                nodesToRemove.Add(bundledType.Declaration);
            }
        }

        return nodesToRemove.Count == 0 ?
            root :
            root.RemoveNodes(nodesToRemove, SyntaxRemoveOptions.KeepNoTrivia)!;
    }

    private static CompilationUnitSyntax AdjustJetBrainsUsingDirectives(
        CompilationUnitSyntax root,
        SourceFileMergeOptions options,
        HashSet<BundledType> reachableTypes
    )
    {
        // When the JetBrains annotations are not bundled, IncludeJetBrainsAnnotationsUsing keeps its meaning
        // so that consumers can supply the annotations from an externally referenced package.
        if (!options.IncludeJetBrainsAnnotations)
        {
            return root;
        }

        var retainsJetBrainsTypes = false;
        var retainsJetBrainsNotNullAttribute = false;
        foreach (var bundledType in reachableTypes)
        {
            if (bundledType.NamespaceName != JetBrainsAnnotationsNamespace)
            {
                continue;
            }

            retainsJetBrainsTypes = true;
            retainsJetBrainsNotNullAttribute |= bundledType.Name == NotNullAttributeName;
        }

        var directivesToRemove = new List<UsingDirectiveSyntax>();
        foreach (var usingDirective in root.Usings)
        {
            // The alias only exists to disambiguate the two NotNullAttribute declarations. Removing the
            // namespace while leaving its import in place would be a compile error, thus both decisions
            // are derived from the trimmed set together.
            if (usingDirective.Alias != null)
            {
                if (!retainsJetBrainsNotNullAttribute &&
                    usingDirective.Alias.Name.Identifier.ValueText == NotNullAttributeName)
                {
                    directivesToRemove.Add(usingDirective);
                }
            }
            else if (!retainsJetBrainsTypes &&
                     usingDirective.NamespaceOrType.ToString() == JetBrainsAnnotationsNamespace)
            {
                directivesToRemove.Add(usingDirective);
            }
        }

        return directivesToRemove.Count == 0 ?
            root :
            root.RemoveNodes(directivesToRemove, SyntaxRemoveOptions.KeepNoTrivia)!;
    }

    private static ISymbol? GetDeclaredSymbol(SemanticModel semanticModel, MemberDeclarationSyntax declaration) =>
        declaration switch
        {
            BaseTypeDeclarationSyntax typeDeclaration => semanticModel.GetDeclaredSymbol(typeDeclaration),
            DelegateDeclarationSyntax delegateDeclaration => semanticModel.GetDeclaredSymbol(delegateDeclaration),
            _ => null,
        };

    private sealed record BundledType(MemberDeclarationSyntax Declaration, string NamespaceName, string Name);

    /// <summary>
    /// Identifies the bundled support types in the merged tree. They cannot be derived from the mechanisms
    /// that brought them there - two of them arrive as always-processed source files that whitelist mode
    /// exempts from reachability analysis, the other two as literal blocks injected by the merger.
    /// </summary>
    private sealed class BundledTypeCatalog
    {
        private readonly List<NamespaceDeclarationSyntax> _bundledNamespaces = new ();
        private readonly List<BundledType> _types = new ();
        private readonly Dictionary<SyntaxNode, BundledType> _typesByDeclaration = new ();
        private readonly Dictionary<string, List<BundledType>> _typesByUnqualifiedName = new (StringComparer.Ordinal);
        private readonly List<BundledType> _typesOutsideBundledNamespaces = new ();

        public IReadOnlyList<BundledType> Types => _types;

        public IReadOnlyList<NamespaceDeclarationSyntax> BundledNamespaces => _bundledNamespaces;

        public IReadOnlyList<BundledType> TypesOutsideBundledNamespaces => _typesOutsideBundledNamespaces;

        public static BundledTypeCatalog Create(CompilationUnitSyntax root, SourceFileMergeOptions options)
        {
            var catalog = new BundledTypeCatalog();
            foreach (var namespaceDeclaration in root.Members.OfType<NamespaceDeclarationSyntax>())
            {
                var namespaceName = namespaceDeclaration.Name.ToString();
                if (BundledNamespaceNames.Contains(namespaceName))
                {
                    catalog._bundledNamespaces.Add(namespaceDeclaration);
                    foreach (var member in namespaceDeclaration.Members)
                    {
                        catalog.Add(member, namespaceName, false);
                    }
                }
                else if (namespaceName == options.BaseNamespace)
                {
                    foreach (var member in namespaceDeclaration.Members)
                    {
                        if (member is BaseTypeDeclarationSyntax
                            {
                                Identifier.ValueText: ValidatedNotNullAttributeName,
                            })
                        {
                            catalog.Add(member, namespaceName, true);
                        }
                    }
                }
            }

            return catalog;
        }

        public bool TryGetBundledType(
            SyntaxNode declaration,
            [MaybeNullWhen(false)] out BundledType bundledType
        ) =>
            _typesByDeclaration.TryGetValue(declaration, out bundledType);

        public bool IsDeclaredInsideBundledType(SyntaxNode node) =>
            node.Ancestors().Any(_typesByDeclaration.ContainsKey);

        public IReadOnlyList<BundledType> FindByUnqualifiedName(NameSyntax attributeName) =>
            _typesByUnqualifiedName.TryGetValue(
                CleanupStep.GetUnqualifiedAttributeName(attributeName),
                out var bundledTypes
            ) ?
                bundledTypes :
                [];

        private void Add(MemberDeclarationSyntax declaration, string namespaceName, bool isOutsideBundledNamespace)
        {
            var name = declaration switch
            {
                BaseTypeDeclarationSyntax typeDeclaration => typeDeclaration.Identifier.ValueText,
                DelegateDeclarationSyntax delegateDeclaration => delegateDeclaration.Identifier.ValueText,
                _ => null,
            };
            if (name == null)
            {
                return;
            }

            var bundledType = new BundledType(declaration, namespaceName, name);
            _types.Add(bundledType);
            _typesByDeclaration.Add(declaration, bundledType);
            if (isOutsideBundledNamespace)
            {
                _typesOutsideBundledNamespaces.Add(bundledType);
            }

            var unqualifiedName = CleanupStep.RemoveAttributeSuffix(name);
            if (!_typesByUnqualifiedName.TryGetValue(unqualifiedName, out var bundledTypes))
            {
                bundledTypes = new ();
                _typesByUnqualifiedName.Add(unqualifiedName, bundledTypes);
            }

            bundledTypes.Add(bundledType);
        }
    }
}