codegen.c

#define _POSIX_C_SOURCE 200809L
#include "codegen.h"
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* Instruction buffer for two-pass assembly */
typedef struct Instruction {
  char text[128];
  struct Instruction *next;
} Instruction;

static Instruction *instr_head = NULL;
static Instruction *instr_tail = NULL;
static int64_t instr_count = 0;

/* Forward declarations */
static void gen_commands(CodeGen *gen, ASTNode *cmds);
static void gen_command(CodeGen *gen, ASTNode *cmd);
static void gen_expression(CodeGen *gen, ASTNode *expr);
static void gen_condition(CodeGen *gen, ASTNode *cond, int64_t *jump_false);
static void gen_value(CodeGen *gen, ASTNode *val);
static void gen_load_address(CodeGen *gen, ASTNode *id);
static void gen_multiply(CodeGen *gen);
static void gen_divide(CodeGen *gen);
static void gen_modulo(CodeGen *gen);
static void gen_procedure(CodeGen *gen, ASTNode *proc);
static void gen_proc_call(CodeGen *gen, ASTNode *call);

/* Procedure info for code generation */
typedef struct ProcInfo {
  char *name;
  int64_t addr;         /* Start address of procedure code */
  int64_t ret_addr_mem; /* Memory location for return address */
  int arg_count;
  int64_t *arg_addrs; /* Memory locations for argument addresses */
  struct ProcInfo *next;
} ProcInfo;

static ProcInfo *proc_list = NULL;

static int64_t alloc_temp(CodeGen *gen, int64_t count) {
  int64_t addr = gen->temp_top;
  gen->temp_top += count;
  return addr;
}

static void free_temp(CodeGen *gen, int64_t count) {
  if (gen->temp_top - count >= gen->temp_base) {
    gen->temp_top -= count;
  }
}

static ProcInfo *get_proc_info(const char *name) {
  for (ProcInfo *p = proc_list; p; p = p->next) {
    if (strcmp(p->name, name) == 0)
      return p;
  }
  return NULL;
}

static void add_instruction(const char *text) {
  Instruction *inst = (Instruction *)malloc(sizeof(Instruction));
  strncpy(inst->text, text, sizeof(inst->text) - 1);
  inst->text[sizeof(inst->text) - 1] = '\0';
  inst->next = NULL;

  if (instr_tail) {
    instr_tail->next = inst;
  } else {
    instr_head = inst;
  }
  instr_tail = inst;
  instr_count++;
}

static void free_instructions(void) {
  Instruction *curr = instr_head;
  while (curr) {
    Instruction *next = curr->next;
    free(curr);
    curr = next;
  }
  instr_head = instr_tail = NULL;
  instr_count = 0;
}

/* Patch jump target address while preserving instruction type */
static void patch_jump(int64_t instr_addr, int64_t target_addr) {
  Instruction *inst = instr_head;
  for (int64_t i = 0; i < instr_addr && inst; i++, inst = inst->next)
    ;
  if (inst) {
    /* Extract instruction type (e.g., "JPOS", "JZERO", "JUMP") */
    char instr_type[16] = {0};
    sscanf(inst->text, "%15s", instr_type);
    snprintf(inst->text, sizeof(inst->text), "%s %lld", instr_type,
             (long long)target_addr);
  }
}

CodeGen *codegen_create(FILE *output, SymbolTable *symtab) {
  CodeGen *gen = (CodeGen *)malloc(sizeof(CodeGen));
  gen->output = output;
  gen->symtab = symtab;
  gen->current_line = 0;
  gen->temp_base = symtab_get_next_address(symtab);
  gen->temp_top = gen->temp_base;
  return gen;
}

void codegen_destroy(CodeGen *gen) {
  /* Free procedure info */
  ProcInfo *p = proc_list;
  while (p) {
    ProcInfo *next = p->next;
    free(p->name);
    if (p->arg_addrs)
      free(p->arg_addrs);
    free(p);
    p = next;
  }
  proc_list = NULL;

  free_instructions();
  free(gen);
}

void emit(CodeGen *gen, const char *fmt, ...) {
  char buf[128];
  va_list args;
  va_start(args, fmt);
  vsnprintf(buf, sizeof(buf), fmt, args);
  va_end(args);

  add_instruction(buf);
  gen->current_line++;
}

int64_t codegen_current_addr(CodeGen *gen) { return gen->current_line; }

