code gen can use registers besides RAX now

1 files changed, 177 insertions, 96 deletions

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);
     }