ast.cup

import "compiler/tokens.cup"
import "compiler/types.cup"
import "std/vector.cup"

enum NodeType {
    // Unary
    AST_NEG,
    AST_NOT,
    AST_BWINV,
    AST_ADDROF,
    AST_DEREF,
    // Binary
    AST_PLUS,
    AST_MINUS,
    AST_MUL,
    AST_DIV,
    AST_MOD,
    AST_LSHIFT,
    AST_RSHIFT,
    AST_AND,
    AST_BWAND,
    AST_OR,
    AST_BWOR,
    AST_XOR,
    // Comparison
    AST_EQ,
    AST_NEQ,
    AST_LT,
    AST_LEQ,
    AST_GT,
    AST_GEQ,
    // Misc.
    AST_ASSIGN,
    AST_MEMBER,
    // AST types
    AST_LITERAL,
    AST_CONSTANT,
    AST_CONVERT,
    AST_FUNCCALL,
    AST_CONDITIONAL,
    AST_IF,
    AST_WHILE,
    AST_DEFER,
    AST_FOR,
    AST_SWITCH,
    AST_CASE,
    AST_VARDECL,
    AST_LOCAL_VAR,
    AST_GLOBAL_VAR,
    AST_RETURN,
    AST_BREAK,
    AST_FUNC,
    AST_BUILTIN,
    AST_PROGRAM,
    AST_BLOCK,

    NUM_NODE_TYPES,
};

struct Variable {
    name: char *;
    typ: Type *;
    offset: int;
};

struct Node {
    typ: int;       // NodeType
    etyp: Type*;    // Expression type

    union {
        binary: struct {
            lhs: Node *;
            rhs: Node *;
        };

        unary: Node *;

        func: struct {
            name: char *;
            body: Node *;
            max_locals_size: int;
            args: Vector *; // Vector<Variable>
            is_defined: int;

            is_method: int;
            method_of: Type*;

            is_constructor: int;
        };

        block: struct {
            children: Vector *; // Vector<Node *>
            locals: Vector *;   // Vector<Variable>
            locals_size: int;
        };

        literal: union {
            as_int: int;
            as_char: char;
            as_string: char *;
        };

        var_decl: struct {
            var: Variable;
            init: Node *;
        };

        assign: struct {
            lhs: Node *;
            rhs: Node *;
        };

        conditional: struct {
            cond: Node *;
            then: Node *;
            els: Node *;
        };

        // `loop` is keyword in rust, syntax highlighting breaks
        looop: struct {
            cond: Node *;
            body: Node *;
            // for loop:
            init: Node *;
            step: Node *;
        };

        variable: Variable *;

        call: struct {
            func: Node *;
            args: Vector *; // Vector<Node *>
        };

        member: struct {
            obj: Node *;
            offset: int;
            is_ptr: int;
        };

        constant: struct {
            name: char *;
            value: Node *; // Must be int literal
        };

        switch_stmt: struct {
            expr: Node *;
            cases: Vector *; // Vector<CaseStatement *>
            defolt: Vector *; // Vector<Node *>
        };

        case_stmt: struct {
            val: i64;
            children: Vector *;
        };
    };
};

let node_counter = 0;

fn node_new(typ: int): Node* {
    let node: Node* = malloc(sizeof(Node));
    ++node_counter;
    node.typ = typ;
    return node;
}

fn node_from_int_literal(val: int): Node* {
    let node: Node* = node_new(AST_LITERAL);
    node.etyp = type_new(TYPE_INT);
    node.literal.as_int = val;
    return node;
}

fn type_check_unary(node: Node*, token: Token*): Node*;
fn type_check_binary(node: Node*, token: Token*): Node*;

fn node_new_binop(typ: int, lhs: Node*, rhs: Node*): Node* {
    let node = node_new(typ);
    node.binary.lhs = lhs;
    node.binary.rhs = rhs;
    return type_check_binary(node, null);
}

fn convert_to_float(node: Node*): Node* {
    if (is_float_type(node.etyp))
        return node;
    let conv = node_new(AST_CONVERT);
    conv.etyp = type_new(TYPE_F64);
    conv.unary = node;
    return conv;
}