/* Generate constant into register a */
static void gen_const(CodeGen *gen, int64_t value) {
  if (value == 0) {
    emit(gen, "RST a");
    return;
  }

  /* Find highest bit */
  int bits[64];
  int bit_count = 0;
  int64_t v = value;
  while (v > 0) {
    bits[bit_count++] = v & 1;
    v >>= 1;
  }

  /* Generate from highest bit down */
  emit(gen, "RST a");
  for (int i = bit_count - 1; i >= 0; i--) {
    if (i < bit_count - 1) {
      emit(gen, "SHL a");
    }
    if (bits[i]) {
      emit(gen, "INC a");
    }
  }
}

/* Load variable value into register a */
static void gen_load_var(CodeGen *gen, Symbol *sym) {
  if (sym->is_reference) {
    /* Parameter passed by reference - load address from memory, then load value
     * from that address */
    emit(gen, "LOAD %lld",
         (long long)sym->address); /* ra = address of actual variable */
    emit(gen, "RLOAD a");          /* ra = value at actual variable */
  } else {
    emit(gen, "LOAD %lld", (long long)sym->address);
  }
}

/* Store register a into variable */
static void gen_store_var(CodeGen *gen, Symbol *sym) {
  if (sym->is_reference) {
    /* Parameter passed by reference */
    emit(gen, "SWP b");
    emit(gen, "LOAD %lld", (long long)sym->address);
    emit(gen, "SWP b");
    emit(gen, "RSTORE b");
  } else {
    emit(gen, "STORE %lld", (long long)sym->address);
  }
}

static const char *current_scope = NULL;

void codegen_program(CodeGen *gen, ASTNode *root) {
  if (root == NULL)
    return;

  /* First pass: collect procedure info */
  int64_t temp_addr = gen->temp_base;

  /* Reserve space for temporaries used in arithmetic */
  /* We need about 10 temp locations */
  temp_addr += 10;

  /* Register all procedures */
  for (ASTNode *proc = root->left; proc; proc = proc->next) {
    if (proc->type == NODE_PROCEDURE) {
      ProcInfo *info = (ProcInfo *)calloc(1, sizeof(ProcInfo));
      info->name = strdup(proc->name);
      info->ret_addr_mem = temp_addr++;

      /* Count and allocate argument address storage */
      info->arg_count = 0;
      for (ASTNode *arg = proc->left; arg; arg = arg->next) {
        info->arg_count++;
      }
      if (info->arg_count > 0) {
        info->arg_addrs = (int64_t *)malloc(info->arg_count * sizeof(int64_t));
        for (int i = 0; i < info->arg_count; i++) {
          info->arg_addrs[i] = temp_addr++;
        }
      }

      info->next = proc_list;
      proc_list = info;
    }
  }

  /* Set temp stack base after reserving procedure storage */
  gen->temp_top = temp_addr;

  /* Emit jump to main */
  emit(gen, "JUMP 0"); /* Placeholder, will patch */
  int64_t jump_to_main = 0;

  /* Generate procedure code */
  for (ASTNode *proc = root->left; proc; proc = proc->next) {
    if (proc->type == NODE_PROCEDURE) {
      gen_procedure(gen, proc);
    }
  }

  /* Patch jump to main */
  int64_t main_addr = gen->current_line;

  /* Generate main code */
  current_scope = NULL;
  symtab_set_scope(gen->symtab, NULL);
  if (root->right) {
    gen_commands(gen, root->right->right);
  }

  emit(gen, "HALT");

  /* Write all instructions to output, patching the first jump */
  Instruction *inst = instr_head;
  int64_t addr = 0;
  while (inst) {
    if (addr == jump_to_main) {
      fprintf(gen->output, "JUMP %lld\n", (long long)main_addr);
    } else {
      fprintf(gen->output, "%s\n", inst->text);
    }
    inst = inst->next;
    addr++;
  }
}

static void gen_procedure(CodeGen *gen, ASTNode *proc) {
  ProcInfo *info = get_proc_info(proc->name);
  if (!info)
    return;

  info->addr = gen->current_line;
  current_scope = proc->name;
  symtab_set_scope(gen->symtab, proc->name);

  /* Update parameter symbols to use arg_addrs locations */
  Symbol *proc_sym = symtab_lookup_procedure(gen->symtab, proc->name);
  if (proc_sym && proc_sym->args) {
    for (int i = 0; i < info->arg_count && i < proc_sym->arg_count; i++) {
      Symbol *param = proc_sym->args[i];
      if (param) {
        param->address = info->arg_addrs[i];
        param->is_reference = true;
      }
    }
  }

  /* Store return address */
  emit(gen, "STORE %lld", (long long)info->ret_addr_mem);

  /* Generate body */
  gen_commands(gen, proc->right);

  /* Return */
  emit(gen, "LOAD %lld", (long long)info->ret_addr_mem);
  emit(gen, "RTRN");

  current_scope = NULL;
  symtab_set_scope(gen->symtab, NULL);
}

