From 39bd29f7dea6426db6fa78b4f18aae9aa76d8bc4 Mon Sep 17 00:00:00 2001
From: Ian Moffett <ian@osmora.org>
Date: Fri, 3 Oct 2025 16:36:06 -0400
Subject: np: piir: Add bitmap based register allocation

Introduce register allocation via a bitmap where each bit corresponds to
a specific register index.

Signed-off-by: Ian Moffett <ian@osmora.org>
---
 src/sys/arch/amd64/np/piir_conv.c | 46 ++++++++++++++++++++++++++++++++++++++-
 src/sys/include/arch/amd64/piir.h | 21 ++++++++++++++++++
 src/sys/include/np/piir.h         |  3 +++
 3 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/src/sys/arch/amd64/np/piir_conv.c b/src/sys/arch/amd64/np/piir_conv.c
index 280e0a7..f949f48 100644
--- a/src/sys/arch/amd64/np/piir_conv.c
+++ b/src/sys/arch/amd64/np/piir_conv.c
@@ -55,9 +55,25 @@ typedef enum {
     R32_ESP,
     R32_EBP,
     R32_ESI,
-    R32_RDI
+    R32_RDI,
+    __R32_MAX
 } r32_t;
 
+/*
+ * Valid 64-bit register IDs
+ */
+typedef enum {
+    R64_RAX,
+    R64_RCX,
+    R64_RDX,
+    R64_RBX,
+    R64_RSP,
+    R64_RBP,
+    R64_RSI,
+    R64_RDI,
+    __R64_MAX
+} r64_t;
+
 /* SYS-V ABI specific */
 #define R32_RETVAL R32_EAX
 
@@ -128,3 +144,31 @@ md_piir_decode(struct np_work *work, struct piir_vm *vm, ir_byte_t input)
     }
     return 0;
 }
+
+reg_t
+md_alloc_reg(struct np_work *work, struct piir_vm *vm, int flags)
+{
+    if (work == NULL || vm == NULL) {
+        return -EINVAL;
+    }
+
+    /* If a bit is unset, it is free */
+    for (int i = 0; i < __R32_MAX; ++i) {
+        if (!ISSET(vm->regset, BIT(i))) {
+            vm->regset |= BIT(i);
+            return i;
+        }
+    }
+
+    return -1;
+}
+
+void
+md_free_reg(struct np_work *work, struct piir_vm *vm, reg_t reg)
+{
+    if (work == NULL || vm == NULL) {
+        return;
+    }
+
+    vm->regset &= ~BIT(reg);
+}
diff --git a/src/sys/include/arch/amd64/piir.h b/src/sys/include/arch/amd64/piir.h
index 75236b9..d994339 100644
--- a/src/sys/include/arch/amd64/piir.h
+++ b/src/sys/include/arch/amd64/piir.h
@@ -54,4 +54,25 @@ ssize_t md_piir_decode(
     ir_byte_t input
 );
 
+/*
+ * Allocate a register ID for use in the IR
+ *
+ * @work: Current work
+ * @vm: Virtual machine state
+ * @flags: Optional flags
+ *
+ * Returns the allocated register ID on success, otherwise
+ * a less than zero value on failure
+ */
+reg_t md_alloc_reg(struct np_work *work, struct piir_vm *vm, int flags);
+
+/*
+ * Free one or more registers
+ *
+ * @work: Current work
+ * @vm: Virtual machine state
+ * @reg: Register to free
+ */
+void md_free_reg(struct np_work *work, struct piir_vm *vm, reg_t reg);
+
 #endif  /* !_MACHINE_PIIR_H_ */
diff --git a/src/sys/include/np/piir.h b/src/sys/include/np/piir.h
index 5625409..b892231 100644
--- a/src/sys/include/np/piir.h
+++ b/src/sys/include/np/piir.h
@@ -42,6 +42,7 @@
  */
 typedef int8_t ir_byte_t;
 typedef int8_t md_byte_t;
+typedef int8_t reg_t;
 
 /*
  * The maxiumum size of the bytecode stack per
@@ -65,11 +66,13 @@ typedef int8_t md_byte_t;
  * @code: Generated machine code
  * @last_ir: Last IR byte used
  * @code_i: Current index into code buffer
+ * @regset: Register set bitmap
  */
 struct piir_vm {
     md_byte_t code[4096];
     ir_byte_t last_ir;
     uint32_t code_i;
+    uint8_t regset;
 };
 
 /*
-- 
cgit v1.2.3