fn convert_to_int(node: Node*): Node* {
    if (is_int_type(node.etyp))
        return node;
    if (!is_float_type(node.etyp))
        die2(here, "Cannot convert non-numeric type to int: ", create_type_string(node.etyp));
    let conv = node_new(AST_CONVERT);
    conv.etyp = type_new(TYPE_INT);
    conv.unary = node;
    return conv;
}

fn variable_new(name: char*, typ: Type*, offset: int): Variable* {
    let v: Variable* = malloc(sizeof(Variable));
    v.name = name;
    v.typ = typ;
    v.offset = offset;
    return v;
}

fn block_add_child(block: Node*, child: Node*) {
    if (block.block.children == null)
        block.block.children = vector_new();
    block.block.children::push(child);
}

// TODO: Careful here, the input type here is the same as `type_to_string`
fn node_type_to_string(typ: int): char* {
    static_assert(NUM_NODE_TYPES == 45, "Exhaustive match in node_type_to_string");
    switch (typ) {
        case AST_NEG:         return "AST_NEG";
        case AST_NOT:         return "AST_NOT";
        case AST_BWINV:       return "AST_BWINV";
        case AST_ADDROF:      return "AST_ADDROF";
        case AST_DEREF:       return "AST_DEREF";
        case AST_PLUS:        return "AST_PLUS";
        case AST_MINUS:       return "AST_MINUS";
        case AST_MUL:         return "AST_MUL";
        case AST_DIV:         return "AST_DIV";
        case AST_MOD:         return "AST_MOD";
        case AST_LSHIFT:      return "AST_LSHIFT";
        case AST_RSHIFT:      return "AST_RSHIFT";
        case AST_AND:         return "AST_AND";
        case AST_BWAND:       return "AST_BWAND";
        case AST_OR:          return "AST_OR";
        case AST_BWOR:        return "AST_BWOR";
        case AST_XOR:         return "AST_XOR";
        case AST_EQ:          return "AST_EQ";
        case AST_NEQ:         return "AST_NEQ";
        case AST_LT:          return "AST_LT";
        case AST_LEQ:         return "AST_LEQ";
        case AST_GT:          return "AST_GT";
        case AST_GEQ:         return "AST_GEQ";
        case AST_ASSIGN:      return "AST_ASSIGN";
        case AST_MEMBER:      return "AST_MEMBER";
        case AST_LITERAL:     return "AST_LITERAL";
        case AST_CONSTANT:    return "AST_CONSTANT";
        case AST_CONVERT:     return "AST_CONVERT";
        case AST_FUNCCALL:    return "AST_FUNCCALL";
        case AST_CONDITIONAL: return "AST_CONDITIONAL";
        case AST_IF:          return "AST_IF";
        case AST_WHILE:       return "AST_WHILE";
        case AST_DEFER:       return "AST_DEFER";
        case AST_FOR:         return "AST_FOR";
        case AST_SWITCH:      return "AST_SWITCH";
        case AST_CASE:        return "AST_CASE";
        case AST_VARDECL:     return "AST_VARDECL";
        case AST_LOCAL_VAR:   return "AST_LOCAL_VAR";
        case AST_GLOBAL_VAR:  return "AST_GLOBAL_VAR";
        case AST_RETURN:      return "AST_RETURN";
        case AST_FUNC:        return "AST_FUNC";
        case AST_BREAK:       return "AST_BREAK";
        case AST_BUILTIN:     return "AST_BUILTIN";
        case AST_PROGRAM:     return "AST_PROGRAM";
        case AST_BLOCK:       return "AST_BLOCK";
    }
    
    puts("Unknown node type in node_type_to_string: ");
    putu(typ); putc('\n');
    exit(1);
}

fn is_binary_op(typ: int): int {
    static_assert(NUM_NODE_TYPES == 45, "Exhaustive match in is_binary_op");
    if (typ == AST_PLUS) return true;
    if (typ == AST_MINUS) return true;
    if (typ == AST_MUL) return true;
    if (typ == AST_DIV) return true;
    if (typ == AST_MOD) return true;
    if (typ == AST_LSHIFT) return true;
    if (typ == AST_RSHIFT) return true;
    if (typ == AST_AND) return true;
    if (typ == AST_BWAND) return true;
    if (typ == AST_OR) return true;
    if (typ == AST_BWOR) return true;
    if (typ == AST_XOR) return true;
    if (typ == AST_EQ) return true;
    if (typ == AST_NEQ) return true;
    if (typ == AST_LT) return true;
    if (typ == AST_LEQ) return true;
    if (typ == AST_GT) return true;
    if (typ == AST_GEQ) return true;
    return false;
}

