ast.jou

# This file defines data structures of the Abstract Syntax Tree. They are
# constructed in parser.jou.
#
# Many classes in this file have .print(), which can be used during debugging. If
# .print() exists, you probably don't need to call .print_with_tree_printer(),
# which is basically an internal detail of the .print() implementation.

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

import "./errors_and_warnings.jou"
import "./types.jou"
import "./types_in_ast.jou"
import "./utils.jou"


@public
enum AstTypeKind:
    Named
    Pointer
    Array


@public
class AstArrayType:
    member_type: AstType*
    length: AstExpression*

    def free(self) -> None:
        self->member_type->free()
        self->length->free()
        free(self->member_type)
        free(self->length)


@public
class AstType:
    kind: AstTypeKind
    location: Location

    union:
        name: byte[100]       # AstTypeKind.Named
        value_type: AstType*  # AstTypeKind.Pointer
        array: AstArrayType   # AstTypeKind.Array

    def is_void(self) -> bool:
        return self->kind == AstTypeKind.Named and strcmp(self->name, "void") == 0

    def is_none(self) -> bool:
        return self->kind == AstTypeKind.Named and strcmp(self->name, "None") == 0

    def is_noreturn(self) -> bool:
        return self->kind == AstTypeKind.Named and strcmp(self->name, "noreturn") == 0

    def print(self, show_lineno: bool) -> None:
        match self->kind:
            case AstTypeKind.Named:
                printf("%s", self->name)
            case AstTypeKind.Pointer:
                self->value_type->print(False)
                printf("*")
            case AstTypeKind.Array:
                self->array.member_type->print(False)
                printf("[<size>]")  # TODO: show the size expression better?

        if show_lineno:
            printf(" [line %d]", self->location.lineno)

    def free(self) -> None:
        match self->kind:
            case AstTypeKind.Pointer:
                self->value_type->free()
                free(self->value_type)
            case AstTypeKind.Array:
                self->array.free()
            case AstTypeKind.Named:
                pass


# Statements and expressions can be printed in a tree.
# To see a tree, run:
#
#   $ jou --parse-only examples/hello.jou
#
class TreePrinter:
    prefix: byte[100]

    # Returned subprinter can be used to print elements "inside" the current line.
    def print_prefix(self, is_last_child: bool) -> TreePrinter:
        subprinter = TreePrinter{}
        if is_last_child:
            printf("%s`--- ", self->prefix)
            snprintf(subprinter.prefix, sizeof subprinter.prefix, "%s  ", self->prefix)
        else:
            printf("%s|--- ", self->prefix)
            snprintf(subprinter.prefix, sizeof subprinter.prefix, "%s| ", self->prefix)
        return subprinter


# Foo.Bar where Foo is an enum
@public
class AstEnumMember:
    enum_name: byte[100]
    member_name: byte[100]


# foo.bar, foo->bar
@public
class AstClassField:
    instance: AstExpression*
    uses_arrow_operator: bool  # distinguishes foo.bar and foo->bar
    field_name: byte[100]

    def free(self) -> None:
        self->instance->free()
        free(self->instance)


# Foo{bar = 1, baz = 2}
@public
class AstInstantiation:
    class_name_location: Location  # TODO: probably not necessary, can use location of the instantiate expression
    class_name: byte[100]
    nfields: int
    field_names: byte[100]*
    field_values: AstExpression*

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("instantiate \"%s\"\n", self->class_name)
        for i = 0; i < self->nfields; i++:
            sub = tp.print_prefix(i == self->nfields - 1)
            printf("field \"%s\": ", self->field_names[i])
            self->field_values[i].print_with_tree_printer(sub)

    def free(self) -> None:
        for i = 0; i < self->nfields; i++:
            self->field_values[i].free()
        free(self->field_names)
        free(self->field_values)


