main.jou

import "stdlib/io.jou"
import "stdlib/mem.jou"
import "stdlib/errno.jou"
import "stdlib/str.jou"
import "stdlib/process.jou"

import "./cf_graph.jou"
import "./command_line_args.jou"
import "./build_cf_graph.jou"
import "./evaluate.jou"
import "./run.jou"
import "./codegen.jou"
import "./llvm.jou"
import "./typecheck/common.jou"
import "./typecheck/step1_create_types.jou"
import "./typecheck/step2_populate_types.jou"
import "./typecheck/step3_function_and_method_bodies.jou"
import "./target.jou"
import "./types.jou"
import "./token.jou"
import "./parser.jou"
import "./paths.jou"
import "./errors_and_warnings.jou"
import "./tokenizer.jou"
import "./ast.jou"


def open_the_file(path: byte*, import_location: Location*) -> FILE*:
    f = fopen(path, "rb")
    if f == NULL:
        msg: byte[500]
        if import_location != NULL:
            snprintf(msg, sizeof(msg), "cannot import from \"%s\": %s", path, strerror(get_errno()))
            fail(*import_location, msg)
        else:
            snprintf(msg, sizeof(msg), "cannot open file: %s", strerror(get_errno()))
            fail(Location{path=path}, msg)
    return f


def defines_main(ast: AstFile*) -> bool:
    for i = 0; i < ast->body.nstatements; i++:
        s = &ast->body.statements[i]
        if s->kind == AstStatementKind.Function and strcmp(s->function.signature.name, "main") == 0:
            return True
    return False


def statement_conflicts_with_an_import(stmt: AstStatement*, importsym: ExportSymbol*) -> bool:
    match stmt->kind:
        case AstStatementKind.Function:
            return (
                importsym->kind == ExportSymbolKind.Function
                and strcmp(importsym->name, stmt->function.signature.name) == 0
            )
        case AstStatementKind.GlobalVariableDeclaration | AstStatementKind.GlobalVariableDefinition:
            return (
                importsym->kind == ExportSymbolKind.GlobalVar
                and strcmp(importsym->name, stmt->var_declaration.name) == 0
            )
        case AstStatementKind.Class:
            return (
                importsym->kind == ExportSymbolKind.Type
                and strcmp(importsym->name, stmt->classdef.name) == 0
            )
        case AstStatementKind.Enum:
            return (
                importsym->kind == ExportSymbolKind.Type
                and strcmp(importsym->name, stmt->enumdef.name) == 0
            )
        case _:
            assert False


def print_llvm_ir(module: LLVMModule*, is_optimized: bool) -> None:
    if is_optimized:
        opt_or_unopt = "Optimized"
    else:
        opt_or_unopt = "Unoptimized"

    len = 0L
    filename = LLVMGetSourceFileName(module, &len)
    printf("===== %s LLVM IR for file \"%.*s\" =====\n", opt_or_unopt, len as int, filename)

    s = LLVMPrintModuleToString(module)
    puts(s)
    LLVMDisposeMessage(s)