static void gen_proc_call(CodeGen *gen, ASTNode *call) {
  ProcInfo *info = get_proc_info(call->name);
  if (!info)
    return;

  Symbol *proc_sym = symtab_lookup_procedure(gen->symtab, call->name);
  if (!proc_sym)
    return;

  /* Pass arguments - store addresses of actual parameters */
  ASTNode *arg = call->left;
  for (int i = 0; i < info->arg_count && arg; i++, arg = arg->next) {
    Symbol *actual = symtab_lookup(gen->symtab, arg->name);
    if (!actual)
      continue;

    Symbol *formal = proc_sym->args[i];

    if (actual->is_reference) {
      /* Already a reference - load the address */
      emit(gen, "LOAD %lld", (long long)actual->address);
    } else if (actual->type == SYM_ARRAY ||
               (formal && formal->param_mode == PARAM_T)) {
      /* Array - compute base address */
      gen_const(gen, actual->address - actual->array_start);
    } else {
      /* Scalar - store its address */
      gen_const(gen, actual->address);
    }
    emit(gen, "STORE %lld", (long long)info->arg_addrs[i]);
  }

  /* Update parameter symbols in callee scope to point to arg storage */
  symtab_set_scope(gen->symtab, call->name);

  /* Find the procedure node to get formal parameter names */
  Symbol *callee = symtab_lookup_procedure(gen->symtab, call->name);
  for (int i = 0; i < info->arg_count && callee && callee->args; i++) {
    Symbol *param = callee->args[i];
    if (param) {
      param->address = info->arg_addrs[i];
      param->is_reference = true;
    }
  }
  symtab_set_scope(gen->symtab, current_scope);

  /* Call procedure */
  emit(gen, "CALL %lld", (long long)info->addr);
}

static void gen_commands(CodeGen *gen, ASTNode *cmds) {
  for (ASTNode *cmd = cmds; cmd; cmd = cmd->next) {
    gen_command(gen, cmd);
  }
}