@public
enum AstExpressionKind:
    String
    Int
    Short
    Long
    Byte
    Float
    Double
    Bool
    Null
    Array
    Call  # function call or method call
    Instantiate   # MyClass{x=1, y=2}
    Self    # not a variable lookup, so you can't use 'self' as variable name outside a class
    GetVariable
    GetEnumMember
    GetClassField   # foo.bar, foo->bar
    As
    # unary operators
    SizeOf      # sizeof x
    AddressOf   # &x
    Dereference # *x
    Negate      # -x
    Not         # not x
    PreIncr     # ++x
    PostIncr    # x++
    PreDecr     # --x
    PostDecr    # x--
    # binary operators
    Add         # x+y
    Sub         # x-y
    Mul         # x*y
    Div         # x/y
    Indexing    # x[y]
    Mod         # x % y
    Eq          # x == y
    Ne          # x != y
    Gt          # x > y
    Lt          # x < y
    Ge          # x >= y
    Le          # x <= y
    And         # x and y
    Or          # x or y

@public
class AstExpression:
    location: Location
    kind: AstExpressionKind
    types: ExpressionTypes  # populated in typecheck, zero-initialized before typecheck runs

    union:
        enum_member: AstEnumMember
        class_field: AstClassField
        string: byte*
        int_value: int
        short_value: short
        long_value: long
        byte_value: byte
        bool_value: bool
        call: AstCall
        instantiation: AstInstantiation
        as_: AstAs*  # Must be pointer, because it contains an AstExpression
        array: AstArray
        varname: byte[100]
        float_or_double_text: byte[100]
        operands: AstExpression*  # Only for operators. Length is arity, see get_arity()

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("[line %d", self->location.lineno)
        if self->types.orig_type != NULL:
            printf(", type %s", self->types.orig_type->name)
            # TODO: print casts and other stuff under self->types?
        printf("] ")

        match self->kind:
            case AstExpressionKind.String:
                printf("string ")
                print_string(self->string, strlen(self->string))
                printf("\n")
            case AstExpressionKind.Short:
                printf("short %hd\n", self->short_value)
            case AstExpressionKind.Int:
                printf("int %d\n", self->int_value)
            case AstExpressionKind.Long:
                printf("long %lld\n", self->long_value)
            case AstExpressionKind.Byte:
                printf("byte %d\n", self->byte_value)
            case AstExpressionKind.Float:
                printf("float %s\n", self->float_or_double_text)
            case AstExpressionKind.Double:
                printf("double %s\n", self->float_or_double_text)
            case AstExpressionKind.Bool:
                if self->bool_value:
                    printf("True\n")
                else:
                    printf("False\n")
            case AstExpressionKind.Null:
                printf("NULL\n")
            case AstExpressionKind.Indexing:
                printf("indexing\n")
            case AstExpressionKind.Array:
                printf("array\n")
                for i = 0; i < self->array.length; i++:
                    self->array.items[i].print_with_tree_printer(tp.print_prefix(i == self->array.length-1))
            case AstExpressionKind.Call:
                self->call.print_with_tree_printer(tp)
            case AstExpressionKind.Instantiate:
                self->instantiation.print_with_tree_printer(tp)
            case AstExpressionKind.Self:
                printf("self\n")
            case AstExpressionKind.GetVariable:
                printf("get variable \"%s\"\n", self->varname)
            case AstExpressionKind.GetEnumMember:
                printf(
                    "get member \"%s\" from enum \"%s\"\n",
                    self->enum_member.member_name,
                    self->enum_member.enum_name,
                )
            case AstExpressionKind.GetClassField:
                printf("get class field \"%s\"", self->class_field.field_name)
                if self->class_field.uses_arrow_operator:
                    printf(" with the arrow operator")
                printf("\n")
                self->class_field.instance->print_with_tree_printer(tp.print_prefix(True))
            case AstExpressionKind.As:
                self->as_->print_with_tree_printer(tp)
            case AstExpressionKind.SizeOf:
                printf("sizeof\n")
            case AstExpressionKind.AddressOf:
                printf("address of\n")
            case AstExpressionKind.Dereference:
                printf("dereference\n")
            case AstExpressionKind.Negate:
                printf("negate\n")
            case AstExpressionKind.Not:
                printf("not\n")
            case AstExpressionKind.PreIncr:
                printf("pre-increment\n")
            case AstExpressionKind.PostIncr:
                printf("post-increment\n")
            case AstExpressionKind.PreDecr:
                printf("pre-decrement\n")
            case AstExpressionKind.PostDecr:
                printf("post-decrement\n")
            case AstExpressionKind.Add:
                printf("add\n")
            case AstExpressionKind.Sub:
                printf("sub\n")
            case AstExpressionKind.Mul:
                printf("mul\n")
            case AstExpressionKind.Div:
                printf("div\n")
            case AstExpressionKind.Mod:
                printf("mod\n")
            case AstExpressionKind.Eq:
                printf("eq\n")
            case AstExpressionKind.Ne:
                printf("ne\n")
            case AstExpressionKind.Gt:
                printf("gt\n")
            case AstExpressionKind.Ge:
                printf("ge\n")
            case AstExpressionKind.Lt:
                printf("lt\n")
            case AstExpressionKind.Le:
                printf("le\n")
            case AstExpressionKind.And:
                printf("and\n")
            case AstExpressionKind.Or:
                printf("or\n")

        for i = 0; i < self->get_arity(); i++:
            self->operands[i].print_with_tree_printer(tp.print_prefix(i == self->get_arity()-1))

    def free(self) -> None:
        match self->kind:
            case AstExpressionKind.String:
                free(self->string)
            case AstExpressionKind.Array:
                self->array.free()
            case AstExpressionKind.Call:
                self->call.free()
            case AstExpressionKind.Instantiate:
                self->instantiation.free()
            case AstExpressionKind.GetClassField:
                self->class_field.free()
            case AstExpressionKind.As:
                self->as_->free()
                free(self->as_)
            case _:
                if self->get_arity() != 0:
                    for i = 0; i < self->get_arity(); i++:
                        self->operands[i].free()
                    free(self->operands)

    # arity = number of operands, e.g. 2 for a binary operator such as "+"
    def get_arity(self) -> int:
        match self->kind:
            case (
                AstExpressionKind.SizeOf
                | AstExpressionKind.AddressOf
                | AstExpressionKind.Dereference
                | AstExpressionKind.Negate
                | AstExpressionKind.Not
                | AstExpressionKind.PreIncr
                | AstExpressionKind.PreDecr
                | AstExpressionKind.PostIncr
                | AstExpressionKind.PostDecr
            ):
                return 1
            case (
                AstExpressionKind.Add
                | AstExpressionKind.Sub
                | AstExpressionKind.Mul
                | AstExpressionKind.Div
                | AstExpressionKind.Indexing
                | AstExpressionKind.Mod
                | AstExpressionKind.Eq
                | AstExpressionKind.Ne
                | AstExpressionKind.Gt
                | AstExpressionKind.Lt
                | AstExpressionKind.Ge
                | AstExpressionKind.Le
                | AstExpressionKind.And
                | AstExpressionKind.Or
            ):
                return 2
            case _:
                return 0

    def is_valid_as_a_statement(self) -> bool:
        return (
            self->kind == AstExpressionKind.Call
            or self->kind == AstExpressionKind.PreIncr
            or self->kind == AstExpressionKind.PreDecr
            or self->kind == AstExpressionKind.PostIncr
            or self->kind == AstExpressionKind.PostDecr
        )


