Enhance register allocation with usage-based spill

Building on the VReg abstraction, this commit adds intelligent spill
decisions based on variable usage patterns and context.
- Track variable lifetime with first_use, last_use, and use_count
- Add loop_depth field for identifying loop-resident variables
- Transform spill cost from simple distance to multi-factor heuristic:
  * Loop variables: +200 cost per nesting level
  * Usage frequency: +5 cost per use
  * Live range length: +20 for long ranges
  * Constants: -50 (prefer spilling for rematerialization)

Author: jserv

src/arm-codegen.c

@@ -485,16 +485,18 @@ void code_generate(void)
         emit_ph2_ir(ph2_ir);
 
     /* prepare 'argc' and 'argv', then proceed to 'main' function */
-    emit(__movw(__AL, __r8, GLOBAL_FUNC->stack_size));
-    emit(__movt(__AL, __r8, GLOBAL_FUNC->stack_size));
-    emit(__add_r(__AL, __r8, __r12, __r8));
-    emit(__lw(__AL, __r0, __r8, 0));
-    emit(__add_i(__AL, __r1, __r8, 4));
-    emit(__bl(__AL, MAIN_BB->elf_offset - elf_code->size));
+    if (MAIN_BB) {
+        emit(__movw(__AL, __r8, GLOBAL_FUNC->stack_size));
+        emit(__movt(__AL, __r8, GLOBAL_FUNC->stack_size));
+        emit(__add_r(__AL, __r8, __r12, __r8));
+        emit(__lw(__AL, __r0, __r8, 0));
+        emit(__add_i(__AL, __r1, __r8, 4));
+        emit(__bl(__AL, MAIN_BB->elf_offset - elf_code->size));
 
-    /* exit with main's return value - r0 already has the return value */
-    emit(__mov_i(__AL, __r7, 1));
-    emit(__svc());
+        /* exit with main's return value - r0 already has the return value */
+        emit(__mov_i(__AL, __r7, 1));
+        emit(__svc());
+    }
 
     for (int i = 0; i < ph2_ir_idx; i++) {
         ph2_ir = PH2_IR_FLATTEN[i];

src/defs.h

@@ -375,6 +375,10 @@ struct var {
     int vreg_id;    /* Virtual register ID */
     int phys_reg;   /* Physical register assignment (-1 if unassigned) */
     int vreg_flags; /* VReg flags */
+    int first_use;  /* First instruction index where variable is used */
+    int last_use;   /* Last instruction index where variable is used */
+    int loop_depth; /* Nesting depth if variable is in a loop */
+    int use_count;  /* Number of times variable is used */
 };
 
 typedef struct {

src/parser.c

@@ -64,6 +64,10 @@ var_t *require_var(block_t *blk)
     var_list->elements[var_list->size++] = var;
     var->consumed = -1;
     var->phys_reg = -1;
+    var->first_use = -1;
+    var->last_use = -1;
+    var->loop_depth = 0;
+    var->use_count = 0;
     var->base = var;
     var->type = TY_int;
     return var;

src/reg-alloc.c

@@ -54,6 +54,19 @@ bool check_live_out(basic_block_t *bb, var_t *var)
     return false;
 }
 
+void track_var_use(var_t *var, int insn_idx)
+{
+    if (!var)
+        return;
+
+    var->use_count++;
+
+    if (var->first_use < 0)
+        var->first_use = insn_idx;
+
+    var->last_use = insn_idx;
+}
+
 void refresh(basic_block_t *bb, insn_t *insn)
 {
     for (int i = 0; i < REG_CNT; i++) {
@@ -93,18 +106,57 @@ ph2_ir_t *bb_add_ph2_ir(basic_block_t *bb, opcode_t op)
     return n;
 }
 
+/* Calculate the cost of spilling a variable from a register.
+ * Higher cost means the variable is more valuable to keep in a register.
+ * The cost is computed based on multiple factors that affect performance.
+ */
 int calculate_spill_cost(var_t *var, basic_block_t *bb, int current_idx)
 {
+    int cost = 0;
+
+    /* Variables that are live-out of the basic block must be spilled anyway,
+     * so give them a high cost to prefer spilling them over others
+     */
     if (check_live_out(bb, var))
-        return 1000;
+        cost += 1000;
 
+    /* Variables that will be used soon should have higher cost.
+     * The closer the next use, the higher the penalty for spilling
+     */
     if (var->consumed > current_idx) {
         int distance = var->consumed - current_idx;
         if (distance < 10)
-            return 100 - distance * 10;
+            cost += 100 - distance * 10; /* Max 100 points for immediate use */
+    }
+
+    /* Frequently used variables should stay in registers.
+     * Each use adds 5 points to the cost
+     */
+    if (var->use_count > 0)
+        cost += var->use_count * 5;
+
+    /* Variables inside loops are accessed repeatedly, so they should have much
+     * higher priority to stay in registers (200 points per level)
+     */
+    if (var->loop_depth > 0)
+        cost += var->loop_depth * 200;
+
+    /* Constants can be easily reloaded, so prefer spilling them by reducing
+     * their cost
+     */
+    if (var->is_const)
+        cost -= 50;
+
+    /* Variables with long live ranges may benefit from spilling to free up
+     * registers for other variables
+     */
+    if (var->first_use >= 0 && var->last_use >= 0) {
+        int range_length = var->last_use - var->first_use;
+        if (range_length > 100)
+            cost += 20; /* Small penalty for very long live ranges */
     }
 
-    return 0;
+    return cost;
 }
 
 int find_best_spill(basic_block_t *bb,
@@ -269,9 +321,8 @@ int prepare_dest(basic_block_t *bb, var_t *var, int operand_0, int operand_1)
         }
     }
 
-    if (REGS[spilled].var) {
+    if (REGS[spilled].var)
         vreg_clear_phys(REGS[spilled].var);
-    }
 
     spill_var(bb, REGS[spilled].var, spilled);
     REGS[spilled].var = var;
@@ -512,6 +563,7 @@ void reg_alloc(void)
 
                 switch (insn->opcode) {
                 case OP_unwound_phi:
+                    track_var_use(insn->rs1, insn->idx);
                     src0 = prepare_operand(bb, insn->rs1, -1);
 
                     if (!insn->rd->offset) {
@@ -610,6 +662,7 @@ void reg_alloc(void)
                     if (insn->rd->consumed == -1)
                         break;
 
+                    track_var_use(insn->rs1, insn->idx);
                     src0 = find_in_regs(insn->rs1);
 
                     /* If operand is loaded from stack, clear the original slot
@@ -754,6 +807,8 @@ void reg_alloc(void)
                 case OP_bit_and:
                 case OP_bit_or:
                 case OP_bit_xor:
+                    track_var_use(insn->rs1, insn->idx);
+                    track_var_use(insn->rs2, insn->idx);
                     src0 = prepare_operand(bb, insn->rs1, -1);
                     src1 = prepare_operand(bb, insn->rs2, src0);
                     dest = prepare_dest(bb, insn->rd, src0, src1);

src/riscv-codegen.c

@@ -467,14 +467,16 @@ void code_generate(void)
 
     /* prepare 'argc' and 'argv', then proceed to 'main' function */
     /* use original sp saved in s0 to get argc/argv */
-    emit(__addi(__t0, __s0, 0));
-    emit(__lw(__a0, __t0, 0));
-    emit(__addi(__a1, __t0, 4));
-    emit(__jal(__ra, MAIN_BB->elf_offset - elf_code->size));
+    if (MAIN_BB) {
+        emit(__addi(__t0, __s0, 0));
+        emit(__lw(__a0, __t0, 0));
+        emit(__addi(__a1, __t0, 4));
+        emit(__jal(__ra, MAIN_BB->elf_offset - elf_code->size));
 
-    /* exit with main's return value in a0 */
-    emit(__addi(__a7, __zero, 93));
-    emit(__ecall());
+        /* exit with main's return value in a0 */
+        emit(__addi(__a7, __zero, 93));
+        emit(__ecall());
+    }
 
     for (int i = 0; i < ph2_ir_idx; i++) {
         ph2_ir = PH2_IR_FLATTEN[i];