static void gen_command(CodeGen *gen, ASTNode *cmd) {
  if (cmd == NULL)
    return;

  switch (cmd->type) {
  case NODE_ASSIGN: {
    /* Generate expression value into ra */
    gen_expression(gen, cmd->right);

    /* Store to target */
    ASTNode *target = cmd->left;
    Symbol *sym = symtab_lookup(gen->symtab, target->name);
    if (!sym)
      return;

    if (target->type == NODE_IDENT) {
      gen_store_var(gen, sym);
    } else {
      /* Array access - need to compute address */
      /* Use rc to preserve value since gen_load_address uses rb */
      emit(gen, "SWP c"); /* Save value in rc */
      gen_load_address(gen, target);
      emit(gen, "SWP b"); /* Address in rb, whatever was in rb goes to ra (don't
                             care) */
      emit(gen, "SWP c"); /* Value back to ra from rc */
      emit(gen, "RSTORE b");
    }
    break;
  }

  case NODE_IF: {
    int64_t jump_false;
    gen_condition(gen, cmd->left, &jump_false);
    gen_commands(gen, cmd->right);

    /* Patch condition jump */
    patch_jump(jump_false, gen->current_line);
    break;
  }

  case NODE_IF_ELSE: {
    int64_t jump_false;
    gen_condition(gen, cmd->left, &jump_false);
    gen_commands(gen, cmd->right);

    /* Jump over else */
    emit(gen, "JUMP 0"); /* Placeholder */
    int64_t jump_end = gen->current_line - 1;

    /* Patch condition jump */
    patch_jump(jump_false, gen->current_line);

    gen_commands(gen, cmd->extra);

    /* Patch end jump */
    patch_jump(jump_end, gen->current_line);
    break;
  }

  case NODE_WHILE: {
    int64_t loop_start = gen->current_line;
    int64_t jump_false;
    gen_condition(gen, cmd->left, &jump_false);
    gen_commands(gen, cmd->right);
    emit(gen, "JUMP %lld", (long long)loop_start);

    /* Patch condition jump */
    patch_jump(jump_false, gen->current_line);
    break;
  }

  case NODE_REPEAT: {
    int64_t loop_start = gen->current_line;
    gen_commands(gen, cmd->right);

    /* Condition - jump back if false (condition NOT met) */
    int64_t jump_false;
    gen_condition(gen, cmd->left, &jump_false);

    /* Patch - REPEAT continues if condition is false */
    patch_jump(jump_false, loop_start);
    break;
  }

  case NODE_FOR_TO:
  case NODE_FOR_DOWNTO: {
    Symbol *iter = cmd->sym ? cmd->sym : symtab_lookup(gen->symtab, cmd->name);
    if (!iter)
      return;

    /* Allocate temp for limit and counter */
    int64_t temp_base = alloc_temp(gen, 2);
    int64_t limit_addr = temp_base;
    int64_t counter_addr = temp_base + 1;

    /* Initialize iterator */
    gen_value(gen, cmd->left);
    emit(gen, "STORE %lld", (long long)iter->address);

    /* Calculate and store limit */
    gen_value(gen, cmd->extra);
    emit(gen, "STORE %lld", (long long)limit_addr);

    /* Skip loop if range is empty */
    emit(gen, "LOAD %lld", (long long)iter->address);
    emit(gen, "SWP b"); /* b = from */
    emit(gen, "LOAD %lld", (long long)limit_addr);
    emit(gen, "SWP c"); /* c = to */
    emit(gen, "RST a");
    if (cmd->type == NODE_FOR_TO) {
      /* from - to > 0 => empty */
      emit(gen, "ADD b");
      emit(gen, "SUB c");
    } else {
      /* to - from > 0 => empty */
      emit(gen, "ADD c");
      emit(gen, "SUB b");
    }
    emit(gen, "JPOS 0"); /* Placeholder for empty range */
    int64_t jump_empty = gen->current_line - 1;

    /* Calculate iteration count = |to - from| + 1 */
    if (cmd->type == NODE_FOR_TO) {
      /* counter = to - from + 1 */
      emit(gen, "LOAD %lld", (long long)limit_addr);
      emit(gen, "SWP b");
      emit(gen, "LOAD %lld", (long long)iter->address);
      emit(gen, "SWP b");
      emit(gen, "SUB b"); /* to - from */
      emit(gen, "INC a");
    } else {
      /* counter = from - to + 1 */
      emit(gen, "LOAD %lld", (long long)iter->address);
      emit(gen, "SWP b");
      emit(gen, "LOAD %lld", (long long)limit_addr);
      emit(gen, "SWP b");
      emit(gen, "SUB b"); /* from - to */
      emit(gen, "INC a");
    }
    emit(gen, "STORE %lld", (long long)counter_addr);

    /* Loop start */
    int64_t loop_start = gen->current_line;

    /* Check counter > 0 */
    emit(gen, "LOAD %lld", (long long)counter_addr);
    emit(gen, "JZERO 0"); /* Placeholder */
    int64_t jump_end = gen->current_line - 1;

    /* Body */
    symtab_push_iterator(gen->symtab, iter);
    gen_commands(gen, cmd->right);
    symtab_pop_iterator(gen->symtab);

    /* Update iterator and counter */
    emit(gen, "LOAD %lld", (long long)iter->address);
    if (cmd->type == NODE_FOR_TO) {
      emit(gen, "INC a");
    } else {
      emit(gen, "DEC a");
    }
    emit(gen, "STORE %lld", (long long)iter->address);

    emit(gen, "LOAD %lld", (long long)counter_addr);
    emit(gen, "DEC a");
    emit(gen, "STORE %lld", (long long)counter_addr);

    emit(gen, "JUMP %lld", (long long)loop_start);

    /* Patch end jump */
    patch_jump(jump_end, gen->current_line);
    patch_jump(jump_empty, gen->current_line);
    free_temp(gen, 2);
    break;
  }

  case NODE_PROC_CALL:
    gen_proc_call(gen, cmd);
    break;

  case NODE_READ: {
    emit(gen, "READ");
    ASTNode *target = cmd->left;
    Symbol *sym = symtab_lookup(gen->symtab, target->name);
    if (!sym)
      return;

    if (target->type == NODE_IDENT) {
      gen_store_var(gen, sym);
    } else {
      /* Preserve read value while computing address */
      emit(gen, "SWP c"); /* value -> rc */
      gen_load_address(gen, target);
      emit(gen, "SWP b"); /* address -> rb */
      emit(gen, "SWP c"); /* value -> ra */
      emit(gen, "RSTORE b");
    }
    break;
  }

  case NODE_WRITE:
    gen_value(gen, cmd->left);
    emit(gen, "WRITE");
    break;

  default:
    break;
  }
}

/* Check if n is a power of 2 and return the log2, or -1 if not */
static int is_power_of_2(int64_t n) {
  if (n <= 0)
    return -1;
  int count = 0;
  while ((n & 1) == 0) {
    n >>= 1;
    count++;
  }
  return (n == 1) ? count : -1;
}


/* --- FRAGMENT: gen_expression --- */