# [foo, bar, baz]
@public
class AstArray:
    length: int
    items: AstExpression*

    def free(self) -> None:
        for i = 0; i < self->length; i++:
            self->items[i].free()
        free(self->items)


# foo as bar
@public
class AstAs:
    value: AstExpression
    type: AstType

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("as ")
        self->type.print(True)
        printf("\n")
        self->value.print_with_tree_printer(tp.print_prefix(True))

    def free(self) -> None:
        self->value.free()
        self->type.free()


# foo(arg1, arg2, arg3)
# foo.bar(arg1, arg2, arg3)
# foo->bar(arg1, arg2, arg3)
@public
class AstCall:
    location: Location
    name: byte[100]  # name of function or method
    method_call_self: AstExpression*  # NULL for function calls, the foo of foo.bar() for method calls
    uses_arrow_operator: bool  # distinguishes foo->bar() and foo.bar()
    nargs: int
    args: AstExpression*
    called_signature: Signature*  # populated in typecheck, NULL before typecheck runs, not owned

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("call %s \"%s\"", self->function_or_method(), self->name)
        if self->uses_arrow_operator:
            printf(" with the arrow operator")
        printf("\n")

        if self->method_call_self != NULL:
            sub = tp.print_prefix(self->nargs == 0)
            printf("self: ")
            self->method_call_self->print_with_tree_printer(sub)

        for i = 0; i < self->nargs; i++:
            sub = tp.print_prefix(i == self->nargs - 1)
            printf("argument %d: ", i)
            self->args[i].print_with_tree_printer(sub)

    def free(self) -> None:
        for i = 0; i < self->nargs; i++:
            self->args[i].free()
        free(self->args)
        if self->method_call_self != NULL:
            self->method_call_self->free()
            free(self->method_call_self)

    # Useful for formatting error messages, but not much else.
    def function_or_method(self) -> byte*:
        if self->method_call_self == NULL:
            return "function"
        else:
            return "method"


