Tests.cs

namespace SpacetimeDB.Codegen.Tests;

using System.Collections.Immutable;
using System.Runtime.CompilerServices;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.MSBuild;
using Microsoft.CodeAnalysis.Text;

public static class GeneratorSnapshotTests
{
    // Note that we can't use assembly path here because it will be put in some deep nested folder.
    // Instead, to get the test project directory, we can use the `CallerFilePath` attribute which will magically give us path to the current file.
    static string GetProjectDir([CallerFilePath] string path = "") => Path.GetDirectoryName(path)!;

    record struct StepOutput(string Key, IncrementalStepRunReason Reason, object Value);

    class Fixture
    {
        private readonly string projectDir;
        private readonly CSharpCompilation sampleCompilation;

        public Fixture(string projectDir, CSharpCompilation sampleCompilation)
        {
            this.projectDir = projectDir;
            this.sampleCompilation = sampleCompilation;
        }

        public CSharpCompilation SampleCompilation => sampleCompilation;

        public static async Task<Fixture> Compile(string name)
        {
            var projectDir = Path.Combine(GetProjectDir(), "fixtures", name);
            using var workspace = MSBuildWorkspace.Create();
            var sampleProject = await workspace.OpenProjectAsync($"{projectDir}/{name}.csproj");
            var compilation = await sampleProject.GetCompilationAsync();
            return new(projectDir, (CSharpCompilation)compilation!);
        }

        public Task Verify(string fileName, object target) =>
            Verifier.Verify(target).UseDirectory($"{projectDir}/snapshots").UseFileName(fileName);

        private static CSharpGeneratorDriver CreateDriver(
            IIncrementalGenerator generator,
            LanguageVersion languageVersion
        )
        {
            return CSharpGeneratorDriver.Create(
                [generator.AsSourceGenerator()],
                driverOptions: new(
                    disabledOutputs: IncrementalGeneratorOutputKind.None,
                    trackIncrementalGeneratorSteps: true
                ),
                // Make sure that generated files are parsed with the same language version.
                parseOptions: new(languageVersion)
            );
        }

        private async Task<IEnumerable<SyntaxTree>> RunAndCheckGenerator(
            IIncrementalGenerator generator
        )
        {
            var driver = CreateDriver(generator, sampleCompilation.LanguageVersion);

            // Store the new driver instance - it contains the results and the cache.
            var driverAfterGen = driver.RunGenerators(sampleCompilation);
            var genResult = driverAfterGen.GetRunResult();

            // Verify the generated code against the snapshots.
            await Verify(generator.GetType().Name, genResult);

            CheckCacheWorking(sampleCompilation, driverAfterGen);

            return genResult.GeneratedTrees;
        }

        public GeneratorDriverRunResult RunGeneratorAndGetResult(IIncrementalGenerator generator)
        {
            var driver = CreateDriver(generator, sampleCompilation.LanguageVersion);
            return driver.RunGenerators(sampleCompilation).GetRunResult();
        }

        public async Task<CSharpCompilation> RunAndCheckGenerators(
            params IIncrementalGenerator[] generators
        ) =>
            sampleCompilation.AddSyntaxTrees(
                (await Task.WhenAll(generators.Select(RunAndCheckGenerator))).SelectMany(output =>
                    output
                )
            );
    }

    private static void CheckCacheWorking(
        CSharpCompilation sampleCompilation,
        GeneratorDriver driverAfterGen
    )
    {
        // Run again with a driver containing the cache and a trivially modified code to verify that the cache is working.
        var modifiedCompilation = sampleCompilation
            .RemoveAllSyntaxTrees()
            .AddSyntaxTrees(
                sampleCompilation.SyntaxTrees.Select(tree =>
                    tree.WithChangedText(
                        SourceText.From(
                            string.Join(
                                "\n",
                                tree.GetText().Lines.Select(line => $"{line} // Modified")
                            )
                        )
                    )
                )
            );

        var driverAfterRegen = driverAfterGen.RunGenerators(modifiedCompilation);

        var regenSteps = driverAfterRegen
            .GetRunResult()
            .Results.SelectMany(result => result.TrackedSteps)
            .Where(step => step.Key.StartsWith("SpacetimeDB."))
            .SelectMany(step =>
                step.Value.SelectMany(value => value.Outputs)
                    .Select(output => new StepOutput(step.Key, output.Reason, output.Value))
            )
            .ToImmutableArray();

        // Ensure that we have tracked steps at all.
        Assert.NotEmpty(regenSteps);

        // Ensure that all steps were cached.
        Assert.Empty(
            regenSteps.Where(step =>
                step.Reason
                    is not (IncrementalStepRunReason.Cached or IncrementalStepRunReason.Unchanged)
            )
        );
    }