static void gen_expression(CodeGen *gen, ASTNode *expr) {
  if (expr == NULL)
    return;

  switch (expr->type) {
  case NODE_NUM:
  case NODE_IDENT:
  case NODE_ARRAY_ACCESS:
  case NODE_ARRAY_NUM_ACCESS:
    gen_value(gen, expr);
    break;

  /* ====== POPRAWIONE ====== */
  case NODE_ADD:
    gen_value(gen, expr->left);
    emit(gen, "SWP c");        /* rc = left */
    gen_value(gen, expr->right);
    emit(gen, "ADD c");        /* ra = right + left */
    break;

  case NODE_SUB:
    gen_value(gen, expr->right);
    emit(gen, "SWP c");        /* rc = right */
    gen_value(gen, expr->left);
    emit(gen, "SUB c");        /* ra = left - right */
    break;
  /* ======================= */

  case NODE_MUL: {
    int shift = -1;
    ASTNode *other = NULL;

    if (expr->left->type == NODE_NUM) {
      shift = is_power_of_2(expr->left->value);
      other = expr->right;
    } else if (expr->right->type == NODE_NUM) {
      shift = is_power_of_2(expr->right->value);
      other = expr->left;
    }

    if (shift >= 0 && other) {
      gen_value(gen, other);
      for (int i = 0; i < shift; i++) {
        emit(gen, "SHL a");
      }
    } else {
      gen_value(gen, expr->left);
      emit(gen, "SWP c");
      gen_value(gen, expr->right);
      emit(gen, "SWP d");
      gen_multiply(gen);
    }
    break;
  }

  case NODE_DIV: {
    if (expr->right->type == NODE_NUM) {
      int shift = is_power_of_2(expr->right->value);
      if (shift >= 0) {
        gen_value(gen, expr->left);
        for (int i = 0; i < shift; i++) {
          emit(gen, "SHR a");
        }
        break;
      }
    }
    gen_value(gen, expr->left);
    emit(gen, "SWP c");
    gen_value(gen, expr->right);
    emit(gen, "SWP d");
    gen_divide(gen);
    break;
  }

  case NODE_MOD:
    gen_value(gen, expr->left);
    emit(gen, "SWP c");
    gen_value(gen, expr->right);
    emit(gen, "SWP d");
    gen_modulo(gen);
    break;

  default:
    break;
  }
}