# assert foo
@public
class AstAssertion:
    condition: AstExpression
    condition_str: byte*

    def free(self) -> None:
        self->condition.free()
        free(self->condition_str)


@public
enum AstStatementKind:
    Import
    ExpressionStatement  # Evaluate an expression. Discard the result.
    Assert
    Pass
    Return
    If
    Match
    WhileLoop
    ForLoop
    Break
    Continue
    DeclareLocalVar # x: SomeType = y (the "= y" is optional)
    ClassField      # x: SomeType (inside a class)
    ClassUnion      # "union:" followed by indented fields inside a class
    Assign          # x = y
    InPlaceAdd      # x += y
    InPlaceSub      # x -= y
    InPlaceMul      # x *= y
    InPlaceDiv      # x /= y
    InPlaceMod      # x %= y
    FunctionDeclare
    FunctionDef
    MethodDef
    Class
    Enum
    GlobalVariableDeclare
    GlobalVariableDef

@public
class AstStatement:
    location: Location
    kind: AstStatementKind

    union:
        expression: AstExpression       # ExpressionStatement, Assert
        if_statement: AstIfStatement
        while_loop: AstConditionAndBody
        for_loop: AstForLoop
        return_value: AstExpression*        # can be NULL
        assignment: AstAssignment           # also used for +=, -= etc
        local_var_declare: AstNameTypeValue # DeclareLocalVar
        global_var_declare: AstGlobalVarDeclare  # GlobalVariableDeclare
        global_var_def: AstGlobalVarDef     # GlobalVariableDef
        class_field: AstNameTypeValue       # ClassField
        union_fields: AstUnionFields        # ClassUnion
        function: AstFunctionOrMethod
        method: AstFunctionOrMethod
        classdef: AstClassDef
        enumdef: AstEnumDef
        assertion: AstAssertion
        match_statement: AstMatchStatement
        import_statement: AstImport

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("[line %d] ", self->location.lineno)
        match self->kind:
            case AstStatementKind.ExpressionStatement:
                printf("expression statement\n")
                self->expression.print_with_tree_printer(tp.print_prefix(True))
            case AstStatementKind.Assert:
                printf("assert \"%s\"\n", self->assertion.condition_str)
                self->assertion.condition.print_with_tree_printer(tp.print_prefix(True))
            case AstStatementKind.Pass:
                printf("pass\n")
            case AstStatementKind.Return:
                printf("return\n")
                if self->return_value != NULL:
                    self->return_value->print_with_tree_printer(tp.print_prefix(True))
            case AstStatementKind.If:
                self->if_statement.print_with_tree_printer(tp)
            case AstStatementKind.ForLoop:
                self->for_loop.print_with_tree_printer(tp)
            case AstStatementKind.Match:
                self->match_statement.print_with_tree_printer(tp)
            case AstStatementKind.WhileLoop:
                printf("while loop\n")
                self->while_loop.print_with_tree_printer(tp, True)
            case AstStatementKind.Break:
                printf("break\n")
            case AstStatementKind.Continue:
                printf("continue\n")
            case AstStatementKind.DeclareLocalVar:
                printf("declare local var ")
                self->local_var_declare.print_with_tree_printer(&tp)
            case AstStatementKind.Assign:
                printf("assign\n")
                self->assignment.print_with_tree_printer(tp)
            case AstStatementKind.InPlaceAdd:
                printf("in-place add\n")
                self->assignment.print_with_tree_printer(tp)
            case AstStatementKind.InPlaceSub:
                printf("in-place sub\n")
                self->assignment.print_with_tree_printer(tp)
            case AstStatementKind.InPlaceMul:
                printf("in-place mul\n")
                self->assignment.print_with_tree_printer(tp)
            case AstStatementKind.InPlaceDiv:
                printf("in-place div\n")
                self->assignment.print_with_tree_printer(tp)
            case AstStatementKind.InPlaceMod:
                printf("in-place mod\n")
                self->assignment.print_with_tree_printer(tp)
            case AstStatementKind.FunctionDeclare:
                printf("declare a function: ")
                self->function.print_with_tree_printer(tp)
            case AstStatementKind.FunctionDef:
                printf("define a function: ")
                self->function.print_with_tree_printer(tp)
            case AstStatementKind.MethodDef:
                printf("define a method: ")
                self->method.print_with_tree_printer(tp)
            case AstStatementKind.Class:
                self->classdef.print_with_tree_printer(tp)
            case AstStatementKind.Enum:
                self->enumdef.print_with_tree_printer(tp)
            case AstStatementKind.GlobalVariableDeclare:
                printf("declare global var ")
                self->global_var_declare.print()
                printf("\n")
            case AstStatementKind.GlobalVariableDef:
                printf("define global var ")
                self->global_var_def.print()
                printf("\n")
            case AstStatementKind.Import:
                self->import_statement.print()
            case AstStatementKind.ClassField:
                printf("class field ")
                self->class_field.print_with_tree_printer(NULL)
                printf("\n")
            case AstStatementKind.ClassUnion:
                printf("union\n")
                self->union_fields.print_with_tree_printer(tp)

    def free(self) -> None:
        match self->kind:
            case AstStatementKind.ExpressionStatement:
                self->expression.free()
            case AstStatementKind.Assert:
                self->assertion.free()
            case AstStatementKind.Return:
                if self->return_value != NULL:
                    self->return_value->free()
                    free(self->return_value)
            case AstStatementKind.If:
                self->if_statement.free()
            case AstStatementKind.WhileLoop:
                self->while_loop.free()
            case AstStatementKind.ForLoop:
                self->for_loop.free()
            case AstStatementKind.Match:
                self->match_statement.free()
            case AstStatementKind.DeclareLocalVar:
                self->local_var_declare.free()
            case AstStatementKind.GlobalVariableDeclare:
                self->global_var_declare.free()
            case AstStatementKind.GlobalVariableDef:
                self->global_var_def.free()
            case AstStatementKind.ClassField:
                self->class_field.free()
            case (
                AstStatementKind.Assign
                | AstStatementKind.InPlaceAdd
                | AstStatementKind.InPlaceSub
                | AstStatementKind.InPlaceMul
                | AstStatementKind.InPlaceDiv
                | AstStatementKind.InPlaceMod
            ):
                self->assignment.free()
            case AstStatementKind.FunctionDeclare | AstStatementKind.FunctionDef:
                self->function.free()
            case AstStatementKind.MethodDef:
                self->method.free()
            case AstStatementKind.Class:
                self->classdef.free()
            case AstStatementKind.Enum:
                self->enumdef.free()
            case AstStatementKind.Import:
                self->import_statement.free()
            case AstStatementKind.ClassUnion:
                self->union_fields.free()
            case AstStatementKind.Pass | AstStatementKind.Break | AstStatementKind.Continue:
                pass