class FileState:
    path: byte*  # owned
    ast: AstFile
    types: FileTypes
    module: LLVMModule*

    def free(self) -> None:
        self->ast.free()
        free(self->path)
        self->types.free()

    def add_imported_symbol(self, es: ExportSymbol*, imp: AstImport*) -> None:
        for i = 0; i < self->ast.body.nstatements; i++:
            if statement_conflicts_with_an_import(&self->ast.body.statements[i], es):
                match es->kind:
                    case ExportSymbolKind.Function:
                        wat = "function"
                    case ExportSymbolKind.GlobalVar:
                        wat = "global variable"
                    case ExportSymbolKind.Type:
                        wat = "type"

                msg: byte[500]
                snprintf(msg, sizeof msg, "a %s named '%s' already exists", wat, es->name)
                fail(self->ast.body.statements[i].location, msg)

        match es->kind:
            case ExportSymbolKind.Function:
                self->types.functions = realloc(self->types.functions, sizeof(self->types.functions[0]) * (self->types.nfunctions + 1))
                assert self->types.functions != NULL
                self->types.functions[self->types.nfunctions++] = SignatureAndUsedPtr{
                    signature = es->funcsignature.copy(),
                    usedptr = &imp->used,
                }
            case ExportSymbolKind.Type:
                self->types.types = realloc(self->types.types, sizeof(self->types.types[0]) * (self->types.ntypes + 1))
                assert self->types.types != NULL
                self->types.types[self->types.ntypes++] = TypeAndUsedPtr{
                    type = es->type,
                    usedptr = &imp->used,
                }
            case ExportSymbolKind.GlobalVar:
                g = GlobalVariable{
                    type = es->type,
                    usedptr = &imp->used,
                }

                assert strlen(es->name) < sizeof g.name
                strcpy(g.name, es->name)

                self->types.globals = realloc(self->types.globals, sizeof(self->types.globals[0]) * (self->types.nglobals + 1))
                assert self->types.globals != NULL
                self->types.globals[self->types.nglobals++] = g

    def warn_about_unused_imports(self) -> None:
        for imp = self->ast.imports; imp < &self->ast.imports[self->ast.nimports]; imp++:
            if not imp->used:
                msg: byte[500]
                snprintf(msg, sizeof msg, "\"%s\" imported but not used", imp->specified_path)
                show_warning(imp->location, msg)

    def build_cf_graphs(self) -> CfGraphFile:
        if command_line_args.verbosity >= 1:
            printf("Building Control Flow Graphs: %s\n", self->path)

        cf_graphs = build_control_flow_graphs(&self->ast, &self->types)
        self->warn_about_unused_imports()

        if command_line_args.verbosity >= 2:
            cf_graphs.print()

        # TODO: implement this
        #if command_line_args.verbosity >= 1:
        #    printf("Analyzing CFGs: %s\n", self->path)
        #simplify_control_flow_graphs(&cf_graphs)
        #if command_line_args.verbosity >= 2:
        #    cf_graphs.print()

        return cf_graphs

    def build_llvm_ir(self, cf_graphs: CfGraphFile*) -> LLVMModule*:
        if command_line_args.verbosity >= 1:
            printf("Building LLVM IR: %s\n", self->path)

        mod = codegen(cf_graphs, &self->types)

        if command_line_args.verbosity >= 2:
            print_llvm_ir(mod, False)

        # If this fails, it is not just users writing dumb code, it is a bug in this compiler.
        # This compiler should always fail with an error elsewhere, or generate valid LLVM IR.
        LLVMVerifyModule(mod, LLVMVerifierFailureAction.AbortProcess, NULL)

        if command_line_args.optlevel != 0:
            if command_line_args.verbosity >= 1:
                printf("Optimizing %s (level %d)\n", self->path, command_line_args.optlevel)
            optimize(mod, command_line_args.optlevel)
            if command_line_args.verbosity >= 2:
                print_llvm_ir(mod, True)

        return mod


