/**
 * Register allocator.
 * Uses a simple stack-based algorithm.
 */
#include <stdio.h>
#include <stdlib.h>

#include "io.h"
#include "ralloc.h"
#include "riscv.h"
#include "types.h"

/* Order of temporary registers to allocate. */
const reg_t ralloc_regs[] = { A0, A1, A2, A3, A4, A5, A6, A7,
                              T0, T1, T2, T3, T4, T5, T6 };

ralloc_t ralloc(void) {
    return (ralloc_t){ .regs = { false } };
}

reg_t ralloc_next(ralloc_t *ra) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        if (!ra->regs[i]) {
            ra->regs[i] = true;
            return ralloc_regs[i];
        }
    }
    bail("out of registers");
}

reg_t ralloc_next_except(ralloc_t *ra, reg_t avoid) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        if (!ra->regs[i] && ralloc_regs[i] != avoid) {
            ra->regs[i] = true;
            return ralloc_regs[i];
        }
    }
    return ralloc_next(ra);
}

void ralloc_free(ralloc_t *ra, reg_t r) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        if (ralloc_regs[i] == r) {
            ra->regs[i] = false;
            break;
        }
    }
}

void ralloc_reserve(ralloc_t *ra, reg_t r) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        if (ralloc_regs[i] == r) {
            ra->regs[i] = true;
            break;
        }
    }
}

bool ralloc_is_free(ralloc_t *ra, reg_t r) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        if (ralloc_regs[i] == r) {
            return !ra->regs[i];
        }
    }
    return false;
}

void ralloc_free_all(ralloc_t *ra) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        ra->regs[i] = false;
    }
}

void ralloc_save(ralloc_t *ra, bool *reserved) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        if (ralloc_regs[i] != A0) {
            reserved[i] = ra->regs[i];
        }
    }
}

void ralloc_restore(ralloc_t *ra, bool *reserved) {
    for (int i = 0; i < RALLOC_NREGS; i++) {
        if (ralloc_regs[i] != A0) {
            ra->regs[i] = reserved[i];
        }
    }
}