# Useful for e.g. "while condition: body", "if condition: body"
@public
class AstConditionAndBody:
    condition: AstExpression
    body: AstBody

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{}, True)

    def print_with_tree_printer(self, tp: TreePrinter, body_is_last_sub_item: bool) -> None:
        sub = tp.print_prefix(False)
        printf("condition: ")
        self->condition.print_with_tree_printer(sub)

        sub = tp.print_prefix(body_is_last_sub_item)
        printf("body:\n")
        self->body.print_with_tree_printer(sub)

    def free(self) -> None:
        self->condition.free()
        self->body.free()


# foo = bar
@public
class AstAssignment:
    target: AstExpression
    value: AstExpression

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        self->target.print_with_tree_printer(tp.print_prefix(False))
        self->value.print_with_tree_printer(tp.print_prefix(True))

    def free(self) -> None:
        self->target.free()
        self->value.free()


# if foo:
#     ...
# elif bar:
#     ...
# elif baz:
#     ...
# else:
#     ...
@public
class AstIfStatement:
    if_and_elifs: AstConditionAndBody*
    n_if_and_elifs: int  # At least 1 (the if statement). The rest, if any, are elifs.
    else_body: AstBody    # Empty if there is no else

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("if\n")
        for i = 0; i < self->n_if_and_elifs; i++:
            self->if_and_elifs[i].print_with_tree_printer(tp, i == self->n_if_and_elifs - 1 and self->else_body.nstatements == 0)

        if self->else_body.nstatements > 0:
            sub = tp.print_prefix(True)
            printf("else body:\n")
            self->else_body.print_with_tree_printer(sub)

    def free(self) -> None:
        for i = 0; i < self->n_if_and_elifs; i++:
            self->if_and_elifs[i].free()
        free(self->if_and_elifs)
        self->else_body.free()