/* Generate condition - emits JZERO placeholder at jump_false */
static void gen_condition(CodeGen *gen, ASTNode *cond, int64_t *jump_false) {
  /* Load both values */
  gen_value(gen, cond->left);
  emit(gen, "SWP b");
  gen_value(gen, cond->right);
  emit(gen, "SWP c"); /* left in rb, right in rc */

  switch (cond->type) {
  case NODE_EQ:
    /* a == b: (a - b == 0) && (b - a == 0) */
    emit(gen, "RST a");
    emit(gen, "ADD b");
    emit(gen, "SUB c");
    emit(gen, "JPOS 0"); /* if a > b, not equal */
    *jump_false = gen->current_line - 1;
    emit(gen, "RST a");
    emit(gen, "ADD c");
    emit(gen, "SUB b");
    /* If a == b, ra == 0, so we need JPOS to skip */
    /* But we want JZERO to skip if NOT equal */
    /* Actually return ra = 1 if equal, ra = 0 if not */
    emit(gen, "JPOS 0"); /* Placeholder for second check */
    {
      int64_t second_check = gen->current_line - 1;
      emit(gen, "RST a");
      emit(gen, "INC a");  /* Equal - set to 1 */
      emit(gen, "JUMP 0"); /* Skip the "not equal" part */
      int64_t skip_not_eq = gen->current_line - 1;

      /* Patch first JPOS to here */
      Instruction *inst = instr_head;
      for (int64_t i = 0; i < *jump_false && inst; i++, inst = inst->next)
        ;
      if (inst) {
        snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
                 (long long)gen->current_line);
      }

      /* Patch second JPOS to here */
      inst = instr_head;
      for (int64_t i = 0; i < second_check && inst; i++, inst = inst->next)
        ;
      if (inst) {
        snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
                 (long long)gen->current_line);
      }

      emit(gen, "RST a"); /* Not equal - set to 0 */

      /* Patch skip jump */
      inst = instr_head;
      for (int64_t i = 0; i < skip_not_eq && inst; i++, inst = inst->next)
        ;
      if (inst) {
        snprintf(inst->text, sizeof(inst->text), "JUMP %lld",
                 (long long)gen->current_line);
      }
    }
    emit(gen, "JZERO 0");
    *jump_false = gen->current_line - 1;
    break;

  case NODE_NEQ:
    /* a != b: (a - b > 0) || (b - a > 0) */
    emit(gen, "RST a");
    emit(gen, "ADD b");
    emit(gen, "SUB c");
    emit(gen, "JPOS 0"); /* If a > b, definitely not equal */
    {
      int64_t first_pos = gen->current_line - 1;
      emit(gen, "RST a");
      emit(gen, "ADD c");
      emit(gen, "SUB b");
      emit(gen, "JPOS 0"); /* If b > a, definitely not equal */
      int64_t second_pos = gen->current_line - 1;

      /* Equal - ra = 0, will trigger JZERO */
      emit(gen, "RST a");
      emit(gen, "JUMP 0");
      int64_t skip = gen->current_line - 1;

      /* Not equal - set ra = 1 */
      int64_t not_equal_addr = gen->current_line;
      emit(gen, "RST a");
      emit(gen, "INC a");

      /* Patch jumps */
      Instruction *inst = instr_head;
      for (int64_t i = 0; i < first_pos && inst; i++, inst = inst->next)
        ;
      if (inst)
        snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
                 (long long)not_equal_addr);

      inst = instr_head;
      for (int64_t i = 0; i < second_pos && inst; i++, inst = inst->next)
        ;
      if (inst)
        snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
                 (long long)not_equal_addr);

      inst = instr_head;
      for (int64_t i = 0; i < skip && inst; i++, inst = inst->next)
        ;
      if (inst)
        snprintf(inst->text, sizeof(inst->text), "JUMP %lld",
                 (long long)gen->current_line);
    }
    emit(gen, "JZERO 0");
    *jump_false = gen->current_line - 1;
    break;

  case NODE_LT:
    /* a < b: b - a > 0 */
    emit(gen, "RST a");
    emit(gen, "ADD c");
    emit(gen, "SUB b");
    emit(gen, "JPOS %lld", (long long)(gen->current_line + 2));
    emit(gen, "JUMP 0");
    *jump_false = gen->current_line - 1;
    break;

  case NODE_GT:
    /* a > b: a - b > 0 */
    emit(gen, "RST a");
    emit(gen, "ADD b");
    emit(gen, "SUB c");
    emit(gen, "JPOS %lld", (long long)(gen->current_line + 2));
    emit(gen, "JUMP 0");
    *jump_false = gen->current_line - 1;
    break;

  case NODE_LEQ:
    /* a <= b: a - b == 0 || b > a, i.e., NOT (a > b) */
    emit(gen, "RST a");
    emit(gen, "ADD b");
    emit(gen, "SUB c");
    emit(gen, "JPOS 0");
    *jump_false = gen->current_line - 1;
    break;

  case NODE_GEQ:
    /* a >= b: NOT (a < b), i.e., NOT (b - a > 0) */
    emit(gen, "RST a");
    emit(gen, "ADD c");
    emit(gen, "SUB b");
    emit(gen, "JPOS 0");
    *jump_false = gen->current_line - 1;
    break;

  default:
    emit(gen, "JZERO 0");
    *jump_false = gen->current_line - 1;
    break;
  }
}

static void gen_value(CodeGen *gen, ASTNode *val) {
  if (val == NULL)
    return;

  if (val->type == NODE_NUM) {
    gen_const(gen, val->value);
    return;
  }

  Symbol *sym = symtab_lookup(gen->symtab, val->name);
  if (!sym)
    return;

  if (val->type == NODE_IDENT) {
    gen_load_var(gen, sym);
  } else {
    /* Array access */
    gen_load_address(gen, val);
    emit(gen, "RLOAD a");
  }
}

/* Load address of identifier into ra */
static void gen_load_address(CodeGen *gen, ASTNode *id) {
  Symbol *sym = symtab_lookup(gen->symtab, id->name);
  if (!sym)
    return;

  if (id->type == NODE_IDENT) {
    if (sym->is_reference) {
      emit(gen, "LOAD %lld", (long long)sym->address);
    } else {
      gen_const(gen, sym->address);
    }
    return;
  }

  /* Array access arr[idx] */
  /* Address = base + idx - array_start (for local arrays) */
  /* For T parameters: address is already base-adjusted */

  if (id->type == NODE_ARRAY_NUM_ACCESS) {
    /* Constant index */
    if (sym->is_reference) {
      /* T parameter - load base address and add index */
      emit(gen, "LOAD %lld", (long long)sym->address);
      if (id->value != 0) {
        emit(gen, "SWP b");
        gen_const(gen, id->value);
        emit(gen, "ADD b");
      }
    } else {
      gen_const(gen, sym->address + id->value - sym->array_start);
    }
  } else {
    /* Variable index arr[x] */
    /* Load index */
    gen_value(gen, id->left);
    emit(gen, "SWP b");

    if (sym->is_reference) {
      /* T parameter */
      emit(gen, "LOAD %lld", (long long)sym->address);
      emit(gen, "ADD b");
    } else {
      int64_t offset = sym->address - sym->array_start;
      if (offset >= 0) {
        gen_const(gen, offset);
        emit(gen, "ADD b");
      } else {
        gen_const(gen, -offset);
        emit(gen, "SWP d");
        emit(gen, "RST a");
        emit(gen, "ADD b");
        emit(gen, "SUB d");
      }
    }
  }
}

