3 files changed, 80 insertions, 9 deletions

diff --git a/src/sys/arch/amd64/np/piir_conv.c b/src/sys/arch/amd64/np/piir_conv.c
index f949f48..3cb49f2 100644
--- a/src/sys/arch/amd64/np/piir_conv.c
+++ b/src/sys/arch/amd64/np/piir_conv.c
@@ -76,6 +76,7 @@ typedef enum {
 
 /* SYS-V ABI specific */
 #define R32_RETVAL R32_EAX
+#define R64_RETVAL R64_RAX
 
 /* Declare an instruction array */
 #define INST_DECL(...) ((inst_t)__VA_ARGS__)
@@ -92,6 +93,10 @@ typedef enum {
 #define OP_LOAD32_R32(IMM32, RD) INST_DECL({0xB8 + (RD), (IMM32)})
 #define OP_LOAD32_R32_LEN 5
 
+/* MOV R64, IMM64 (REX.W | 0xB8 + rd) */
+#define OP_LOAD64_R64(IMM64, RD) INST_DECL({0x48, 0xB8 + (RD), (IMM64)})
+#define OP_LOAD64_R64_LEN 10
+
 /*
  * Push an instruction byte into the virtual
  * machine descriptor's code buffer
@@ -109,11 +114,32 @@ vm_push(struct piir_vm *vm, inst_t inst, size_t len)
     return len;
 }
 
+/*
+ * Push a 64-bit value
+ */
+static void
+vm_push64(struct piir_vm *vm, uint64_t v)
+{
+    size_t max_len;
+    uint64_t *p;
+
+    max_len = (vm->code_i + sizeof(v));
+    if (max_len >= sizeof(vm->code) - 1) {
+        return;
+    }
+
+    p = (uint64_t *)&vm->code[vm->code_i];
+    *p = v;
+    vm->code_i += sizeof(v);
+}
+
 ssize_t
 md_piir_decode(struct np_work *work, struct piir_vm *vm, ir_byte_t input)
 {
+    struct symbol *sym;
     struct piir_stack *stack;
     ssize_t len;
+    int error;
 
     if (vm == NULL) {
         return -EINVAL;
@@ -137,6 +163,34 @@ md_piir_decode(struct np_work *work, struct piir_vm *vm, ir_byte_t input)
             vm, OP_LOAD32_R32(input, R32_RETVAL),
             OP_LOAD32_R32_LEN
         );
+    case PIIR_RET_SYMBOL:
+        /* Get the symbol ID */
+        if ((input = piir_pop(stack)) < 0) {
+            pr_error("failed to pop element from stack\n");
+            return input;
+        }
+
+        /* Lookup the symbol */
+        error = symbol_lookup_id(&work->symlist, input, &sym);
+        if (error < 0) {
+            pr_error("failed to lookup symbol %d\n", input);
+            return error;
+        }
+
+        /* Load some padding into our register */
+        error = vm_push(vm, OP_LOAD64_R64(0, R64_RETVAL), OP_LOAD64_R64_LEN);
+        if (error < 0) {
+            pr_error("failed to push load op\n");
+            return error;
+        }
+
+        /*
+         * Move back 8 bytes, load the symbol address
+         * and overwrite the padding
+         */
+        vm->code_i -= sizeof(sym->addr);
+        vm_push64(vm, (uintptr_t)sym->addr);
+        break;
     case PIIR_NOP:
         return vm_push(vm, OP_NOP, OP_NOP_LEN);
     case PIIR_RET_NIL:
diff --git a/src/sys/include/np/piir.h b/src/sys/include/np/piir.h
index b892231..46f88d1 100644
--- a/src/sys/include/np/piir.h
+++ b/src/sys/include/np/piir.h
@@ -58,6 +58,7 @@ typedef int8_t reg_t;
 #define PIIR_LOAD_R64 0x04  /* Load 64-bit register */
 #define PIIR_RET_NIL  0x05  /* Return nothing */
 #define PIIR_RET_NUM  0x06  /* Return a number */
+#define PIIR_RET_SYMBOL  0x07  /* Return a symbol */
 
 /*
  * Represents the PIIR virtual machine for storing
diff --git a/src/sys/np/core/np_parse.c b/src/sys/np/core/np_parse.c
index 0605a64..3b3ffdb 100644
--- a/src/sys/np/core/np_parse.c
+++ b/src/sys/np/core/np_parse.c
@@ -184,20 +184,36 @@ parse_type(struct np_work *work, struct lex_token *tok)
 static int
 parse_return(struct np_work *work, struct lex_token *tok)
 {
+    struct symbol *sym = NULL;
+    int error;
     tt_t tt;
 
-    /*
-     * For now we'll only accept numbers as return values
-     * so we must see one after the return statement
-     */
-    tt = parse_expect(work, "return", TT_NUMBER, tok);
-    if (tt == TT_NONE) {
-        return -1;
+    if ((error = lex_nom(work, tok)) < 0) {
+        return error;
     }
 
 #define PIIR_PUSH(BYTE) piir_push(work->piir_stack, (BYTE))
-    PIIR_PUSH(PIIR_RET_NUM);
-    PIIR_PUSH(tok->val);
+    switch (tok->token) {
+    case TT_NUMBER:
+        PIIR_PUSH(PIIR_RET_NUM);
+        PIIR_PUSH(tok->val);
+        break;
+    case TT_STR:
+        sym = symbol_alloc(&work->symlist, NULL, tok->val_str);
+        if (sym == NULL) {
+            pr_error("internal error, failed to alloc symbol\n");
+            return -1;
+        }
+        PIIR_PUSH(PIIR_RET_SYMBOL);
+        PIIR_PUSH(sym->id);
+        break;
+    default:
+        pr_error(
+            "line %d: unexpected token %s after 'return'\n",
+            work->line_no, stoktab[tok->token]
+        );
+        return -1;
+    }
 #undef PIIR_PUSH
     return 0;
 }