# match match_obj with func:
#     case ...:
#         ...
#     case ...:
#         ...
@public
class AstMatchStatement:
    match_obj: AstExpression
    func_name: byte[100]  # empty if there's no "with foo"
    func_signature: Signature  # populated in typecheck, zero-initialized before typecheck runs
    cases: AstCase*
    ncases: int
    case_underscore: AstBody*  # body of "case _" (always last), NULL if no "case _"
    case_underscore_location: Location  # not meaningful if case_underscore == NULL

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("match\n")
        for i = 0; i < self->ncases; i++:
            self->cases[i].print_with_tree_printer(tp, i == self->ncases - 1 and self->case_underscore == NULL)

        if self->case_underscore != NULL:
            sub = tp.print_prefix(True)
            printf("[line %d] body of case _:\n", self->case_underscore_location.lineno)
            self->case_underscore->print_with_tree_printer(sub)

    def free(self) -> None:
        self->match_obj.free()
        if self->func_signature.name[0] != '\0':
            self->func_signature.free()
        for i = 0; i < self->ncases; i++:
            self->cases[i].free()
        free(self->cases)
        if self->case_underscore != NULL:
            self->case_underscore->free()
            free(self->case_underscore)


# case case_obj1 | case_obj2 | case_obj3:
#     body
@public
class AstCase:
    case_objs: AstExpression*
    n_case_objs: int
    body: AstBody

    def print_with_tree_printer(self, tp: TreePrinter, is_last_case: bool) -> None:
        for i = 0; i < self->n_case_objs; i++:
            sub = tp.print_prefix(False)
            printf("case_obj: ")
            self->case_objs[i].print_with_tree_printer(sub)

        sub = tp.print_prefix(is_last_case)
        printf("body:\n")
        self->body.print_with_tree_printer(sub)

    def free(self) -> None:
        for i = 0; i < self->n_case_objs; i++:
            self->case_objs[i].free()
        free(self->case_objs)
        self->body.free()