/* Russian Peasant Multiplication - O(log n) */
/* Input: rc = A, rd = B */
/* Output: ra = A * B */
static void gen_multiply(CodeGen *gen) {
  /* result in re = 0 */
  emit(gen, "RST e");

  /* Loop while B (rd) > 0 */
  int64_t loop_start = gen->current_line;
  emit(gen, "RST a");
  emit(gen, "ADD d");
  emit(gen, "JZERO 0"); /* Exit if B == 0 */
  int64_t exit_jump = gen->current_line - 1;

  /* Check if B is odd: B - (B/2)*2 > 0 */
  emit(gen, "RST a");
  emit(gen, "ADD d");
  emit(gen, "SHR a");
  emit(gen, "SHL a");
  emit(gen, "SWP f"); /* f = (B/2)*2 */
  emit(gen, "RST a");
  emit(gen, "ADD d");
  emit(gen, "SUB f");
  emit(gen, "JZERO 0"); /* Skip add if even */
  int64_t skip_add = gen->current_line - 1;

  /* B is odd: result += A */
  emit(gen, "RST a");
  emit(gen, "ADD e");
  emit(gen, "ADD c");
  emit(gen, "SWP e");

  /* Patch skip_add to here */
  Instruction *inst = instr_head;
  for (int64_t i = 0; i < skip_add && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JZERO %lld",
             (long long)gen->current_line);

  /* A = A * 2 (shift left) */
  emit(gen, "SHL c");

  /* B = B / 2 (shift right) */
  emit(gen, "SHR d");

  emit(gen, "JUMP %lld", (long long)loop_start);

  /* Patch exit_jump */
  inst = instr_head;
  for (int64_t i = 0; i < exit_jump && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JZERO %lld",
             (long long)gen->current_line);

  /* Move result to ra */
  emit(gen, "RST a");
  emit(gen, "ADD e");
}

/* Binary Long Division - O(log n) */
/* Input: rc = dividend, rd = divisor */
/* Output: ra = quotient */
/* Binary Long Division - O(log n) */
/* Input: rc = dividend, rd = divisor */
/* Output: ra = quotient */
static void gen_divide(CodeGen *gen) {
  /* Check divisor == 0 */
  emit(gen, "RST a");
  emit(gen, "ADD d");
  emit(gen, "JPOS 0"); /* If divisor > 0, continue */
  int64_t nonzero = gen->current_line - 1;

  /* Divisor is 0, return 0 */
  emit(gen, "RST a");
  emit(gen, "JUMP 0"); /* Jump to end */
  int64_t zero_exit = gen->current_line - 1;

  /* Patch nonzero check */
  Instruction *inst = instr_head;
  for (int64_t i = 0; i < nonzero && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
             (long long)gen->current_line);

  /* re = quotient = 0 */
  emit(gen, "RST e");

  /* rf = bit weight = 1 */
  emit(gen, "RST f");
  emit(gen, "INC f");

  /* Copy divisor to rg for scaling */
  emit(gen, "RST a");
  emit(gen, "ADD d");
  emit(gen, "SWP g");

  /* Scale up: while g <= c, g *= 2, f *= 2 */
  int64_t scale_up = gen->current_line;
  emit(gen, "RST a");
  emit(gen, "ADD c");
  emit(gen, "SUB g");
  emit(gen, "JZERO 0"); /* if c < g, stop scaling */
  int64_t stop_scale = gen->current_line - 1;

  emit(gen, "SHL g");
  emit(gen, "SHL f");
  emit(gen, "JUMP %lld", (long long)scale_up);

  /* Patch stop_scale */
  inst = instr_head;
  for (int64_t i = 0; i < stop_scale && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JZERO %lld",
             (long long)gen->current_line);

  /* Scale down: while f >= 1 */
  int64_t scale_down = gen->current_line;
  emit(gen, "RST a");
  emit(gen, "ADD f");
  emit(gen, "JZERO 0"); /* if f == 0, done */
  int64_t done = gen->current_line - 1;

  /* Check if c >= g (equivalent to NOT g > c) */
  /* g - c using max(g-c, 0). If > 0 then g > c. */
  emit(gen, "RST a");
  emit(gen, "ADD g");
  emit(gen, "SUB c");
  emit(gen, "JPOS 0"); /* if g > c, skip subtraction */
  int64_t skip_sub = gen->current_line - 1;

  /* c -= g */
  emit(gen, "RST a");
  emit(gen, "ADD c");
  emit(gen, "SUB g");
  emit(gen, "SWP c");

  /* quotient += f */
  emit(gen, "RST a");
  emit(gen, "ADD e");
  emit(gen, "ADD f");
  emit(gen, "SWP e");

  /* Patch skip_sub to here */
  inst = instr_head;
  for (int64_t i = 0; i < skip_sub && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
             (long long)gen->current_line);

  /* g /= 2, f /= 2 */
  emit(gen, "SHR g");
  emit(gen, "SHR f");

  emit(gen, "JUMP %lld", (long long)scale_down);

  /* Patch done */
  inst = instr_head;
  for (int64_t i = 0; i < done && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JZERO %lld",
             (long long)gen->current_line);

  /* Move quotient to ra */
  emit(gen, "RST a");
  emit(gen, "ADD e");
  emit(gen, "JUMP 0");
  int64_t final_end = gen->current_line - 1;

  /* Patch zero_exit */
  inst = instr_head;
  for (int64_t i = 0; i < zero_exit && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JUMP %lld",
             (long long)gen->current_line);

  /* Patch final_end */
  inst = instr_head;
  for (int64_t i = 0; i < final_end && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JUMP %lld",
             (long long)gen->current_line);
}