    static IEnumerable<Diagnostic> GetCompilationErrors(Compilation compilation)
    {
        return compilation
            .Emit(Stream.Null)
            .Diagnostics.Where(diag => diag.Severity != DiagnosticSeverity.Hidden)
            // The order of diagnostics is not predictable, sort them by location to make the test deterministic.
            .OrderBy(diag => diag.GetMessage() + diag.Location.ToString());
    }

    static void AssertGeneratedCodeDoesNotUseInternalBound(CSharpCompilation compilation)
    {
        var generatedText = string.Join(
            "\n\n",
            compilation.SyntaxTrees.Select(tree => tree.GetText().ToString())
        );

        Assert.DoesNotContain("global::SpacetimeDB.Internal.Bound<", generatedText);
        Assert.Contains("global::SpacetimeDB.Bound<", generatedText);
    }

    static void AssertPublicBoundIsAvailableInRuntime(Compilation compilation)
    {
        var bound = compilation.GetTypeByMetadataName("SpacetimeDB.Bound`1");
        Assert.NotNull(bound);
        Assert.Equal(Accessibility.Public, bound!.DeclaredAccessibility);
    }

    static void AssertRuntimeDoesNotDefineLocal(Compilation compilation)
    {
        var runtimeAssembly = compilation
            .References.Select(r => compilation.GetAssemblyOrModuleSymbol(r))
            .OfType<IAssemblySymbol>()
            .FirstOrDefault(a => a.Name == "SpacetimeDB.Runtime");

        Assert.NotNull(runtimeAssembly);

        // These types are generated per-module by SpacetimeDB.Codegen.Module.
        // If Runtime defines any of them too, user projects can hit CS0436 warnings.
        var codegenOwnedTypes = new[]
        {
            "SpacetimeDB.Local",
            "SpacetimeDB.ProcedureContext",
            "SpacetimeDB.ProcedureTxContext",
            "SpacetimeDB.ReducerContext",
            "SpacetimeDB.ViewContext",
            "SpacetimeDB.AnonymousViewContext",
        };

        foreach (var name in codegenOwnedTypes)
        {
            Assert.Null(runtimeAssembly!.GetTypeByMetadataName(name));
        }
    }

    static void AssertNoCs0436Diagnostics(Compilation compilation)
    {
        var diagnostics = compilation
            .Emit(Stream.Null)
            .Diagnostics.Where(diag => diag.Severity != DiagnosticSeverity.Hidden);

        Assert.DoesNotContain(diagnostics, d => d.Id == "CS0436");
    }

    [Fact]
    public static async Task TypeGeneratorOnClient()
    {
        var fixture = await Fixture.Compile("client");

        var compilationAfterGen = await fixture.RunAndCheckGenerators(
            new SpacetimeDB.Codegen.Type()
        );

        Assert.Empty(GetCompilationErrors(compilationAfterGen));
    }

    [Fact]
    public static async Task TypeAndModuleGeneratorsOnServer()
    {
        var fixture = await Fixture.Compile("server");

        var compilationAfterGen = await fixture.RunAndCheckGenerators(
            new SpacetimeDB.Codegen.Type(),
            new SpacetimeDB.Codegen.Module()
        );

        Assert.Empty(GetCompilationErrors(compilationAfterGen));

        AssertPublicBoundIsAvailableInRuntime(compilationAfterGen);
        AssertRuntimeDoesNotDefineLocal(compilationAfterGen);
        AssertGeneratedCodeDoesNotUseInternalBound(compilationAfterGen);

        // Regression guard for user-reported warning spam:
        // make sure a downstream "user" file that references SpacetimeDB.Local doesn't trigger CS0436.
        var userCode =
            "namespace User; public sealed class UseLocal { public SpacetimeDB.Local Db; }";
        var userTree = CSharpSyntaxTree.ParseText(
            userCode,
            new CSharpParseOptions(compilationAfterGen.LanguageVersion)
        );
        var compilationWithUserCode = compilationAfterGen.AddSyntaxTrees(userTree);
        AssertNoCs0436Diagnostics(compilationWithUserCode);
    }

    [Fact]
    public static async Task SettingsAndExplicitNames()
    {
        var fixture = await Fixture.Compile("explicitnames");

        var compilationAfterGen = await fixture.RunAndCheckGenerators(
            new SpacetimeDB.Codegen.Type(),
            new SpacetimeDB.Codegen.Module()
        );

        Assert.Empty(GetCompilationErrors(compilationAfterGen));

        AssertPublicBoundIsAvailableInRuntime(compilationAfterGen);
        AssertRuntimeDoesNotDefineLocal(compilationAfterGen);
        AssertGeneratedCodeDoesNotUseInternalBound(compilationAfterGen);
    }

