code gen can use registers besides RAX now

6 files changed, 311 insertions, 140 deletions

diff --git a/README.md b/README.md
index 2cfb498..3ded28b 100644
--- a/README.md
+++ b/README.md
@@ -3,14 +3,14 @@
 fuck it, we ball
 
 TODO soon:
-- [ ] let cgen use more than one register
-  - [ ] basic used/unused register tracking so allocate_temporary can allocate
+- [x] let cgen use more than one register
+  - [x] basic used/unused register tracking so allocate_temporary can allocate
         a register
-  - [ ] proper deallocation of temporaries when they are no longer used
-  - [ ] spilling of values occupying calling convention registers when a
+  - [x] proper deallocation of temporaries when they are no longer used
+  - [x] spilling of values occupying calling convention registers when a
         function is called
-  - [ ] allowing variable values to be kept in registers temporarily and then
-        spilled to their permanent memory locations
+  - [ ] ~~allowing variable values to be kept in registers temporarily and then
+        spilled to their permanent memory locations~~ PUNTED
 - [ ] use evaluated types to make code gen around math better
 - [ ] support for functions/function pointers as types
-  - [ ] move the "expression not callable" stuff to the type checker
\ No newline at end of file
+  - [ ] move the "expression not callable" stuff to the type checker
diff --git a/ccc.h b/ccc.h
index 3542e60..37bcbf1 100644
--- a/ccc.h
+++ b/ccc.h
@@ -5,6 +5,7 @@
 #define ARRAY_SZ(x) (sizeof(x) / sizeof(x[0]))
 
 typedef unsigned long long integral_t;
+typedef long long sintegral_t;
 typedef double floating_t;
 
 void* ccc_alloc(integral_t sz);
diff --git a/codegen.c b/codegen.c
index 5cba20f..3d96d33 100644
--- a/codegen.c
+++ b/codegen.c
@@ -19,18 +19,13 @@ struct lval_def {
     struct storage_location loc;
 };
 
-static const struct storage_location RV_LOC = {
-    .type = STO_REG,
-    .reg = &RAX,
-};
-static const struct storage_location MULDIV_LOC = RV_LOC;
-static const struct storage_location MULDIV_OVERFLOW_LOC = {
-    .type = STO_REG,
-    .reg = &RDX,
-};
+static struct reg* RV_REG = &RAX;
+static struct reg* MULDIV_REG = &RAX;
+static struct reg* MULDIV_OVERFLOW_REG = &RDX;
 
-#define RETURN_LABEL_FMT "%s@coda"
+#define RETURN_LABEL_FMT "%s@ret"
 #define FULL_REG_SZ 8
+#define WORD_SZ 2
 
 static struct scope* scope;
 static const struct fn_decl_node* active_fn;
@@ -57,14 +52,30 @@ static void exit_scope(struct scope* child_scope, bool save_bp_offset) {
     if (save_bp_offset) scope->bp_offset = child_scope->bp_offset;
 }
 