fn is_unary_op(typ: int): int {
    static_assert(NUM_NODE_TYPES == 45, "Exhaustive match in is_unary_op");
    if (typ == AST_NEG) return true;
    if (typ == AST_NOT) return true;
    if (typ == AST_BWINV) return true;
    if (typ == AST_ADDROF) return true;
    if (typ == AST_DEREF) return true;
    return false;
}

fn is_lvalue(typ: int): int {
    static_assert(NUM_NODE_TYPES == 45, "Exhaustive match in is_lvalue");
    if (typ == AST_LOCAL_VAR) return true;
    if (typ == AST_GLOBAL_VAR) return true;
    if (typ == AST_MEMBER) return true;
    if (typ == AST_DEREF) return true;
    return false;
}

fn binary_token_to_op(token_typ: int): int {
    static_assert(NUM_TOKEN_TYPES == 75, "Exhaustive match in binary_token_to_op");
    if (token_typ == TOKEN_PLUS)      return AST_PLUS;
    if (token_typ == TOKEN_MINUS)     return AST_MINUS;
    if (token_typ == TOKEN_STAR)      return AST_MUL;
    if (token_typ == TOKEN_SLASH)     return AST_DIV;
    if (token_typ == TOKEN_PERCENT)   return AST_MOD;
    if (token_typ == TOKEN_LSHIFT)    return AST_LSHIFT;
    if (token_typ == TOKEN_RSHIFT)    return AST_RSHIFT;
    if (token_typ == TOKEN_AND)       return AST_AND;
    if (token_typ == TOKEN_OR)        return AST_OR;
    if (token_typ == TOKEN_XOR)       return AST_XOR;
    if (token_typ == TOKEN_EQ)        return AST_EQ;
    if (token_typ == TOKEN_NEQ)       return AST_NEQ;
    if (token_typ == TOKEN_LT)        return AST_LT;
    if (token_typ == TOKEN_LEQ)       return AST_LEQ;
    if (token_typ == TOKEN_GT)        return AST_GT;
    if (token_typ == TOKEN_GEQ)       return AST_GEQ;
    if (token_typ == TOKEN_AMPERSAND) return AST_BWAND;
    if (token_typ == TOKEN_BAR)       return AST_BWOR;
    if (token_typ == TOKEN_CARET)     return AST_XOR;

    puts("Unknown token in binary_token_to_op: ");
    putsln(token_type_to_string(token_typ));
    exit(1);
}