class CompileState:
    stdlib_path: byte*
    files: FileState*
    nfiles: int

    def free(self) -> None:
        for fs = self->files; fs < &self->files[self->nfiles]; fs++:
            fs->free()
        free(self->files)

    def find_file(self, path: byte*) -> FileState*:
        for fs = self->files; fs < &self->files[self->nfiles]; fs++:
            if strcmp(fs->path, path) == 0:
                return fs
        return NULL

    def parse(self, filename: byte*, import_location: Location*) -> None:
        if self->find_file(filename) != NULL:
            # already parsed
            return

        fs = FileState{path = strdup(filename)}

        if command_line_args.verbosity >= 1:
            printf("Tokenizing %s\n", filename)
        tokens = tokenize(fs.path, import_location)

        if command_line_args.verbosity >= 2:
            print_tokens(tokens)

        if command_line_args.verbosity >= 1:
            printf("Parsing %s\n", filename)
        fs.ast = parse(tokens, self->stdlib_path)
        free_tokens(tokens)

        if command_line_args.verbosity >= 1:
            printf("Evaluating compile-time if statements in %s\n", filename)
        evaluate_compile_time_if_statements(&fs.ast.body)

        if command_line_args.verbosity >= 2:
            fs.ast.print()

        # If it's not the file passed on command line, it shouldn't define main()
        if strcmp(filename, command_line_args.infile) != 0 and defines_main(&fs.ast):
            # Set error location to import, so user immediately knows which file
            # imports something that defines main().
            assert import_location != NULL
            fail(*import_location, "imported file should not have `main` function")

        self->files = realloc(self->files, sizeof(self->files[0]) * (self->nfiles + 1))
        assert self->files != NULL
        self->files[self->nfiles++] = fs

        for imp = fs.ast.imports; imp < &fs.ast.imports[fs.ast.nimports]; imp++:
            self->parse(imp->resolved_path, &imp->location)  # recursive call

    def parse_from_stdlib(self, filename: byte*) -> None:
        path = malloc(strlen(self->stdlib_path) + strlen(filename) + 123)
        sprintf(path, "%s/%s", self->stdlib_path, filename)
        self->parse(path, NULL)
        free(path)

    # Each step of type checking produces exported symbols that other files can import.
    # This method hands them over to the importing files.
    def exports_to_imports(self, pending_exports: ExportSymbol**) -> None:
        for to = self->files; to < &self->files[self->nfiles]; to++:
            seen_before: FileState** = NULL
            seen_before_len = 0

            for imp = to->ast.imports; imp < &to->ast.imports[to->ast.nimports]; imp++:
                from = self->find_file(imp->resolved_path)
                assert from != NULL
                from_index = ((from as long) - (self->files as long)) / sizeof(self->files[0])

                if from == to:
                    fail(imp->location, "the file itself cannot be imported")

                for i = 0; i < seen_before_len; i++:
                    if seen_before[i] == from:
                        msg: byte[500]
                        snprintf(msg, sizeof msg, "file \"%s\" is imported twice", imp->specified_path)
                        fail(imp->location, msg)

                seen_before = realloc(seen_before, sizeof(seen_before[0]) * (seen_before_len + 1))
                seen_before[seen_before_len++] = from

                for es = pending_exports[from_index]; es->name[0] != '\0'; es++:
                    if command_line_args.verbosity >= 2:
                        match es->kind:
                            case ExportSymbolKind.Function:
                                kindstr = "function"
                            case ExportSymbolKind.GlobalVar:
                                kindstr = "global var"
                            case ExportSymbolKind.Type:
                                kindstr = "type"
                        printf("Adding imported %s %s: %s --> %s\n",
                            kindstr, es->name, from->path, to->path)

                    to->add_imported_symbol(es, imp)

            free(seen_before)

        for i = 0; i < self->nfiles; i++:
            for es = pending_exports[i]; es->name[0] != '\0'; es++:
                es->free()
            free(pending_exports[i])
            pending_exports[i] = NULL

    def typecheck_all_files(self) -> None:
        if command_line_args.verbosity >= 1:
            printf("Type-checking...\n")

        pending_exports: ExportSymbol** = malloc(sizeof(pending_exports[0]) * self->nfiles)

        for i = 0; i < self->nfiles; i++:
            if command_line_args.verbosity >= 1:
                printf("  step 1: %s\n", self->files[i].path)
            pending_exports[i] = typecheck_step1_create_types(&self->files[i].types, &self->files[i].ast)

        self->exports_to_imports(pending_exports)

        for i = 0; i < self->nfiles; i++:
            if command_line_args.verbosity >= 1:
                printf("  step 2: %s\n", self->files[i].path)
            pending_exports[i] = typecheck_step2_populate_types(&self->files[i].types, &self->files[i].ast)

        self->exports_to_imports(pending_exports)

        for i = 0; i < self->nfiles; i++:
            if command_line_args.verbosity >= 1:
                printf("  step 3: %s\n", self->files[i].path)
            typecheck_step3_function_and_method_bodies(&self->files[i].types, &self->files[i].ast)
            if command_line_args.verbosity >= 3:
                # AST now contains types. If very verbose, print them.
                self->files[i].ast.print()

        free(pending_exports)