-static struct lval_def allocate_register(const struct type* type) {
-    return (struct lval_def) {
-        .loc = {
-            .type = STO_REG,
-            .reg = &RAX, /* TODO: no real register coloring happening LOL */
-        },
-        .type = type,
-    };
+static struct reg* allocate_register() {
+    for (integral_t i = 0; DATA_REGS[i] != NULL; i++) {
+        if (!DATA_REGS[i]->is_occupied) {
+            DATA_REGS[i]->is_occupied = true;
+            return DATA_REGS[i];
+        }
+    }
+    return NULL;
+}
+
+static void release_register(struct reg* reg) {
+    reg->is_occupied = false;
+}
+
+static void spill_register(FILE* outfile, struct reg* reg) {
+    fprintf(outfile, "\tpush %s\n", reg->qword);
+    reg->is_occupied = false;
+    scope->bp_offset += FULL_REG_SZ;
+}
+
+static void unspill_register(FILE* outfile, struct reg* reg) {
+    fprintf(outfile, "\tpop %s\n", reg->qword);
+    reg->is_occupied = true;
+    scope->bp_offset -= FULL_REG_SZ;
 }
 
 static const struct data_type* get_effective_data_type(
@@ -99,16 +110,29 @@ static struct lval_def allocate_temporary(
     FILE* outfile,
     const struct type* type
 ) {
-    return allocate_stack(outfile, type);
+    struct reg* reg = allocate_register();
+    if (reg == NULL) return allocate_stack(outfile, type);
+    return (struct lval_def) {
+        .type = type,
+        .loc = {
+            .type = STO_REG,
+            .reg = reg,
+        },
+    };
 }
 
-static void deallocate_temporary(FILE* outfile, const struct lval_def* tmp) {
-    if (tmp->loc.type == STO_STACK) {
-        integral_t type_sz = get_effective_data_type(tmp->type)->sz;
-        fprintf(outfile, "\tadd rsp, %llu\n", type_sz);
-        scope->bp_offset -= type_sz;
-    } else if (tmp->loc.type == STO_REG) {
-        /* TOOD: release the register back to the algo */
+static void release_temporary(FILE* outfile, const struct lval_def* tmp) {
+    switch (tmp->loc.type) {
+        case STO_REG:
+            release_register(tmp->loc.reg);
+            break;
+        case STO_STACK:
+        case STO_IMM:
+        case STO_FN:
+        case STO_LABEL:
+            break;
+        case STO_UNRESOLVED:
+            CGEN_PANIC("can't release unresolved storage");
     }
 }
 
@@ -201,9 +225,36 @@ static void emit_mov(
         case STO_STACK:
             if (src->type == STO_STACK) {
                 /* `mov mem, mem` is illegal in x86_64 */
-                struct lval_def tmp = allocate_register(dst->type);
+                if (dst_sz == FULL_REG_SZ || dst_sz == WORD_SZ) {
+                    /* if we can swing it, use the stack as the intermediary */
+                    fprintf(outfile, "\tpush ");
+                    emit_size_const(outfile, dst_sz);
+                    emit_storage_loc(outfile, &dst->loc, dst_sz);
+                    fprintf(outfile, "\n\tpop ");
+                    emit_size_const(outfile, dst_sz);
+                    emit_storage_loc(outfile, src, dst_sz);
+                    break;
+                }
+
+                struct reg* tmp_reg = allocate_register();
+                bool spill_reg = tmp_reg == NULL;
+                if (spill_reg) {
+                    spill_register(outfile, &RCX);
+                    tmp_reg = &RCX;
+                }
+
+                struct lval_def tmp = {
+                    .type = dst->type,
+                    .loc = {
+                        .type = STO_REG,
+                        .reg = tmp_reg,
+                    }
+                };
                 emit_mov(outfile, &tmp, src);
                 emit_mov(outfile, dst, &tmp.loc);
+
+                if (spill_reg) unspill_register(outfile, &RCX);
+                else release_register(tmp_reg);
                 return;
             }
 
@@ -319,7 +370,7 @@ static void emit_decl(
         allocate_stack(outfile, node->def_ref->type);
     node->def_ref->loc = var_dst.loc;
 
-    fprintf(outfile, "\t; '%s' lives in: ", node->def_ref->name);
+    fprintf(outfile, "\t; %s", node->def_ref->name);
     integral_t dst_sz = get_effective_data_type(var_dst.type)->sz;
     emit_storage_loc(outfile, &var_dst.loc, dst_sz);
     fprintf(outfile, "\n");
@@ -359,67 +410,71 @@ static void emit_assignment(
     if (dst != NULL) emit_mov(outfile, dst, &lval_def.loc);
 }
 
+static integral_t push_stack_args(FILE* outfile, struct expr_list_node* arg) {
+    if (arg == NULL) return 0;
+
+    integral_t args_sz =
+        get_effective_data_type(arg->resolved_type)->sz
+        + push_stack_args(outfile, arg->next);
+
+    struct lval_def arg_dst = allocate_stack(outfile, arg->resolved_type);
+    emit_expr(outfile, arg->expr, &arg_dst);
+
+    return args_sz;
+}
+
 static void emit_call(
     FILE* outfile,
     const struct call_node* node,
     const struct lval_def* dst
 ) {
-    integral_t orig_bp_offset = scope->bp_offset;
-    integral_t arg_bp_offset = orig_bp_offset;
- 
-    struct arg_decl_node* arg_decl = node->fn_ref->args;
-    struct expr_list_node* arg_eval = node->args;
-    while (arg_decl != NULL && arg_eval != NULL) {
-        struct lval_def arg_dst =
-            allocate_stack(outfile, arg_decl->def_ref->type);
-        emit_expr(outfile, arg_eval->expr, &arg_dst);
-
-        arg_decl = arg_decl->next;
-        arg_eval = arg_eval->next;
+    /* 1. spill all existing data registers besides dst */
+    const struct reg* dst_reg = dst->loc.type == STO_REG ? dst->loc.reg : NULL;
+    integral_t reg_occupied = 0, n_data_regs = 0;
+    for (; DATA_REGS[n_data_regs] != NULL; n_data_regs++) {
+        if (!DATA_REGS[n_data_regs]->is_occupied) continue;
+        if (DATA_REGS[n_data_regs] == dst_reg) continue;
+
+        reg_occupied |= 1 << n_data_regs;
+        spill_register(outfile, DATA_REGS[n_data_regs]);
     }
 
-    unsigned char arg_regnum = 0;
-    arg_decl = node->fn_ref->args;
-    while (arg_decl != NULL) {
-        const struct type* arg_type = arg_decl->def_ref->type;
-        arg_bp_offset += get_effective_data_type(arg_type)->sz;
-
-        struct lval_def arg_dst;
-        /* TODO: if the convention register is used,
-         * will need to spill the current value */
-        if (arg_regnum < CC_N_REGS)
-            arg_dst = (struct lval_def) {
-                .type = arg_type,
-                .loc = (struct storage_location) {
-                    .type = STO_REG,
-                    .reg = CALLING_CONV[arg_regnum++],
-                },
-            };
-        else
-            arg_dst = allocate_stack(outfile, arg_type);
-
-        emit_mov(outfile, &arg_dst, &(struct storage_location) {
-            .type = STO_STACK,
-            .bp_offset = arg_bp_offset,
+    /* 2. evaluate arguments in reverse order into their respective locations */
+    struct expr_list_node* arg = node->args;    
+    for (integral_t i = 0;
+            arg != NULL && CALLING_CONV[i] != NULL;
+            arg = arg->next, i++) {
+        CALLING_CONV[i]->is_occupied = true;
+        emit_expr(outfile, arg->expr, &(struct lval_def) {
+            .type = arg->resolved_type,
+            .loc = {
+                .type = STO_REG,
+                .reg = CALLING_CONV[i],
+            }
         });
-        arg_decl = arg_decl->next;
     }
+    integral_t arg_stack_space = push_stack_args(outfile, arg);
 
+    /* 3. `call <label>` */
     fprintf(outfile, "\tcall %s\n", node->fn_ref->name);
-    if (dst != NULL) {
-        if (get_effective_data_type(
-                &node->fn_ref->return_type) == &void_type)
-            CGEN_PANIC("can't assign the result of a void function");
 
-        emit_mov(outfile, dst, &RV_LOC);
+    /* 4. `mov dst, rax` */
+    emit_mov(outfile, dst, &(struct storage_location) {
+        .type = STO_REG,
+        .reg = &RAX,
+    });
+
+    /* 5. destroy stack-based arg holders */
+    if (arg_stack_space > 0)
+        fprintf(outfile, "\tadd rsp, %llu\n", arg_stack_space);
+
+    /* 6. unspill all data registers in reverse order */
+    for (integral_t i = n_data_regs; i > 0; i--) {
+        integral_t reg_idx = i - 1;
+        DATA_REGS[reg_idx]->is_occupied = false;
+        if (!(reg_occupied & (1 << reg_idx))) continue;
+        unspill_register(outfile, DATA_REGS[reg_idx]);
     }
-
-    /* mass-pop our argument temporaries off the stack */
-    scope->bp_offset = orig_bp_offset;
-    if (orig_bp_offset > 0)
-        fprintf(outfile, "\tlea rsp, [rbp - %llu]\n", orig_bp_offset);
-    else
-        fprintf(outfile, "\tmov rsp, rbp\n");
 }
 
 static void emit_unary(
@@ -457,14 +512,32 @@ static void emit_binary(
     }
 
     /* LHS goes in RAX explicitly because imul and idiv are weird */
-    struct lval_def rhs_dst = allocate_temporary(outfile, dst->type);
-    emit_expr(outfile, node->rhs, &rhs_dst);
-    struct lval_def lhs_dst = (struct lval_def) {
-        .loc = MULDIV_LOC,
+    bool clobber_dst = dst->loc.type == STO_REG && dst->loc.reg == MULDIV_REG;
+    bool spill_numerator = MULDIV_REG->is_occupied && !clobber_dst;
+    if (spill_numerator) spill_register(outfile, MULDIV_REG);
+    MULDIV_REG->is_occupied = true;
+    struct lval_def lhs_dst = {
         .type = dst->type,
+        .loc = {
+            .type = STO_REG,
+            .reg = MULDIV_REG,
+        }
     };
     emit_expr(outfile, node->lhs, &lhs_dst);
 
+    /* overflow is going to clobber RDX */
+    bool spill_overflow = MULDIV_OVERFLOW_REG->is_occupied;
+    if (spill_overflow) spill_register(outfile, MULDIV_OVERFLOW_REG);
+    MULDIV_OVERFLOW_REG->is_occupied = true;
+    struct storage_location overflow_loc = {
+        .type = STO_REG,
+        .reg = MULDIV_OVERFLOW_REG,
+    };
+
+    /* RHS can go wherever */
+    struct lval_def rhs_dst = allocate_temporary(outfile, dst->type);
+    emit_expr(outfile, node->rhs, &rhs_dst);
+
     integral_t dst_sz = get_effective_data_type(dst->type)->sz;
     switch (node->op) {
         case BINARY_ADD:
@@ -485,9 +558,9 @@ static void emit_binary(
         case BINARY_DIV:
             /* nothing in the top half reg */
             fprintf(outfile, "\txor ");
-            emit_storage_loc(outfile, &MULDIV_OVERFLOW_LOC, FULL_REG_SZ);
+            emit_storage_loc(outfile, &overflow_loc, FULL_REG_SZ);
             fprintf(outfile, ", ");
-            emit_storage_loc(outfile, &MULDIV_OVERFLOW_LOC, FULL_REG_SZ);
+            emit_storage_loc(outfile, &overflow_loc, FULL_REG_SZ);
             fprintf(outfile, "\n");
 
             fprintf(outfile, "\tidiv ");
@@ -499,11 +572,16 @@ static void emit_binary(
     emit_storage_loc(outfile, &rhs_dst.loc, dst_sz);
     fprintf(outfile, "\n");
 
-    /* TODO: deal with RDX overflow shit for imul and idiv */
+    /* we don't care about overflow */
+    if (spill_overflow) unspill_register(outfile, MULDIV_OVERFLOW_REG);
+    else release_register(MULDIV_OVERFLOW_REG);
+
+    /* answer's already in RAX */
+    if (clobber_dst) return;
     emit_mov(outfile, dst, &lhs_dst.loc);
 
-    deallocate_temporary(outfile, &lhs_dst);
-    deallocate_temporary(outfile, &rhs_dst);
+    if (spill_numerator) unspill_register(outfile, MULDIV_REG);
+    else release_register(MULDIV_REG);
 }
 
 static void emit_expr_list(
@@ -573,8 +651,11 @@ static void emit_return(FILE* outfile, const struct return_node* node) {
             outfile,
             node->ret_val,
             &(struct lval_def) {
-                .loc = RV_LOC,
                 .type = &active_fn->return_type,
+                .loc = {
+                    .type = STO_REG,
+                    .reg = RV_REG,
+                },
             });
 
     fprintf(outfile, "\tjmp " RETURN_LABEL_FMT "\n", active_fn->name);
@@ -605,6 +686,7 @@ static void emit_if(FILE* outfile, const struct if_node* node) {
         allocate_temporary(outfile, node->cond->resolved_type);
     emit_expr_list(outfile, node->cond, &cond_result);
     emit_cmp_zero(outfile, &cond_result);
+    release_temporary(outfile, &cond_result);
 
     integral_t branch_num = ++branch_counter;
     fprintf(outfile, "\tjz branch_false@%lld\n", branch_num);
@@ -645,6 +727,7 @@ static void emit_loop(FILE* outfile, const struct loop_node* node) {
         fprintf(outfile, "loop_head@%lld:\n", loop_num);
         emit_expr_list(outfile, node->cond, &cond_dst);
         emit_cmp_zero(outfile, &cond_dst);
+        release_temporary(outfile, &cond_dst);
         fprintf(outfile, "\tjz loop_done@%lld\n", loop_num);
     } else {
         fprintf(outfile, "loop_head@%lld:\n", loop_num);
@@ -685,8 +768,6 @@ static void emit_stmt(FILE* outfile, const struct stmt_node* node) {
 }
 
 static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
-    enter_scope(node->scope, 0);
-
     if (active_fn != NULL)
         CGEN_PANIC(
             "can't define function %s inside function %s",
@@ -699,30 +780,30 @@ static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
     fprintf(outfile, "\tpush rbp\n");
     fprintf(outfile, "\tmov rbp, rsp\n");
 
-    scope->bp_offset = 0;
+    enter_scope(node->scope, 0);
 
-    long long spilled_bp_ofs = -16; // return address + old bp
-    unsigned char arg_regnum = 0;
+    sintegral_t spilled_bp_ofs = -16; // return address + spilled rbp
     struct arg_decl_node* arg_decl = node->args;
-    for (; arg_decl != NULL; arg_decl = arg_decl->next) {
+    for (integral_t i = 0; arg_decl != NULL; arg_decl = arg_decl->next, i++) {
         struct var_def* arg_def = arg_decl->def_ref;
+        fprintf(outfile, "\t; %s\n", arg_def->name);
         struct lval_def arg_dst =
             allocate_stack(outfile, arg_def->type);
         arg_def->loc = arg_dst.loc;
 
         struct storage_location arg_src;
-        if (arg_regnum < CC_N_REGS) {
+        if (CALLING_CONV[i] != NULL) {
             arg_src = (struct storage_location) {
                 .type = STO_REG,
-                .reg = CALLING_CONV[arg_regnum++]
+                .reg = CALLING_CONV[i]
             };
+            CALLING_CONV[i]->is_occupied = false;
         } else {
             arg_src = (struct storage_location) {
                 .type = STO_STACK,
                 .bp_offset = spilled_bp_ofs,
             };
-            integral_t arg_sz = get_effective_data_type(arg_def->type)->sz;
-            spilled_bp_ofs -= arg_sz;
+            spilled_bp_ofs -= get_effective_data_type(arg_def->type)->sz;
         }
         emit_mov(outfile, &arg_dst, &arg_src);
     }
diff --git a/register.c b/register.c
index fe1a718..f197513 100644
--- a/register.c
+++ b/register.c
@@ -1,48 +1,123 @@
 #include "register.h"
+#include <stddef.h>
 
-const struct reg RAX = {
+struct reg RAX = {
     .qword = "rax",
     .dword = "eax",
     .word = "ax",
     .byte = "al",
+    .is_occupied = false,
 };
-const struct reg RDI = {
+struct reg RBX = {
+    .qword = "rbx",
+    .dword = "ebx",
+    .word = "bx",
+    .byte = "bl",
+    .is_occupied = false,
+};
+struct reg RCX = {
+    .qword = "rcx",
+    .dword = "ecx",
+    .word = "cx",
+    .byte = "cl",
+    .is_occupied = false,
+};
+struct reg RDX = {
+    .qword = "rdx",
+    .dword = "edx",
+    .word = "dx",
+    .byte = "dl",
+    .is_occupied = false,
+};
+
+struct reg RDI = {
     .qword = "rdi",
     .dword = "edi",
     .word = "di",
     .byte = "dil",
+    .is_occupied = false,
 };
-const struct reg RSI = {
+struct reg RSI = {
     .qword = "rsi",
     .dword = "esi",
     .word = "si",
     .byte = "sil",
+    .is_occupied = false,
 };
-const struct reg RDX = {
-    .qword = "rdx",
-    .dword = "edx",
-    .word = "dx",
-    .byte = "dl",
+struct reg RSP = {
+    .qword = "rsp",
+    .dword = "esp",
+    .word = "sp",
+    .byte = "spl",
+    .is_occupied = true,
 };
-const struct reg R10 = {
-    .qword = "r10",
-    .dword = "r10d",
-    .word = "r10w",
-    .byte = "r10b",
+struct reg RBP = {
+    .qword = "rbp",
+    .dword = "ebp",
+    .word = "bp",
+    .byte = "bpl",
+    .is_occupied = true,
+};
+
+struct reg R8 = {
+    .qword = "r8",
+    .dword = "r8d",
+    .word = "r8w",
+    .byte = "r8b",
+    .is_occupied = false,
 };
-const struct reg R9 = {
+struct reg R9 = {
     .qword = "r9",
     .dword = "r9d",
     .word = "r9w",
     .byte = "r9b",
+    .is_occupied = false,
 };
-const struct reg R8 = {
-    .qword = "r8",
-    .dword = "r8d",
-    .word = "r8w",
-    .byte = "r8b",
+struct reg R10 = {
+    .qword = "r10",
+    .dword = "r10d",
+    .word = "r10w",
+    .byte = "r10b",
+    .is_occupied = false,
+};
+struct reg R11 = {
+    .qword = "r11",
+    .dword = "r11d",
+    .word = "r11w",
+    .byte = "r11b",
+    .is_occupied = false,
+};
+struct reg R12 = {
+    .qword = "r12",
+    .dword = "r12d",
+    .word = "r12w",
+    .byte = "r12b",
+    .is_occupied = false,
+};
+struct reg R13 = {
+    .qword = "r13",
+    .dword = "r13d",
+    .word = "r13w",
+    .byte = "r13b",
+    .is_occupied = false,
+};
+struct reg R14 = {
+    .qword = "r14",
+    .dword = "r14d",
+    .word = "r14w",
+    .byte = "r14b",
+    .is_occupied = false,
+};
+struct reg R15 = {
+    .qword = "r15",
+    .dword = "r15d",
+    .word = "r15w",
+    .byte = "r15b",
+    .is_occupied = false,
 };
 
-const struct reg* const CALLING_CONV[] = {&RDI, &RSI, &RDX, &R10, &R9, &R8};
-const unsigned char CC_N_REGS =
-    sizeof(CALLING_CONV) / sizeof(const struct reg* const);
+struct reg* const CALLING_CONV[] = {&RDI, &RSI, &RDX, &R10, &R9, &R8, NULL};
+struct reg* const DATA_REGS[] = {
+    &RDI, &RSI, &RAX, &RBX, &RCX, &RDX,
+    &R8, &R9, &R10, &R11, &R12, &R13, &R14, &R15, NULL
+};
diff --git a/register.h b/register.h
index 943e839..8a8ae3c 100644
--- a/register.h
+++ b/register.h
@@ -2,21 +2,35 @@
 #define REGISTER_H
 
 struct reg {
-    const char* qword;
-    const char* dword;
-    const char* word;
-    const char* byte;
+    const char* const qword;
+    const char* const dword;
+    const char* const word;
+    const char* const byte;
+    bool is_occupied;
 };
 
-extern const struct reg RAX;
-extern const struct reg RDI;
-extern const struct reg RSI;
-extern const struct reg RDX;
-extern const struct reg R10;
-extern const struct reg R9;
-extern const struct reg R8;
+/* TODO: maybe don't need these in the header file */
+extern struct reg RAX;
+extern struct reg RBX;
+extern struct reg RCX;
+extern struct reg RDX;
 
-extern const struct reg* const CALLING_CONV[];
-extern const unsigned char CC_N_REGS;
+extern struct reg RDI;
+extern struct reg RSI;
+extern struct reg RSP;
+extern struct reg RBP;
+
+extern struct reg R8;
+extern struct reg R9;
+extern struct reg R10;
+extern struct reg R11;
+extern struct reg R12;
+extern struct reg R13;
+extern struct reg R14;
+extern struct reg R15;
+/* END TODO */
+
+extern struct reg* const CALLING_CONV[];
+extern struct reg* const DATA_REGS[];
 
 #endif
diff --git a/scope.h b/scope.h
index 8a3c80f..72b342c 100644
--- a/scope.h
+++ b/scope.h
@@ -16,9 +16,9 @@ struct storage_location {
         STO_UNRESOLVED,
     } type;
     union {
-        const struct reg* reg;
+        struct reg* reg;
         const char* label;
-        long long bp_offset;
+        sintegral_t bp_offset;
         integral_t value;
         struct fn_decl_node* decl;
     };