# for init; cond; incr:
#     ...body...
@public
class AstForLoop:
    init: AstStatement*   # may be NULL
    cond: AstExpression*  # may be NULL
    incr: AstStatement*   # may be NULL
    body: AstBody

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("for loop\n")

        if self->init != NULL:
            sub = tp.print_prefix(False)
            printf("init: ")
            self->init->print_with_tree_printer(sub)

        if self->cond != NULL:
            sub = tp.print_prefix(False)
            printf("cond: ")
            self->cond->print_with_tree_printer(sub)

        if self->incr != NULL:
            sub = tp.print_prefix(False)
            printf("incr: ")
            self->incr->print_with_tree_printer(sub)

        sub = tp.print_prefix(True)
        printf("body:\n")
        self->body.print_with_tree_printer(sub)

    def free(self) -> None:
        if self->init != NULL:
            self->init->free()
            free(self->init)
        if self->cond != NULL:
            self->cond->free()
            free(self->cond)
        if self->incr != NULL:
            self->incr->free()
            free(self->incr)
        self->body.free()


# name: type = value
@public
class AstNameTypeValue:
    name: byte[100]
    name_location: Location
    type: AstType
    value: AstExpression*  # can be NULL

    def print(self) -> None:
        tp = TreePrinter{}
        self->print_with_tree_printer(&tp)

    # tp can be set to NULL, in that case no trailing newline is printed
    def print_with_tree_printer(self, tp: TreePrinter*) -> None:
        printf("%s: ", self->name)
        self->type.print(True)
        if tp == NULL:
            assert self->value == NULL
        else:
            printf("\n")
            if self->value != NULL:
                sub = tp->print_prefix(True)
                printf("initial value: ")
                self->value->print_with_tree_printer(sub)

    def free(self) -> None:
        self->type.free()
        if self->value != NULL:
            self->value->free()
            free(self->value)


# declare global foo: int
class AstGlobalVarDeclare:
    name: byte[100]
    type: AstType
    public: bool  # is it decorated with @public
    used: bool  # used to detect unused global variables

    def print(self) -> None:
        printf("%s: ", self->name)
        self->type.print(True)
        printf("\n")

    def free(self) -> None:
        self->type.free()


# global foo: int
# TODO: support specifying an initial value
@public
class AstGlobalVarDef:
    name: byte[100]
    type: AstType
    public: bool  # is it decorated with @public
    used: bool  # used to detect unused global variables

    def print(self) -> None:
        printf("%s: ", self->name)
        self->type.print(True)
        printf("\n")

    def free(self) -> None:
        self->type.free()


# typically multiple indented lines after a ":" at end of line
@public
class AstBody:
    statements: AstStatement*
    nstatements: int

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        for i = 0; i < self->nstatements; i++:
            self->statements[i].print_with_tree_printer(tp.print_prefix(i == self->nstatements - 1))

    def free(self) -> None:
        for i = 0; i < self->nstatements; i++:
            self->statements[i].free()
        free(self->statements)


