nub-lang/compiler/ModuleGraph.cs

using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;

namespace Compiler;

public class ModuleGraph(Dictionary<string, ModuleGraph.Module> modules)
{
    public static Builder Create() => new();

    public List<Module> GetModules()
    {
        return modules.Values.ToList();
    }

    public bool TryResolveModule(string moduleName, [NotNullWhen(true)] out Module? module)
    {
        module = modules.GetValueOrDefault(moduleName);
        return module != null;
    }

    public sealed class Module(string name)
    {
        public string Name { get; } = name;
        private readonly Dictionary<string, NubType> customTypes = new();
        private readonly Dictionary<string, NubType> identifierTypes = new();

        public List<NubType> GetCustomTypes()
        {
            return customTypes.Values.ToList();
        }

        public bool TryResolveCustomType(string name, [NotNullWhen(true)] out NubType? customType)
        {
            customType = customTypes.GetValueOrDefault(name);
            return customType != null;
        }

        public bool TryResolveIdentifierType(string name, [NotNullWhen(true)] out NubType? identifier)
        {
            identifier = identifierTypes.GetValueOrDefault(name);
            return identifier != null;
        }

        public void AddCustomType(string name, NubType type)
        {
            customTypes.Add(name, type);
        }

        public void AddIdentifier(string name, NubType identifier)
        {
            identifierTypes.Add(name, identifier);
        }
    }

    public class Builder
    {
        private readonly List<Ast> asts = [];

        public void AddAst(Ast ast)
        {
            asts.Add(ast);
        }

        public ModuleGraph Build(out List<Diagnostic> diagnostics)
        {
            diagnostics = [];

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

            // First pass: Register modules
            foreach (var ast in asts)
            {
                var moduleDefinitions = ast.Definitions.OfType<NodeDefinitionModule>().ToList();

                if (moduleDefinitions.Count == 0)
                    diagnostics.Add(Diagnostic.Error("Missing module declaration").At(ast.FileName, 1, 1, 1).Build());

                foreach (var extraModuleDefinition in moduleDefinitions.Skip(1))
                    diagnostics.Add(Diagnostic.Warning("Duplicate module declaration will be ignored").At(ast.FileName, extraModuleDefinition).Build());

                if (moduleDefinitions.Count >= 1)
                {
                    var currentModule = moduleDefinitions[0].Name.Ident;

                    if (!modules.ContainsKey(currentModule))
                    {
                        var module = new Module(currentModule);
                        modules.Add(currentModule, module);
                        astModuleCache[ast] = module;
                    }
                }
            }

            // Second pass: Register struct types without fields
            foreach (var ast in asts)
            {
                var module = astModuleCache.GetValueOrDefault(ast);
                if (module == null) continue;

                foreach (var structDef in ast.Definitions.OfType<NodeDefinitionStruct>())
                    module.AddCustomType(structDef.Name.Ident, new NubTypeStruct(module.Name, structDef.Name.Ident));
            }

            // Third pass: Resolve struct fields
            foreach (var ast in asts)
            {
                var module = astModuleCache.GetValueOrDefault(ast);
                if (module == null) continue;

                foreach (var structDef in ast.Definitions.OfType<NodeDefinitionStruct>())
                {
                    if (!module.TryResolveCustomType(structDef.Name.Ident, out var customType))
                        throw new UnreachableException($"{nameof(customType)} should always be registered");

                    var fields = structDef.Fields.Select(f => new NubTypeStruct.Field(f.Name.Ident, ResolveType(f.Type))).ToList();
                    ((NubTypeStruct)customType).ResolveFields(fields);
                }
            }

            // Fourth pass: Register identifiers
            foreach (var ast in asts)
            {
                var module = astModuleCache.GetValueOrDefault(ast);
                if (module == null) continue;

                foreach (var funcDef in ast.Definitions.OfType<NodeDefinitionFunc>())
                {
                    var parameters = funcDef.Parameters.Select(x => ResolveType(x.Type)).ToList();
                    var returnType = ResolveType(funcDef.ReturnType);
                    var funcType = NubTypeFunc.Get(parameters, returnType);
                    module.AddIdentifier(funcDef.Name.Ident, funcType);
                }
            }

            return new ModuleGraph(modules);

            NubType ResolveType(NodeType node)
            {
                return node switch
                {
                    NodeTypeBool => NubTypeBool.Instance,
                    NodeTypeCustom type => ResolveCustomType(type),
                    NodeTypeFunc type => NubTypeFunc.Get(type.Parameters.Select(ResolveType).ToList(), ResolveType(type.ReturnType)),
                    NodeTypePointer type => NubTypePointer.Get(ResolveType(type.To)),
                    NodeTypeSInt type => NubTypeSInt.Get(type.Width),
                    NodeTypeUInt type => NubTypeUInt.Get(type.Width),
                    NodeTypeString => NubTypeString.Instance,
                    NodeTypeVoid => NubTypeVoid.Instance,
                    _ => throw new ArgumentOutOfRangeException(nameof(node))
                };
            }

            NubType ResolveCustomType(NodeTypeCustom type)
            {
                var module = modules.GetValueOrDefault(type.Module.Ident);
                if (module == null)
                    throw new CompileException(Diagnostic.Error($"Unknown module: {type.Module.Ident}").Build());

                if (!module.TryResolveCustomType(type.Name.Ident, out var customType))
                    throw new CompileException(Diagnostic.Error($"Unknown custom type: {type.Module.Ident}::{type.Name.Ident}").Build());

                return customType;
            }
        }
    }
}