    [Fact]
    public static async Task CSharpKeywordIdentifiersAreEscapedInGeneratedCode()
    {
        var fixture = await Fixture.Compile("server");

        const string source =
            """
            using SpacetimeDB;

            [SpacetimeDB.Table]
            public partial struct KeywordTable
            {
                [SpacetimeDB.PrimaryKey]
                public ulong @class;

                public int @params;
            }

            [SpacetimeDB.Table(Accessor = "class")]
            public partial struct AccessorKeywordTable
            {
                [SpacetimeDB.PrimaryKey]
                [SpacetimeDB.Index.BTree(Accessor = "class")]
                public int Id;
            }

            [SpacetimeDB.Table]
            public partial struct @class
            {
                [SpacetimeDB.PrimaryKey]
                public int Id;
            }

            public static partial class KeywordApis
            {
                [SpacetimeDB.Reducer]
                public static void KeywordReducer(ReducerContext ctx, int @params, string @class)
                {
                    _ = @params;
                    _ = @class;
                }

                [SpacetimeDB.Reducer]
                public static void @class(ReducerContext ctx)
                {
                }

                [SpacetimeDB.Procedure]
                public static int KeywordProcedure(ProcedureContext ctx, int @params, int @class)
                {
                    return @params + @class;
                }

                [SpacetimeDB.Procedure]
                public static void @params(ProcedureContext ctx)
                {
                }
            }
            """;

        var parseOptions = new CSharpParseOptions(fixture.SampleCompilation.LanguageVersion);
        var tree = CSharpSyntaxTree.ParseText(source, parseOptions, path: "KeywordNames.cs");
        var compilation = fixture.SampleCompilation.AddSyntaxTrees(tree);

        var driver = CSharpGeneratorDriver.Create(
            [new SpacetimeDB.Codegen.Type().AsSourceGenerator(), new SpacetimeDB.Codegen.Module().AsSourceGenerator()],
            driverOptions: new(
                disabledOutputs: IncrementalGeneratorOutputKind.None,
                trackIncrementalGeneratorSteps: true
            ),
            parseOptions: parseOptions
        );

        var runResult = driver.RunGenerators(compilation).GetRunResult();
        var compilationAfterGen = compilation.AddSyntaxTrees(runResult.GeneratedTrees);

        Assert.Empty(GetCompilationErrors(compilationAfterGen));
    }

    [Fact]
    public static async Task TestDiagnostics()
    {
        var fixture = await Fixture.Compile("diag");

        var compilationAfterGen = await fixture.RunAndCheckGenerators(
            new SpacetimeDB.Codegen.Type(),
            new SpacetimeDB.Codegen.Module()
        );

        // Unlike in regular tests, we don't expect this compilation to succeed - it's supposed to be full of errors.
        // We already reported the useful ones from the generator, but let's snapshot those emitted by the compiler as well.
        // This way we can notice when they get particularly noisy and improve our codegen for the case of a broken code.
        await fixture.Verify("ExtraCompilationErrors", GetCompilationErrors(compilationAfterGen));

        AssertPublicBoundIsAvailableInRuntime(compilationAfterGen);
        AssertRuntimeDoesNotDefineLocal(compilationAfterGen);
        AssertGeneratedCodeDoesNotUseInternalBound(compilationAfterGen);
    }

    [Fact]
    public static async Task ViewInvalidReturnHighlightsReturnType()
    {
        var fixture = await Fixture.Compile("diag");

        var runResult = fixture.RunGeneratorAndGetResult(new SpacetimeDB.Codegen.Module());

        var method = fixture
            .SampleCompilation.SyntaxTrees.Select(tree => new
            {
                Tree = tree,
                Root = tree.GetRoot(),
            })
            .SelectMany(entry =>
                entry
                    .Root.DescendantNodes()
                    .OfType<MethodDeclarationSyntax>()
                    .Select(method => new
                    {
                        entry.Tree,
                        entry.Root,
                        Method = method,
                    })
            )
            .Single(entry => entry.Method.Identifier.Text == "ViewDefIEnumerableReturnFromIter");

        var returnTypeSpan = method.Method.ReturnType.Span;
        var diagnostics = runResult
            .Results.SelectMany(result => result.Diagnostics)
            .Where(d => d.Id == "STDB0024")
            .ToList();
        var diagnostic = diagnostics.FirstOrDefault(d =>
            d.GetMessage().Contains("ViewDefIEnumerableReturnFromIter")
            && d.Location.SourceTree == method.Tree
        );

        Assert.NotNull(diagnostic);

        Assert.Equal(returnTypeSpan, diagnostic!.Location.SourceSpan);

        var returnTypeText = method
            .Root.ToFullString()
            .Substring(returnTypeSpan.Start, returnTypeSpan.Length);
        Assert.Contains("IEnumerable", returnTypeText);
    }
}