fn dump_ast(node: Node*, depth: int) {
    for (let i = 0; i < 2*depth; ++i)
        putc(' ');
    if (node.typ == AST_PROGRAM || node.typ == AST_BLOCK) {
        putsln(node_type_to_string(node.typ));
        for (let i = 0; i < node.block.children.size; ++i) {
            dump_ast(node.block.children.data[i], depth + 1);
        }
    } else if (is_binary_op(node.typ)) {
        putsln(node_type_to_string(node.typ));
        dump_ast(node.binary.lhs, depth + 1);
        dump_ast(node.binary.rhs, depth + 1);
    } else if (is_unary_op(node.typ) || node.typ == AST_RETURN) {
        putsln(node_type_to_string(node.typ));
        dump_ast(node.unary, depth + 1);

    } else if (node.typ == AST_CONDITIONAL || node.typ == AST_IF) {
        putsln(node_type_to_string(node.typ));
        dump_ast(node.conditional.cond, depth + 1);
        dump_ast(node.conditional.then, depth + 1);
        if (node.conditional.els)
            dump_ast(node.conditional.els, depth + 1);

    } else if (node.typ == AST_LITERAL) {
        if (node.etyp.typ == TYPE_INT) {
            putu(node.literal.as_int); putc('\n');
        } else if (node.etyp.typ == TYPE_PTR) {
            putc('"'); puts(node.literal.as_string); putc('"'); putc('\n');
        } else if (node.etyp.typ == TYPE_CHAR) {
            putc('\''); putc(node.literal.as_char); putc('\''); putc('\n');
        } else if (node.etyp.typ == TYPE_F64) {
            putsln(node.literal.as_string);
        } else {
            die(here, "Unknown literal type in dump_ast");
        }
    } else if (node.typ == AST_FUNC) {
        puts("func "); puts(node.func.name); puts("()\n");
        if (node.func.body)
            dump_ast(node.func.body, depth + 1);
    } else if (node.typ == AST_VARDECL) {
        puts("let "); puts(node.var_decl.var.name);
        if (node.var_decl.var.typ == TYPE_PTR) {
            puts(": ");
            puts(create_type_string(node.var_decl.var.typ));
        }
        if (node.var_decl.init) {
            puts(" =\n");
            dump_ast(node.var_decl.init, depth + 1);
        } else {
            putc('\n');
        }
    } else if (node.typ == AST_ASSIGN) {
        putsln(node_type_to_string(node.typ));
        dump_ast(node.assign.lhs, depth + 1);
        dump_ast(node.assign.rhs, depth + 1);
    } else if (node.typ == AST_LOCAL_VAR || node.typ == AST_GLOBAL_VAR) {
        putsln(node.variable.name);
    } else if (node.typ == AST_CONVERT) {
        puts("("); puts(create_type_string(node.etyp)); putsln(")");
        dump_ast(node.unary, depth + 1);
    } else if (node.typ == AST_CASE) {
        puts("case "); putu(node.case_stmt.val); putsln(":");
        for (let i = 0; i < node.case_stmt.children.size; ++i) {
            dump_ast(node.case_stmt.children.data[i], depth + 1);
        }
    } else if (node.typ == AST_SWITCH) {
        putsln("SWITCH");
        dump_ast(node.switch_stmt.expr, depth + 1);
        for (let i = 0; i < node.switch_stmt.cases.size; ++i) {
            dump_ast(node.switch_stmt.cases.data[i], depth + 2);
        }
        if (node.switch_stmt.defolt) {
            for (let i = 0; i < 2*(depth+2); ++i)
                putc(' ');
            putsln("default:");
            for (let i = 0; i < node.switch_stmt.defolt.size; ++i) {
                dump_ast(node.switch_stmt.defolt.data[i], depth + 3);
            }
        }
    } else {
        putsln(node_type_to_string(node.typ));
    }
}

// Defined below.
fn type_check_unary(node: Node*, token: Token*): Node*;

fn decay_array_to_pointer(node: Node*, token: Token*): Node* {
    // We can only take the address of an lvalue, so we need to ensure that
    if (is_lvalue(node.typ) && node.etyp.typ == TYPE_ARRAY) {
        let address = node_new(AST_ADDROF);
        address.unary = node;
        address = type_check_unary(address, token);
        node = address;
    }
    return node;
}

fn type_check_unary(node: Node*, token: Token*): Node* {
    static_assert(NUM_NODE_TYPES == 45, "Exhaustive match in type_check_unary");

    let old_type = node.unary.etyp;

    if (node.typ != AST_ADDROF && old_type.typ == TYPE_STRUCT)
        die_loc(here, &token.loc, "Performing invalid unary operation on struct type");

    if (is_float_type(old_type)) {
        if (node.typ != AST_ADDROF && node.typ != AST_NEG)
            die_loc2(here, &token.loc, "Performing invalid unary operation on float type: ", node_type_to_string(node.typ));
    }

    if (node.typ == AST_NOT) {
        node.etyp = type_new(TYPE_INT);
    } else if (node.typ == AST_ADDROF) {
        let ptr = type_new(TYPE_PTR);
        // The address of an array is a pointer to the first element
        ptr.ptr = old_type.typ == TYPE_ARRAY ? old_type.ptr : old_type;
        node.etyp = ptr;
    } else if (node.typ == AST_DEREF) {
        if (old_type.typ != TYPE_PTR)
            die_loc(here, &token.loc, "Cannot dereference non-pointer type");
        node.etyp = old_type.ptr;
        // If the dereferenced type is an array, we need to decay it to a
        // pointer to the first element.
        node = decay_array_to_pointer(node, token);
    } else if (node.typ == AST_NEG) {
        if (!is_int_type(old_type) && !is_float_type(old_type))
            die_loc(here, &token.loc, "Cannot negate non-integer type");
        node.etyp = old_type;
    } else {
        // Default to not changing the type
        node.etyp = old_type;
    }
    return node;
}

