nub-lang/compiler/NubLang/Modules/ModuleRepository.cs

using System.Diagnostics.CodeAnalysis;
using NubLang.Ast;
using NubLang.Diagnostics;
using NubLang.Syntax;
using NubLang.Types;

namespace NubLang.Modules;

public sealed class ModuleRepository
{
    public static ModuleRepository Create(List<SyntaxTree> syntaxTrees)
    {
        var structTypes = new Dictionary<(string module, string name), NubStructType>();
        var enumTypes = new Dictionary<(string module, string name), NubIntType>();

        foreach (var syntaxTree in syntaxTrees)
        {
            var module = syntaxTree.TopLevelSyntaxNodes.OfType<ModuleSyntax>().FirstOrDefault();
            if (module == null)
            {
                throw new CompileException(Diagnostic.Error("Module declaration missing").WithHelp("module \"main\"").Build());
            }

            foreach (var structSyntax in syntaxTree.TopLevelSyntaxNodes.OfType<StructSyntax>())
            {
                // note(nub31): Since not all struct types are registered yet, we cannot register field types as they might reference unregistered structs
                var key = (module.NameToken.Value, structSyntax.NameToken.Value);
                structTypes.Add(key, new NubStructType(module.NameToken.Value, structSyntax.NameToken.Value, structSyntax.Packed, []));
            }

            foreach (var enumSyntax in syntaxTree.TopLevelSyntaxNodes.OfType<EnumSyntax>())
            {
                NubIntType? underlyingType = null;
                if (enumSyntax.Type != null)
                {
                    if (enumSyntax.Type is not IntTypeSyntax intType)
                    {
                        throw new CompileException(Diagnostic.Error("Underlying type of enum must be an integer type").At(enumSyntax.Type).Build());
                    }

                    underlyingType = new NubIntType(intType.Signed, intType.Width);
                }

                underlyingType ??= new NubIntType(false, 64);

                var key = (module.NameToken.Value, enumSyntax.NameToken.Value);
                enumTypes.Add(key, underlyingType);
            }
        }

        // note(nub31): Since all struct types are now registered, we can safely resolve the field types
        foreach (var syntaxTree in syntaxTrees)
        {
            var module = syntaxTree.TopLevelSyntaxNodes.OfType<ModuleSyntax>().FirstOrDefault();
            if (module == null)
            {
                throw new CompileException(Diagnostic.Error("Module declaration missing").WithHelp("module \"main\"").Build());
            }

            foreach (var structSyntax in syntaxTree.TopLevelSyntaxNodes.OfType<StructSyntax>())
            {
                var key = (module.NameToken.Value, structSyntax.NameToken.Value);

                structTypes[key].Fields = structSyntax.Fields
                    .Select(x => new NubStructFieldType(x.NameToken.Value, ResolveType(x.Type, module.NameToken.Value), x.Value != null))
                    .ToList();
            }
        }

        var modules = new Dictionary<string, Module>();

        foreach (var syntaxTree in syntaxTrees)
        {
            var moduleDecl = syntaxTree.TopLevelSyntaxNodes.OfType<ModuleSyntax>().FirstOrDefault();
            if (moduleDecl == null)
            {
                throw new CompileException(Diagnostic.Error("Module declaration missing").WithHelp("module \"main\"").Build());
            }

            var functionPrototypes = new List<FuncPrototypeNode>();

            foreach (var funcSyntax in syntaxTree.TopLevelSyntaxNodes.OfType<FuncSyntax>())
            {
                var returnType = ResolveType(funcSyntax.Prototype.ReturnType, moduleDecl.NameToken.Value);
                var parameters = funcSyntax.Prototype.Parameters.Select(x => new FuncParameterNode(x.Tokens, x.NameToken, ResolveType(x.Type, moduleDecl.NameToken.Value))).ToList();
                functionPrototypes.Add(new FuncPrototypeNode(funcSyntax.Prototype.Tokens, funcSyntax.Prototype.NameToken, funcSyntax.Prototype.ExternSymbolToken, parameters, returnType));
            }

            var module = new Module
            {
                Name = moduleDecl.NameToken.Value,
                StructTypes = structTypes.Where(x => x.Key.module == moduleDecl.NameToken.Value).Select(x => x.Value).ToList(),
                EnumTypes = enumTypes
                    .Where(x => x.Key.module == moduleDecl.NameToken.Value)
                    .ToDictionary(x => x.Key.name, x => x.Value),
                FunctionPrototypes = functionPrototypes
            };

            modules.Add(moduleDecl.NameToken.Value, module);
        }

        return new ModuleRepository(modules);

        NubType ResolveType(TypeSyntax type, string currentModule)
        {
            return type switch
            {
                ArrayTypeSyntax arr => new NubArrayType(ResolveType(arr.BaseType, currentModule)),
                BoolTypeSyntax => new NubBoolType(),
                IntTypeSyntax i => new NubIntType(i.Signed, i.Width),
                FloatTypeSyntax f => new NubFloatType(f.Width),
                FuncTypeSyntax func => new NubFuncType(func.Parameters.Select(x => ResolveType(x, currentModule)).ToList(), ResolveType(func.ReturnType, currentModule)),
                SliceTypeSyntax slice => new NubSliceType(ResolveType(slice.BaseType, currentModule)),
                ConstArrayTypeSyntax arr => new NubConstArrayType(ResolveType(arr.BaseType, currentModule), arr.Size),
                PointerTypeSyntax ptr => new NubPointerType(ResolveType(ptr.BaseType, currentModule)),
                StringTypeSyntax => new NubStringType(),
                CustomTypeSyntax c => ResolveCustomType(c, currentModule),
                VoidTypeSyntax => new NubVoidType(),
                _ => throw new NotSupportedException($"Unknown type syntax: {type}")
            };
        }

        NubType ResolveCustomType(CustomTypeSyntax customType, string currentModule)
        {
            var customTypeKey = (customType.ModuleToken?.Value ?? currentModule, customType.NameToken.Value);

            var resolvedStructType = structTypes.GetValueOrDefault(customTypeKey);
            if (resolvedStructType != null)
            {
                return resolvedStructType;
            }

            var resolvedEnumType = enumTypes.GetValueOrDefault(customTypeKey);
            if (resolvedEnumType != null)
            {
                return resolvedEnumType;
            }

            throw new CompileException(Diagnostic
                .Error($"Type {customType.NameToken.Value} not found in module {customType.ModuleToken?.Value ?? currentModule}")
                .At(customType)
                .Build());
        }
    }

