nub-lang/compiler/ModuleGraph.cs

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

namespace Compiler;

public class ModuleGraph
{
    public static Builder CreateBuilder() => new();

    private ModuleGraph(Dictionary<string, Module> modules)
    {
        this.modules = modules;
    }

    private readonly Dictionary<string, Module> modules;

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

    public bool TryResolveIdentifier(string moduleName, string identifierName, bool searchPrivate, [NotNullWhen(true)] out Module.IdentifierInfo? info)
    {
        if (!TryResolveModule(moduleName, out var module))
        {
            info = null;
            return false;
        }

        if (!module.TryResolveIdentifier(identifierName, searchPrivate, out info))
            return false;

        return true;
    }

    public bool TryResolveType(string moduleName, string typeName, bool searchPrivate, [NotNullWhen(true)] out Module.TypeInfo? info)
    {
        if (!TryResolveModule(moduleName, out var module))
        {
            info = null;
            return false;
        }

        if (!module.TryResolveType(typeName, searchPrivate, out info))
            return false;

        return true;
    }

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

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

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

        public void AddManifest(Manifest manifest)
        {
            manifests.Add(manifest);
        }

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

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

            foreach (var manifest in manifests)
            {
                foreach (var (moduleName, manifestModule) in manifest.Modules)
                {
                    var module = GetOrCreateModule(moduleName);

                    foreach (var (name, type) in manifestModule.Types)
                    {
                        switch (type)
                        {
                            case Manifest.Module.TypeInfoStruct s:
                            {
                                var info = new Module.TypeInfoStruct(Module.DefinitionSource.Imported, s.Exported, s.Packed);
                                var fields = s.Fields.Select(x => new Module.TypeInfoStruct.Field(x.Name, x.Type)).ToList();
                                info.SetFields(fields);
                                module.AddType(name, info);
                                break;
                            }
                            case Manifest.Module.TypeInfoEnum e:
                            {
                                var info = new Module.TypeInfoEnum(Module.DefinitionSource.Imported, e.Exported);
                                var variants = e.Variants.Select(v => new Module.TypeInfoEnum.Variant(v.Name, v.Type)).ToList();
                                info.SetVariants(variants);
                                module.AddType(name, info);
                                break;
                            }
                            default:
                                throw new ArgumentOutOfRangeException(nameof(type));
                        }
                    }

                    foreach (var (name, identifier) in manifestModule.Identifiers)
                    {
                        module.AddIdentifier(name, new Module.IdentifierInfo(Module.DefinitionSource.Imported, identifier.Exported, identifier.Extern, identifier.Type, identifier.MangledName));
                    }
                }
            }

            foreach (var ast in asts)
            {
                var module = GetOrCreateModule(ast.ModuleName.Ident);

                foreach (var structDef in ast.Definitions.OfType<NodeDefinitionStruct>())
                {
                    module.AddType(structDef.Name.Ident, new Module.TypeInfoStruct(Module.DefinitionSource.Internal, structDef.Exported, structDef.Packed));
                }

                foreach (var enumDef in ast.Definitions.OfType<NodeDefinitionEnum>())
                {
                    module.AddType(enumDef.Name.Ident, new Module.TypeInfoEnum(Module.DefinitionSource.Internal, enumDef.Exported));
                }
            }

            foreach (var ast in asts)
            {
                var module = GetOrCreateModule(ast.ModuleName.Ident);

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

                    if (typeInfo is Module.TypeInfoStruct structType)
                    {
                        var fields = structDef.Fields.Select(f => new Module.TypeInfoStruct.Field(f.Name.Ident, ResolveType(f.Type, module.Name))).ToList();
                        structType.SetFields(fields);
                    }
                }