# function name and parameters in "def" or "declare"
@public
class AstSignature:
    name_location: Location
    name: byte[100]  # name of function or method, after "def" keyword
    args: AstNameTypeValue*
    nargs: int
    takes_varargs: bool  # True for functions like printf()
    return_type: AstType

    def print(self) -> None:
        printf("%s(", self->name)
        for i = 0; i < self->nargs; i++:
            if i != 0:
                printf(", ")

            if (
                strcmp(self->args[i].name, "self") == 0
                and self->args[i].type.kind == AstTypeKind.Named
                and self->args[i].type.name[0] == '\0'
            ):
                # self with implicitly given type
                printf("self")
            else:
                self->args[i].print_with_tree_printer(NULL)

        if self->takes_varargs:
            if self->nargs != 0:
                printf(", ")
            printf("...")

        printf(") -> ")
        self->return_type.print(True)
        printf("\n")

    def free(self) -> None:
        for i = 0; i < self->nargs; i++:
            self->args[i].free()
        free(self->args)
        self->return_type.free()


# import "./foo.jou"
@public
class AstImport:
    specified_path: byte*  # Path in jou code e.g. "stdlib/io.jou"
    resolved_path: byte*  # Absolute path or relative to current working directory e.g. "/home/akuli/jou/stdlib/io.jou"
    used: bool  # For warning messages

    def print(self) -> None:
        printf(
            "import \"%s\", which resolves to \"%s\"\n",
            self->specified_path, self->resolved_path,
        )

    def free(self) -> None:
        free(self->specified_path)
        free(self->resolved_path)


# Represents the AST of one Jou file.
@public
class AstFile:
    path: byte*  # not owned
    is_main_file: bool  # is this the file passed in Jou compiler command? False means imported file
    body: AstBody
    types: FileTypes  # populated in typecheck, zero-initialized before typecheck runs

    def print(self) -> None:
        printf("===== AST for file \"%s\" =====\n", self->path)
        for i = 0; i < self->body.nstatements; i++:
            self->body.statements[i].print()

    def free(self) -> None:
        self->body.free()
        self->types.free()  # also works if self->types is zero-initialized


# def foo() -> bar:
#     ...
@public
class AstFunctionOrMethod:
    ast_signature: AstSignature
    body: AstBody
    types: FunctionOrMethodTypes  # populated in typecheck, zero-initialized before typecheck runs
    public: bool  # is it decorated with @public
    used: bool  # used to detect unused functions

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        self->ast_signature.print()
        self->body.print_with_tree_printer(tp)

    def free(self) -> None:
        self->ast_signature.free()
        self->body.free()
        if self->types.signature.name[0] != '\0':
            # type-checking ran, the types are set and they need freeing
            self->types.free()


# union:
#     foo: type1
#     bar: type2
@public
class AstUnionFields:
    fields: AstNameTypeValue*
    nfields: int

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        for i = 0; i < self->nfields; i++:
            subprinter = tp.print_prefix(i == self->nfields-1)
            self->fields[i].print_with_tree_printer(&subprinter)  # TODO: does this need to be optional/pointer?

    def free(self) -> None:
        for i = 0; i < self->nfields; i++:
            self->fields[i].free()
        free(self->fields)


# class Foo:
#     ...members...
@public
class AstClassDef:
    name: byte[100]
    body: AstBody*
    type: Type*  # populated in typecheck, NULL before typecheck runs
    public: bool  # is it decorated with @public
    used: bool  # used to detect unused classes

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("class \"%s\"\n", self->name)
        self->body->print_with_tree_printer(tp)

    def free(self) -> None:
        self->body->free()
        free(self->body)


# enum Foo:
#     Member1
#     Member2
#     Member3
@public
class AstEnumDef:
    name: byte[100]
    name_location: Location
    member_count: int
    member_names: byte[100]*
    type: Type*  # populated in typecheck, NULL before typecheck runs
    public: bool  # is it decorated with @public
    used: bool  # used to detect unused enums

    def print(self) -> None:
        self->print_with_tree_printer(TreePrinter{})

    def print_with_tree_printer(self, tp: TreePrinter) -> None:
        printf("enum \"%s\" with %d members\n", self->name, self->member_count)
        for i = 0; i < self->member_count; i++:
            tp.print_prefix(i == self->member_count-1)
            puts(self->member_names[i])

    def free(self) -> None:
        free(self->member_names)