nub-lang/compiler/Generator.cs

using System.Text;

namespace Compiler;

public sealed class Generator(List<TypedNodeDefinitionFunc> functions, ModuleGraph moduleGraph)
{
    public static string Emit(List<TypedNodeDefinitionFunc> functions, ModuleGraph moduleGraph)
    {
        return new Generator(functions, moduleGraph).Emit();
    }

    private readonly IndentedTextWriter writer = new();
    private readonly Dictionary<NubTypeStruct, string> structTypeNames = new();

    private string Emit()
    {
        writer.WriteLine("""
                         #include <float.h>
                         #include <stdarg.h>
                         #include <stddef.h>
                         #include <stdbool.h>
                         #include <stdint.h>

                         struct string {
                            const char *data;
                            int length;
                         };

                         """);

        foreach (var (i, structType) in moduleGraph.GetStructTypes().Index())
            structTypeNames[structType] = $"s{i}";

        foreach (var typeName in structTypeNames)
            writer.WriteLine($"struct {typeName.Value};");

        writer.WriteLine();

        foreach (var typeName in structTypeNames)
        {
            writer.WriteLine($"struct {typeName.Value}");
            writer.WriteLine("{");
            using (writer.Indent())
            {
                foreach (var field in typeName.Key.Fields)
                    writer.WriteLine($"{CType(field.Type, field.Name)};");
            }

            writer.WriteLine("};");
        }

        writer.WriteLine();

        foreach (var node in functions)
        {
            var parameters = node.Parameters.Select(x => CType(x.Type, x.Name.Ident));
            writer.WriteLine($"{CType(node.ReturnType, node.Name.Ident)}({string.Join(", ", parameters)});");
        }

        writer.WriteLine();

        foreach (var node in functions)
        {
            var parameters = node.Parameters.Select(x => CType(x.Type, x.Name.Ident));
            writer.WriteLine($"{CType(node.ReturnType, node.Name.Ident)}({string.Join(", ", parameters)})");
            writer.WriteLine("{");
            using (writer.Indent())
            {
                EmitStatement(node.Body);
            }

            writer.WriteLine("}");
            writer.WriteLine();
        }

        return writer.ToString();
    }

    private void EmitStatement(TypedNodeStatement node)
    {
        switch (node)
        {
            case TypedNodeStatementBlock statement:
                EmitStatementBlock(statement);
                break;
            case TypedNodeStatementFuncCall statement:
                EmitStatementFuncCall(statement);
                break;
            case TypedNodeStatementReturn statement:
                EmitStatementReturn(statement);
                break;
            case TypedNodeStatementVariableDeclaration statement:
                EmitStatementVariableDeclaration(statement);
                break;
            case TypedNodeStatementAssignment statement:
                EmitStatementAssignment(statement);
                break;
            case TypedNodeStatementIf statement:
                EmitStatementIf(statement);
                break;
            case TypedNodeStatementWhile statement:
                EmitStatementWhile(statement);
                break;
            default:
                throw new ArgumentOutOfRangeException(nameof(node), node, null);
        }
    }

    private void EmitStatementBlock(TypedNodeStatementBlock node)
    {
        writer.WriteLine("{");
        using (writer.Indent())
        {
            foreach (var statement in node.Statements)
                EmitStatement(statement);
        }

        writer.WriteLine("}");
    }

    private void EmitStatementFuncCall(TypedNodeStatementFuncCall node)
    {
        var name = EmitExpression(node.Target);
        var parameterValues = node.Parameters.Select(EmitExpression).ToList();
        writer.WriteLine($"{name}({string.Join(", ", parameterValues)});");
    }

    private void EmitStatementReturn(TypedNodeStatementReturn statement)
    {
        var value = EmitExpression(statement.Value);
        writer.WriteLine($"return {value};");
    }

    private void EmitStatementVariableDeclaration(TypedNodeStatementVariableDeclaration statement)
    {
        var value = EmitExpression(statement.Value);
        writer.WriteLine($"{CType(statement.Type)} {statement.Name.Ident} = {value};");
    }

    private void EmitStatementAssignment(TypedNodeStatementAssignment statement)
    {
        var target = EmitExpression(statement.Target);
        var value = EmitExpression(statement.Value);
        writer.WriteLine($"{target} = {value};");
    }

    private void EmitStatementIf(TypedNodeStatementIf statement)
    {
        var condition = EmitExpression(statement.Condition);
        writer.WriteLine($"if ({condition})");
        writer.WriteLine("{");
        using (writer.Indent())
        {
            EmitStatement(statement.ThenBlock);
        }

        writer.WriteLine("}");

        if (statement.ElseBlock != null)
        {
            writer.Write("else");
            if (statement.ElseBlock is TypedNodeStatementIf)
                writer.Write(" ");
            else
                writer.WriteLine();

            writer.WriteLine("{");
            using (writer.Indent())
            {
                EmitStatement(statement.ElseBlock);
            }

            writer.WriteLine("}");
        }
    }

    private void EmitStatementWhile(TypedNodeStatementWhile statement)
    {
        var condition = EmitExpression(statement.Condition);
        writer.WriteLine($"while ({condition})");
        writer.WriteLine("{");
        using (writer.Indent())
        {
            EmitStatement(statement.Block);
        }

        writer.WriteLine("}");
    }