fn type_check_binary(node: Node*, token: Token*): Node* {
    static_assert(NUM_NODE_TYPES == 45, "Exhaustive match in type_check_binary");

    let lhs = node.binary.lhs.etyp;
    let rhs = node.binary.rhs.etyp;

    if (lhs.typ == TYPE_STRUCT || rhs.typ == TYPE_STRUCT)
        die_loc(here, &token.loc, "Performing invalid binary operation on struct type");

    if (node.typ == AST_PLUS) {
        if (is_int_type(lhs) && is_int_type(rhs)) {
            node.etyp = type_new(TYPE_INT);
        } else if (lhs.typ == TYPE_PTR) {
            if (!is_int_type(rhs))
                die_loc(here, &token.loc, "Cannot add non-integer to pointer");

            node.etyp = lhs;
            // Pointer arithmetic!
            if (size_for_type(lhs.ptr) != 1) {
                let mul = node_new_binop(AST_MUL, node.binary.rhs, node_from_int_literal(size_for_type(lhs.ptr)));
                node.binary.rhs = mul;
            }
        } else if (rhs.typ == TYPE_PTR) {
            if (!is_int_type(lhs))
                die_loc(here, &token.loc, "Cannot add non-integer to pointer");

            node.etyp = rhs;
            // Pointer arithmetic!
            if (size_for_type(rhs.ptr) != 1) {
                let mul = node_new_binop(AST_MUL, node.binary.lhs, node_from_int_literal(size_for_type(rhs.ptr)));
                node.binary.lhs = mul;
            }
        } else if (is_float_type(lhs) || is_float_type(rhs)) {
            // TODO: Handle different sized floats
            node.etyp = type_new(TYPE_F64);
            node.binary.lhs = convert_to_float(node.binary.lhs);
            node.binary.rhs = convert_to_float(node.binary.rhs);
        } else {
            puts("lhs: "); puts(create_type_string(lhs)); putc('\n');
            puts("rhs: "); puts(create_type_string(rhs)); putc('\n');
            die_loc(here, &token.loc, "Incompatible types for addition");
        }
    } else if (node.typ == AST_MINUS) {
        if (is_int_type(lhs) && is_int_type(rhs)) {
            node.etyp = type_new(TYPE_INT);
        } else if (lhs.typ == TYPE_PTR && is_int_type(rhs)) {
            node.etyp = lhs;
            // Pointer arithmetic!
            if (size_for_type(lhs.ptr) != 1) {
                let mul = node_new_binop(AST_MUL, node.binary.rhs, node_from_int_literal(size_for_type(lhs.ptr)));
                node.binary.rhs = mul;
            }
        } else if (is_int_type(lhs) && rhs.typ == TYPE_PTR) {
            node.etyp = rhs;
            // Pointer arithmetic!
            if (size_for_type(rhs.ptr) != 1) {
                let mul = node_new_binop(AST_MUL, node.binary.lhs, node_from_int_literal(size_for_type(rhs.ptr)));
                node.binary.lhs = mul;
            }
        } else if (lhs.typ == TYPE_PTR && rhs.typ == TYPE_PTR) {
            node.etyp = type_new(TYPE_INT);
            if (!types_equal(lhs.ptr, rhs.ptr))
                die_loc(here, &token.loc, "Cannot subtract pointers of different types");

            // Pointer arithmetic! (Divide by size of element)
            if (size_for_type(lhs.ptr) != 1) {
                let mul = node_new_binop(AST_MUL, node.binary.lhs, node_from_int_literal(size_for_type(lhs.ptr)));
                node.binary.lhs = mul;
            }
        } else if (is_float_type(lhs) || is_float_type(rhs)) {
            // TODO: Handle different sized floats
            node.etyp = type_new(TYPE_F64);
            node.binary.lhs = convert_to_float(node.binary.lhs);
            node.binary.rhs = convert_to_float(node.binary.rhs);
        } else {
            die_loc(here, &token.loc, "Cannot subtract non-integer types");
        }
    } else if (node.typ == AST_MUL || node.typ == AST_DIV || node.typ == AST_MOD) {
        if (is_int_type(lhs) && is_int_type(rhs)) {
            node.etyp = lhs;
        } else if (is_float_type(lhs) || is_float_type(rhs)) {
            if (node.typ == AST_MOD)
                die_loc(here, &token.loc, "Cannot modulo floats");

            // TODO: Handle different sized floats
            node.etyp = type_new(TYPE_F64);
            node.binary.lhs = convert_to_float(node.binary.lhs);
            node.binary.rhs = convert_to_float(node.binary.rhs);
        } else {
            die_loc2(here, &token.loc, "Cannot do operation non-integer types:", node_type_to_string(node.typ));
        }
    } else {
        // FIXME: This isn't very correct, but it's probably good enough for now
        node.etyp = type_new(TYPE_INT);
        if (is_float_type(lhs) || is_float_type(rhs)) {
            node.binary.lhs = convert_to_float(node.binary.lhs);
            node.binary.rhs = convert_to_float(node.binary.rhs);
        }
    }
    return node;
}