    public ModuleRepository(Dictionary<string, Module> modules)
    {
        _modules = modules;
    }

    private readonly Dictionary<string, Module> _modules;

    public Module Get(IdentifierToken ident)
    {
        var module = _modules.GetValueOrDefault(ident.Value);
        if (module == null)
        {
            throw new CompileException(Diagnostic.Error($"Module {ident.Value} was not found").At(ident).Build());
        }

        return module;
    }

    public bool TryGet(IdentifierToken ident, [NotNullWhen(true)] out Module? module)
    {
        module = _modules.GetValueOrDefault(ident.Value);
        return module != null;
    }

    public bool TryGet(string name, [NotNullWhen(true)] out Module? module)
    {
        module = _modules.GetValueOrDefault(name);
        return module != null;
    }

    public List<Module> GetAll()
    {
        return _modules.Values.ToList();
    }

    public sealed class Module
    {
        public required string Name { get; init; }
        public required List<FuncPrototypeNode> FunctionPrototypes { get; init; } = [];
        public required List<NubStructType> StructTypes { get; init; } = [];
        public required Dictionary<string, NubIntType> EnumTypes { get; init; } = [];

        public bool TryResolveFunc(string name, [NotNullWhen(true)] out FuncPrototypeNode? value, [NotNullWhen(false)] out Diagnostic? diagnostic)
        {
            value = FunctionPrototypes.FirstOrDefault(x => x.NameToken.Value == name);

            if (value == null)
            {
                value = null;
                diagnostic = Diagnostic.Error($"Func {name} not found in module {Name}").Build();
                return false;
            }

            diagnostic = null;
            return true;
        }

        public bool TryResolveFunc(IdentifierToken name, [NotNullWhen(true)] out FuncPrototypeNode? value, [NotNullWhen(false)] out Diagnostic? diagnostic)
        {
            value = FunctionPrototypes.FirstOrDefault(x => x.NameToken.Value == name.Value);

            if (value == null)
            {
                value = null;
                diagnostic = Diagnostic.Error($"Func {name.Value} not found in module {Name}").At(name).Build();
                return false;
            }

            diagnostic = null;
            return true;
        }

        public FuncPrototypeNode ResolveFunc(IdentifierToken name)
        {
            if (!TryResolveFunc(name, out var value, out var diagnostic))
            {
                throw new CompileException(diagnostic);
            }

            return value;
        }

        public bool TryResolveStruct(string name, [NotNullWhen(true)] out NubStructType? value, [NotNullWhen(false)] out Diagnostic? diagnostic)
        {
            value = StructTypes.FirstOrDefault(x => x.Name == name);

            if (value == null)
            {
                value = null;
                diagnostic = Diagnostic.Error($"Struct {name} not found in module {Name}").Build();
                return false;
            }

            diagnostic = null;
            return true;
        }

        public bool TryResolveStruct(IdentifierToken name, [NotNullWhen(true)] out NubStructType? value, [NotNullWhen(false)] out Diagnostic? diagnostic)
        {
            value = StructTypes.FirstOrDefault(x => x.Name == name.Value);

            if (value == null)
            {
                value = null;
                diagnostic = Diagnostic.Error($"Struct {name.Value} not found in module {Name}").At(name).Build();
                return false;
            }

            diagnostic = null;
            return true;
        }

        public NubStructType ResolveStruct(IdentifierToken name)
        {
            if (!TryResolveStruct(name, out var value, out var diagnostic))
            {
                throw new CompileException(diagnostic);
            }

            return value;
        }

        public bool TryResolveEnum(string name, [NotNullWhen(true)] out NubIntType? value, [NotNullWhen(false)] out Diagnostic? diagnostic)
        {
            value = EnumTypes.GetValueOrDefault(name);

            if (value == null)
            {
                value = null;
                diagnostic = Diagnostic.Error($"Enum {name} not found in module {Name}").Build();
                return false;
            }

            diagnostic = null;
            return true;
        }

        public bool TryResolveEnum(IdentifierToken name, [NotNullWhen(true)] out NubIntType? value, [NotNullWhen(false)] out Diagnostic? diagnostic)
        {
            value = EnumTypes.GetValueOrDefault(name.Value);

            if (value == null)
            {
                value = null;
                diagnostic = Diagnostic.Error($"Enum {name.Value} not found in module {Name}").At(name).Build();
                return false;
            }

            diagnostic = null;
            return true;
        }

        public NubIntType ResolveEnum(IdentifierToken name)
        {
            if (!TryResolveEnum(name, out var value, out var diagnostic))
            {
                throw new CompileException(diagnostic);
            }

            return value;
        }
    }
}