    private string EmitExpression(TypedNodeExpression node)
    {
        return node switch
        {
            TypedNodeExpressionBinary expression => EmitExpressionBinary(expression),
            TypedNodeExpressionUnary expression => EmitExpressionUnary(expression),
            TypedNodeExpressionBoolLiteral expression => expression.Value.Value ? "true" : "false",
            TypedNodeExpressionIntLiteral expression => expression.Value.Value.ToString(),
            TypedNodeExpressionStringLiteral expression => $"(struct string){{ \"{expression.Value.Value}\", {expression.Value.Value.Length} }}",
            TypedNodeExpressionStructLiteral expression => EmitExpressionStructLiteral(expression),
            TypedNodeExpressionMemberAccess expression => EmitExpressionMemberAccess(expression),
            TypedNodeExpressionLocalIdent expression => expression.Value.Ident,
            TypedNodeExpressionModuleIdent expression => expression.Value.Ident,
            _ => throw new ArgumentOutOfRangeException(nameof(node), node, null)
        };
    }

    private string EmitExpressionBinary(TypedNodeExpressionBinary expression)
    {
        var left = EmitExpression(expression.Left);
        var right = EmitExpression(expression.Right);

        return expression.Operation switch
        {
            TypedNodeExpressionBinary.Op.Add => $"({left} + {right})",
            TypedNodeExpressionBinary.Op.Subtract => $"({left} - {right})",
            TypedNodeExpressionBinary.Op.Multiply => $"({left} * {right})",
            TypedNodeExpressionBinary.Op.Divide => $"({left} / {right})",
            TypedNodeExpressionBinary.Op.Modulo => $"({left} % {right})",
            TypedNodeExpressionBinary.Op.Equal => $"({left} == {right})",
            TypedNodeExpressionBinary.Op.NotEqual => $"({left} != {right})",
            TypedNodeExpressionBinary.Op.LessThan => $"({left} < {right})",
            TypedNodeExpressionBinary.Op.LessThanOrEqual => $"({left} <= {right})",
            TypedNodeExpressionBinary.Op.GreaterThan => $"({left} > {right})",
            TypedNodeExpressionBinary.Op.GreaterThanOrEqual => $"({left} >= {right})",
            TypedNodeExpressionBinary.Op.LeftShift => $"({left} << {right})",
            TypedNodeExpressionBinary.Op.RightShift => $"({left} >> {right})",
            TypedNodeExpressionBinary.Op.LogicalAnd => $"({left} && {right})",
            TypedNodeExpressionBinary.Op.LogicalOr => $"({left} || {right})",
            _ => throw new ArgumentOutOfRangeException()
        };
    }

    private string EmitExpressionUnary(TypedNodeExpressionUnary expression)
    {
        var target = EmitExpression(expression.Target);

        return expression.Operation switch
        {
            TypedNodeExpressionUnary.Op.Negate => $"(-{target})",
            TypedNodeExpressionUnary.Op.Invert => $"(!{target})",
            _ => throw new ArgumentOutOfRangeException()
        };
    }

    private string EmitExpressionStructLiteral(TypedNodeExpressionStructLiteral expression)
    {
        var initializerValues = new Dictionary<string, string>();

        foreach (var initializer in expression.Initializers)
        {
            var values = EmitExpression(initializer.Value);
            initializerValues[initializer.Name.Ident] = values;
        }

        var initializerStrings = initializerValues.Select(x => $".{x.Key} = {x.Value}");

        return $"(struct {structTypeNames[(NubTypeStruct)expression.Type]}){{ {string.Join(", ", initializerStrings)} }}";
    }

    private string EmitExpressionMemberAccess(TypedNodeExpressionMemberAccess expression)
    {
        var target = EmitExpression(expression.Target);
        return $"{target}.{expression.Name.Ident}";
    }

    private string CType(NubType node, string? varName = null)
    {
        return node switch
        {
            NubTypeVoid => "void" + (varName != null ? $" {varName}" : ""),
            NubTypeBool => "bool" + (varName != null ? $" {varName}" : ""),
            NubTypeStruct type => $"struct {structTypeNames[type]}" + (varName != null ? $" {varName}" : ""),
            NubTypeSInt type => $"int{type.Width}_t" + (varName != null ? $" {varName}" : ""),
            NubTypeUInt type => $"uint{type.Width}_t" + (varName != null ? $" {varName}" : ""),
            NubTypePointer type => CType(type.To) + (varName != null ? $" *{varName}" : "*"),
            NubTypeString => "struct string" + (varName != null ? $" {varName}" : ""),
            NubTypeFunc type => $"{CType(type.ReturnType)} (*{varName})({string.Join(", ", type.Parameters.Select(p => CType(p)))})",
            _ => throw new ArgumentOutOfRangeException(nameof(node), node, null)
        };
    }
}

internal class IndentedTextWriter
{
    private readonly StringBuilder builder = new();
    private int indentLevel;

    public IDisposable Indent()
    {
        indentLevel++;
        return new IndentScope(this);
    }

    public void WriteLine(string text)
    {
        WriteIndent();
        builder.AppendLine(text);
    }

    public void Write(string text)
    {
        WriteIndent();
        builder.Append(text);
    }

    public void WriteLine()
    {
        builder.AppendLine();
    }

    public override string ToString()
    {
        return builder.ToString();
    }

    private void WriteIndent()
    {
        if (builder.Length > 0)
        {
            var lastChar = builder[^1];
            if (lastChar != '\n' && lastChar != '\r')
                return;
        }

        for (var i = 0; i < indentLevel; i++)
        {
            builder.Append("    ");
        }
    }

    private class IndentScope(IndentedTextWriter writer) : IDisposable
    {
        private bool disposed;

        public void Dispose()
        {
            if (disposed) return;
            writer.indentLevel--;
            disposed = true;
        }
    }
}