                foreach (var enumDef in ast.Definitions.OfType<NodeDefinitionEnum>())
                {
                    if (!module.TryResolveType(enumDef.Name.Ident, true, out var typeInfo))
                        throw new UnreachableException($"{nameof(typeInfo)} should always be registered");

                    if (typeInfo is Module.TypeInfoEnum enumType)
                    {
                        var variants = enumDef.Variants.Select(v => new Module.TypeInfoEnum.Variant(v.Name.Ident, ResolveType(v.Type, module.Name))).ToList();
                        enumType.SetVariants(variants);
                    }
                }
            }

            foreach (var ast in asts)
            {
                var module = GetOrCreateModule(ast.ModuleName.Ident);

                foreach (var funcDef in ast.Definitions.OfType<NodeDefinitionFunc>())
                {
                    var parameters = funcDef.Parameters.Select(x => ResolveType(x.Type, module.Name)).ToList();
                    var returnType = funcDef.ReturnType == null ? NubTypeVoid.Instance : ResolveType(funcDef.ReturnType, module.Name);
                    var funcType = NubTypeFunc.Get(parameters, returnType);
                    var info = new Module.IdentifierInfo(Module.DefinitionSource.Internal, funcDef.Exported, false, funcType, NameMangler.Mangle(module.Name, funcDef.Name.Ident, funcType));
                    module.AddIdentifier(funcDef.Name.Ident, info);
                }

                foreach (var funcDef in ast.Definitions.OfType<NodeDefinitionExternFunc>())
                {
                    var parameters = funcDef.Parameters.Select(x => ResolveType(x.Type, module.Name)).ToList();
                    var returnType = funcDef.ReturnType == null ? NubTypeVoid.Instance : ResolveType(funcDef.ReturnType, module.Name);
                    var funcType = NubTypeFunc.Get(parameters, returnType);
                    var info = new Module.IdentifierInfo(Module.DefinitionSource.Internal, funcDef.Exported, true, funcType, funcDef.Name.Ident);
                    module.AddIdentifier(funcDef.Name.Ident, info);
                }

                foreach (var globalVariable in ast.Definitions.OfType<NodeDefinitionGlobalVariable>())
                {
                    var type = ResolveType(globalVariable.Type, module.Name);
                    var info = new Module.IdentifierInfo(Module.DefinitionSource.Internal, globalVariable.Exported, false, type, NameMangler.Mangle(module.Name, globalVariable.Name.Ident, type));
                    module.AddIdentifier(globalVariable.Name.Ident, info);
                }

                foreach (var globalVariable in ast.Definitions.OfType<NodeDefinitionExternGlobalVariable>())
                {
                    var type = ResolveType(globalVariable.Type, module.Name);
                    var info = new Module.IdentifierInfo(Module.DefinitionSource.Internal, globalVariable.Exported, true, type, NameMangler.Mangle(module.Name, globalVariable.Name.Ident, type));
                    module.AddIdentifier(globalVariable.Name.Ident, info);
                }
            }

            if (diagnostics.Any(x => x.Severity == Diagnostic.DiagnosticSeverity.Error))
                return null;

            return new ModuleGraph(modules);

            NubType ResolveType(NodeType node, string currentModule)
            {
                return node switch
                {
                    NodeTypeBool => NubTypeBool.Instance,
                    NodeTypeNamed type => ResolveNamedType(type, currentModule),
                    NodeTypeAnonymousStruct type => NubTypeAnonymousStruct.Get(type.Fields.Select(x => new NubTypeAnonymousStruct.Field(x.Name.Ident, ResolveType(x.Type, currentModule))).ToList()),
                    NodeTypeFunc type => NubTypeFunc.Get(type.Parameters.Select(x => ResolveType(x, currentModule)).ToList(), ResolveType(type.ReturnType, currentModule)),
                    NodeTypePointer type => NubTypePointer.Get(ResolveType(type.To, currentModule)),
                    NodeTypeSInt type => NubTypeSInt.Get(type.Width),
                    NodeTypeUInt type => NubTypeUInt.Get(type.Width),
                    NodeTypeString => NubTypeString.Instance,
                    NodeTypeVoid => NubTypeVoid.Instance,
                    _ => throw new ArgumentOutOfRangeException(nameof(node))
                };
            }

            NubType ResolveNamedType(NodeTypeNamed type, string currentModule)
            {
                return type.Sections.Count switch
                {
                    3 => ResolveThreePartType(type.Sections[0], type.Sections[1], type.Sections[2], currentModule),
                    2 => ResolveTwoPartType(type.Sections[0], type.Sections[1], currentModule),
                    1 => ResolveOnePartType(type.Sections[0], currentModule),
                    _ => throw BasicError("Invalid type name")
                };
            }

            NubType ResolveThreePartType(TokenIdent first, TokenIdent second, TokenIdent third, string currentModule)
            {
                var module = ResolveModule(first);
                if (!module.TryResolveType(second.Ident, currentModule == module.Name, out var typeInfo))
                    throw BasicError($"Named type '{module.Name}::{second.Ident}' not found");

                if (typeInfo is not Module.TypeInfoEnum enumInfo)
                    throw BasicError($"'{module.Name}::{second.Ident}' is not an enum");

                var variant = enumInfo.Variants.FirstOrDefault(v => v.Name == third.Ident);
                if (variant == null)
                    throw BasicError($"Enum '{module.Name}::{second.Ident}' does not have a variant named '{third.Ident}'");

                return NubTypeEnumVariant.Get(NubTypeEnum.Get(module.Name, second.Ident), third.Ident);
            }

            NubType ResolveTwoPartType(TokenIdent first, TokenIdent second, string currentModule)
            {
                if (TryResolveEnumVariant(currentModule, first.Ident, second.Ident, out var variant))
                    return variant;

                var module = ResolveModule(first);
                if (!module.TryResolveType(second.Ident, currentModule == module.Name, out var typeInfo))
                    throw BasicError($"Named type '{module.Name}::{second.Ident}' not found");

                return typeInfo switch
                {
                    Module.TypeInfoStruct => NubTypeStruct.Get(module.Name, second.Ident),
                    Module.TypeInfoEnum => NubTypeEnum.Get(module.Name, second.Ident),
                    _ => throw new ArgumentOutOfRangeException(nameof(typeInfo))
                };
            }

            NubType ResolveOnePartType(TokenIdent name, string currentModule)
            {
                if (!modules.TryGetValue(currentModule, out var module))
                    throw BasicError($"Module '{currentModule}' not found");

                if (!module.TryResolveType(name.Ident, true, out var typeInfo))
                    throw BasicError($"Named type '{module.Name}::{name.Ident}' not found");

                return typeInfo switch
                {
                    Module.TypeInfoStruct => NubTypeStruct.Get(module.Name, name.Ident),
                    Module.TypeInfoEnum => NubTypeEnum.Get(module.Name, name.Ident),
                    _ => throw new ArgumentOutOfRangeException(nameof(typeInfo))
                };
            }

            Module ResolveModule(TokenIdent name)
            {
                if (!modules.TryGetValue(name.Ident, out var module))
                    throw BasicError($"Module '{name.Ident}' not found");

                return module;
            }

            bool TryResolveEnumVariant(string moduleName, string enumName, string variantName, [NotNullWhen(true)] out NubType? result)
            {
                result = null;

                if (!modules.TryGetValue(moduleName, out var module))
                    return false;

                if (!module.TryResolveType(enumName, true, out var typeInfo))
                    return false;

                if (typeInfo is not Module.TypeInfoEnum enumInfo)
                    return false;

                var variant = enumInfo.Variants.FirstOrDefault(v => v.Name == variantName);
                if (variant == null)
                    return false;

                result = NubTypeEnumVariant.Get(
                    NubTypeEnum.Get(moduleName, enumName),
                    variantName);

                return true;
            }

            Exception BasicError(string message)
            {
                return new CompileException(Diagnostic.Error(message).Build());
            }

            Module GetOrCreateModule(string name)
            {
                if (!modules.TryGetValue(name, out var module))
                {
                    module = new Module(name);
                    modules.Add(name, module);
                }

                return module;
            }
        }
    }
}