/* Modulo - O(log n) */
/* Input: rc = dividend, rd = divisor */
/* Output: ra = remainder */
/* Modulo - O(log n) */
/* Input: rc = dividend, rd = divisor */
/* Output: ra = remainder */
static void gen_modulo(CodeGen *gen) {
  /* Check divisor == 0 */
  emit(gen, "RST a");
  emit(gen, "ADD d");
  emit(gen, "JPOS 0");
  int64_t nonzero = gen->current_line - 1;

  /* Divisor is 0, return 0 */
  emit(gen, "RST a");
  emit(gen, "JUMP 0");
  int64_t zero_exit = gen->current_line - 1;

  /* Patch nonzero */
  Instruction *inst = instr_head;
  for (int64_t i = 0; i < nonzero && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
             (long long)gen->current_line);

  /* rf = bit weight = 1 */
  emit(gen, "RST f");
  emit(gen, "INC f");

  /* Copy divisor to rg */
  emit(gen, "RST a");
  emit(gen, "ADD d");
  emit(gen, "SWP g");

  /* Scale up */
  int64_t scale_up = gen->current_line;
  emit(gen, "RST a");
  emit(gen, "ADD c");
  emit(gen, "SUB g");
  emit(gen, "JZERO 0");
  int64_t stop_scale = gen->current_line - 1;

  emit(gen, "SHL g");
  emit(gen, "SHL f");
  emit(gen, "JUMP %lld", (long long)scale_up);

  inst = instr_head;
  for (int64_t i = 0; i < stop_scale && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JZERO %lld",
             (long long)gen->current_line);

  /* Scale down - same as divide but don't track quotient */
  int64_t scale_down = gen->current_line;
  emit(gen, "RST a");
  emit(gen, "ADD f");
  emit(gen, "JZERO 0");
  int64_t done = gen->current_line - 1;

  /* Check if c >= g */
  emit(gen, "RST a");
  emit(gen, "ADD g");
  emit(gen, "SUB c");
  emit(gen, "JPOS 0");
  int64_t skip_sub = gen->current_line - 1;

  emit(gen, "RST a");
  emit(gen, "ADD c");
  emit(gen, "SUB g");
  emit(gen, "SWP c");

  /* Patch skip_sub */
  inst = instr_head;
  for (int64_t i = 0; i < skip_sub && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JPOS %lld",
             (long long)gen->current_line);

  emit(gen, "SHR g");
  emit(gen, "SHR f");
  emit(gen, "JUMP %lld", (long long)scale_down);

  inst = instr_head;
  for (int64_t i = 0; i < done && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JZERO %lld",
             (long long)gen->current_line);

  /* Remainder is in rc */
  emit(gen, "RST a");
  emit(gen, "ADD c");
  emit(gen, "JUMP 0");
  int64_t final_end = gen->current_line - 1;

  inst = instr_head;
  for (int64_t i = 0; i < zero_exit && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JUMP %lld",
             (long long)gen->current_line);

  inst = instr_head;
  for (int64_t i = 0; i < final_end && inst; i++, inst = inst->next)
    ;
  if (inst)
    snprintf(inst->text, sizeof(inst->text), "JUMP %lld",
             (long long)gen->current_line);
}