fn dump_rectype(typ: Type*, depth: int) {
    while (typ.typ == TYPE_PTR || typ.typ == TYPE_ARRAY) {
             if (typ.typ == TYPE_PTR) puts("*");
        else if (typ.typ == TYPE_ARRAY) puts("[]");
        typ = typ.ptr;
    }
         if (typ.typ == TYPE_INT) { putsln("int"); return; }
    else if (typ.typ == TYPE_CHAR) { putsln("char"); return; }
    else if (typ.typ == TYPE_VOID) { putsln("void"); return; }
    else if (typ.typ == TYPE_ANY) { putsln("any"); return; }

    puts(typ.typ == TYPE_STRUCT ? "struct " : "union ");
    puts(typ.struct_name);
    putsln(" {");
    for (let i = 0; i < typ.fields.size; ++i) {
        let field: Variable* = typ.fields.data[i];
        for (let j = 0; j < depth; ++j) puts("  ");
        puts("- "); puts(field.name); puts(" ("); putu(field.offset); puts("): ");
        dump_rectype(field.typ, depth + 2);
        puts("\n");
    }
    for (let i = 0; i < depth; ++i) puts("  ");
    putsln("}");
}

fn convert_type(to: Type*, from_node: Node*): Node* {
    let from = from_node.etyp;

    if (is_float_type(from) && is_int_type(to))
        return convert_to_int(from_node);

    if (is_float_type(to) && is_int_type(from))
        return convert_to_float(from_node);

    if (from.typ == TYPE_ANY || to.typ == TYPE_ANY)
        return from_node;

    // Allow converstions to and from void* to any other pointer type
    if (from.typ == TYPE_PTR && to.typ == TYPE_PTR)
        if (from.ptr.typ == TYPE_VOID || to.ptr.typ == TYPE_VOID)
            return from_node;

    // TODO: Should we raise a warning if the target type is narrower
    //       than the source type?
    if (is_int_type(from) && is_int_type(to))
        return from_node;

    if (types_equal(from, to))
        return from_node;

    return null;
}

// FIXME: These should be in `types.cup` ideally, but `Variable` is not defined
//        there and we can't forward-declare types.
fn compound_push_field(compound: Type*, name: char*, typ: Type*, base_offset: int): int {
    if (compound.typ != TYPE_STRUCT && compound.typ != TYPE_UNION)
        die(here, "compound_push_field: not a compound type");

    let is_union = compound.typ == TYPE_UNION;

    let field_size = size_for_type(typ);
    let offset_factor = min(field_size, 8);
    let offset = is_union ? 0 : align_up(compound.size, offset_factor);
    compound.size = is_union ? max(field_size, compound.size) : offset + field_size;

    compound.fields::push(variable_new(name, typ, base_offset + offset));
    return offset;
}

fn compound_find_field(typ: Type*, name: char*): Variable* {
    for (let i = 0; i < typ.fields.size; ++i) {
        let field: Variable* = typ.fields.data[i];
        if (streq(field.name, name)) {
            return field;
        }

        // If this is an anonymous field, we look inside it:
        if (streq(field.name, "<anon>")) {
            if (field = compound_find_field(field.typ, name))
                return field;
        }
    }
    return null;
}

fn compound_find_method(typ: Type*, name: char*): Node* {
    for (let i = 0; i < typ.methods.size; ++i) {
        let _method: Node* = typ.methods.data[i];
        if (streq(_method.func.name, name)) {
            return _method;
        }
    }
    return null;
}