public class Module(string name)
{
    public string Name { get; } = name;

    private Dictionary<string, TypeInfo> types = new();
    private Dictionary<string, IdentifierInfo> identifiers = new();

    public IReadOnlyDictionary<string, TypeInfo> GetTypes() => types;
    public IReadOnlyDictionary<string, IdentifierInfo> GetIdentifiers() => identifiers;

    public bool TryResolveType(string name, bool searchPrivate, [NotNullWhen(true)] out TypeInfo? customType)
    {
        var info = types.GetValueOrDefault(name);
        if (info == null)
        {
            customType = null;
            return false;
        }

        if (searchPrivate || info.Source == DefinitionSource.Internal)
        {
            customType = info;
            return true;
        }

        customType = null;
        return false;
    }

    public bool TryResolveIdentifier(string name, bool searchPrivate, [NotNullWhen(true)] out IdentifierInfo? identifierType)
    {
        var info = identifiers.GetValueOrDefault(name);
        if (info == null)
        {
            identifierType = null;
            return false;
        }

        if (searchPrivate || info.Source == DefinitionSource.Internal)
        {
            identifierType = info;
            return true;
        }

        identifierType = null;
        return false;
    }

    public void AddType(string name, TypeInfo info)
    {
        types.Add(name, info);
    }

    public void AddIdentifier(string name, IdentifierInfo info)
    {
        identifiers.Add(name, info);
    }

    public enum DefinitionSource
    {
        Internal,
        Imported,
    }

    public class IdentifierInfo(DefinitionSource source, bool exported, bool @extern, NubType type, string mangledName)
    {
        public DefinitionSource Source { get; } = source;
        public bool Exported { get; } = exported;
        public bool Extern { get; } = @extern;
        public NubType Type { get; } = type;
        public string MangledName { get; } = mangledName;
    }

    public abstract class TypeInfo(DefinitionSource source, bool exported)
    {
        public DefinitionSource Source { get; } = source;
        public bool Exported { get; } = exported;
    }

    public class TypeInfoStruct(DefinitionSource source, bool exported, bool packed) : TypeInfo(source, exported)
    {
        private IReadOnlyList<Field>? fields;

        public bool Packed { get; } = packed;
        public IReadOnlyList<Field> Fields => fields ?? throw new InvalidOperationException("Fields has not been set yet");

        public void SetFields(IReadOnlyList<Field> fields)
        {
            this.fields = fields;
        }

        public class Field(string name, NubType type)
        {
            public string Name { get; } = name;
            public NubType Type { get; } = type;
        }
    }

    public class TypeInfoEnum(DefinitionSource source, bool exported) : TypeInfo(source, exported)
    {
        private IReadOnlyList<Variant>? variants;

        public IReadOnlyList<Variant> Variants => variants ?? throw new InvalidOperationException("Fields has not been set yet");

        public void SetVariants(IReadOnlyList<Variant> variants)
        {
            this.variants = variants;
        }

        public class Variant(string name, NubType type)
        {
            public string Name { get; } = name;
            public NubType Type { get; } = type;
        }
    }
}