def optimize(module: LLVMModule*, level: int) -> None:
    assert 1 <= level and level <= 3

    pm = LLVMCreatePassManager()

    # The default settings should be fine for Jou because they work well for
    # C and C++, and Jou is quite similar to C.
    pmbuilder = LLVMPassManagerBuilderCreate()
    LLVMPassManagerBuilderSetOptLevel(pmbuilder, level)
    LLVMPassManagerBuilderPopulateModulePassManager(pmbuilder, pm)
    LLVMPassManagerBuilderDispose(pmbuilder)

    LLVMRunPassManager(pm, module)
    LLVMDisposePassManager(pm)


def compile_llvm_ir_to_object_file(module: LLVMModule*) -> byte*:
    len = 0L
    objname = get_filename_without_jou_suffix(LLVMGetSourceFileName(module, &len))

    objname = realloc(objname, strlen(objname) + 10)
    if WINDOWS:
        strcat(objname, ".obj")
    else:
        strcat(objname, ".o")

    path = get_path_to_file_in_jou_compiled(objname)
    free(objname)

    if command_line_args.verbosity >= 1:
        printf("Emitting object file: %s\n", path)

    tmppath = strdup(path)
    error: byte* = NULL
    if LLVMTargetMachineEmitToFile(target.target_machine, module, tmppath, LLVMCodeGenFileType.ObjectFile, &error) != 0:
        assert error != NULL
        fprintf(stderr, "failed to emit object file \"%s\": %s\n", path, error)
        exit(1)
    free(tmppath)

    assert error == NULL
    return path


def tokenize_only(path: byte*) -> None:
    tokens = tokenize(command_line_args.infile, NULL)
    print_tokens(tokens)
    free_tokens(tokens)


def parse_only(path: byte*, stdlib_path: byte*) -> None:
    tokens = tokenize(command_line_args.infile, NULL)
    ast = parse(tokens, stdlib_path)
    ast.print()
    ast.free()
    free_tokens(tokens)


def main(argc: int, argv: byte**) -> int:
    init_target()
    init_types()
    stdlib = find_stdlib()

    parse_command_line_args(argc, argv)

    if command_line_args.verbosity >= 2:
        printf("Target triple: %s\n", target.triple)
        printf("Data layout: %s\n", target.data_layout)

    if command_line_args.tokenize_only:
        tokenize_only(command_line_args.infile)
        return 0

    if command_line_args.parse_only:
        parse_only(command_line_args.infile, stdlib)
        return 0

    compst = CompileState{ stdlib_path = stdlib }

    # This part recursively finds and parses imported files.
    compst.parse_from_stdlib("_assert_fail.jou")
    if WINDOWS or MACOS or NETBSD:
        compst.parse_from_stdlib("_jou_startup.jou")
    compst.parse(command_line_args.infile, NULL)

    compst.typecheck_all_files()

    objpaths: byte** = calloc(sizeof objpaths[0], compst.nfiles + 1)
    for i = 0; i < compst.nfiles; i++:
        cf_graphs = compst.files[i].build_cf_graphs()

        llvm_ir = compst.files[i].build_llvm_ir(&cf_graphs)
        cf_graphs.free()

        objpaths[i] = compile_llvm_ir_to_object_file(llvm_ir)
        LLVMDisposeModule(llvm_ir)

    # Check for missing main() as late as possible, so that other errors come first.
    # This way Jou users can work on other functions before main() function is written.
    mainfile = compst.find_file(command_line_args.infile)
    assert mainfile != NULL
    if not defines_main(&mainfile->ast):
        fail(Location{path=mainfile->path}, "missing `main` function to execute the program")

    compst.free()
    free(stdlib)

    exepath = decide_exe_path()
    run_linker(objpaths, exepath)

    for i = 0; objpaths[i] != NULL; i++:
        free(objpaths[i])
    free(objpaths)

    ret = 0
    if command_line_args.outfile == NULL:
        if command_line_args.verbosity >= 1:
            printf("Run: %s\n", exepath)
        ret = run_exe(exepath, command_line_args.valgrind)

    free(exepath)

    # not really necessary, but makes valgrind much happier
    free_global_type_state()
    cleanup_target()

    return ret