From e875af8a12dfe3ff1d7960464fc1bf991429136a Mon Sep 17 00:00:00 2001
From: Nick Gasson <nick@nickg.me.uk>
Date: Fri, 28 Oct 2022 10:57:30 +0100
Subject: [PATCH] Add a simple assembler for JIT IR

---
 src/jit/jit-core.c | 115 +++++++++++++++++++++++++++++++++++++++++++++
 src/jit/jit.h      |   1 +
 test/test_jit.c    |  23 +++++++++
 3 files changed, 139 insertions(+)

diff --git a/src/jit/jit-core.c b/src/jit/jit-core.c
index 94c31811..831bd36e 100644
--- a/src/jit/jit-core.c
+++ b/src/jit/jit-core.c
@@ -934,3 +934,118 @@ ident_t jit_get_name(jit_t *j, jit_handle_t handle)
 {
    return jit_get_func(j, handle)->name;
 }
+
+jit_handle_t jit_assemble(jit_t *j, ident_t name, const char *text)
+{
+   jit_func_t *f = hash_get(j->index, name);
+   if (f != NULL)
+      return f->handle;
+
+   f = xcalloc(sizeof(jit_func_t));
+
+   f->name      = name;
+   f->jit       = j;
+   f->handle    = j->funcs.count;
+   f->next_tier = j->tiers;
+   f->hotness   = f->next_tier ? f->next_tier->threshold : 0;
+   f->entry     = jit_interp;
+
+   hash_put(j->index, name, f);
+   APUSH(j->funcs, f);
+
+   enum { INS, RESULT, ARG1, ARG2, NEWLINE } state = INS;
+
+   static const struct {
+      const char *name;
+      jit_op_t    op;
+      int         nresult;
+      int         nargs;
+   } optab[] = {
+      { "MOV",  J_MOV,  1, 1 },
+      { "ADD",  J_ADD,  1, 2 },
+      { "RECV", J_RECV, 1, 1 },
+      { "SEND", J_SEND, 0, 2 },
+      { "RET",  J_RET,  0, 0 },
+   };
+
+   f->bufsz = 128;
+   f->irbuf = xcalloc_array(f->bufsz, sizeof(jit_ir_t));
+
+   jit_ir_t *ir = NULL;
+   int nresult = 0, nargs = 0;
+   char *copy LOCAL = xstrdup(text);
+   for (char *tok = strtok(copy, " \r\t"); tok; tok = strtok(NULL, " \r\t")) {
+   again:
+      switch (state) {
+      case INS:
+         {
+            int ipos = 0;
+            for (; ipos < ARRAY_LEN(optab)
+                    && strcmp(optab[ipos].name, tok); ipos++)
+               ;
+
+            if (ipos == ARRAY_LEN(optab))
+               fatal_trace("illegal instruction %s", tok);
+
+            nargs = optab[ipos].nargs;
+            nresult = optab[ipos].nresult;
+
+            assert(f->nirs < f->bufsz);
+            ir = &(f->irbuf[f->nirs++]);
+
+            ir->op = optab[ipos].op;
+            ir->size = JIT_SZ_UNSPEC;
+            ir->result = JIT_REG_INVALID;
+
+            state = nresult > 0 ? RESULT : (nargs > 0 ? ARG1 : NEWLINE);
+         }
+         break;
+
+      case RESULT:
+         if (tok[0] != 'R')
+            fatal_trace("expected register name but got '%s'", tok);
+
+         ir->result = atoi(tok + 1);
+         f->nregs = MAX(ir->result + 1, f->nregs);
+         state = ARG1;
+         break;
+
+      case ARG1:
+      case ARG2:
+         {
+            jit_value_t arg;
+            if (tok[0] == 'R') {
+               arg.kind = JIT_VALUE_REG;
+               arg.reg  = atoi(tok + 1);
+               f->nregs = MAX(arg.reg + 1, f->nregs);
+            }
+            else if (tok[0] == '#') {
+               arg.kind  = JIT_VALUE_INT64;
+               arg.int64 = strtoll(tok + 1, NULL, 0);
+            }
+            else if (tok[0] == '\n')
+               fatal_trace("got newline, expecting argument");
+            else
+               fatal_trace("cannot parse argument '%s'", tok);
+
+            if (state == ARG1)
+               ir->arg1 = arg;
+            else
+               ir->arg2 = arg;
+
+            state = state == ARG1 && nargs > 1 ? ARG2 : NEWLINE;
+         }
+         break;
+
+      case NEWLINE:
+         if (*tok != '\n')
+            fatal_trace("expected newline, got '%s'", tok);
+         state = INS;
+         if (*++tok != '\0')
+            goto again;
+         break;
+      }
+   }
+
+   return f->handle;
+}
diff --git a/src/jit/jit.h b/src/jit/jit.h
index 8e243dc0..8a18a686 100644
--- a/src/jit/jit.h
+++ b/src/jit/jit.h
@@ -57,6 +57,7 @@ jit_t *jit_new(void);
 void jit_free(jit_t *j);
 jit_handle_t jit_compile(jit_t *j, ident_t name);
 jit_handle_t jit_lazy_compile(jit_t *j, ident_t name);
+jit_handle_t jit_assemble(jit_t *j, ident_t name, const char *text);
 void *jit_link(jit_t *j, jit_handle_t handle);
 void *jit_get_frame_var(jit_t *j, jit_handle_t handle, uint32_t var);
 void jit_set_lower_fn(jit_t *j, jit_lower_fn_t fn, void *ctx);
diff --git a/test/test_jit.c b/test/test_jit.c
index 1f6b4f99..3ecb4914 100644
--- a/test/test_jit.c
+++ b/test/test_jit.c
@@ -1163,6 +1163,28 @@ START_TEST(test_ffi1)
 }
 END_TEST
 
+START_TEST(test_assemble1)
+{
+   jit_t *j = jit_new();
+
+   const char *text1 =
+      "RECV  R0, #0       \n"
+      "MOV   R1, #2       \n"
+      "ADD   R1, R0, R1   \n"
+      "SEND  #0, R1       \n"
+      "RET                \n";
+
+   jit_handle_t h1 = jit_assemble(j, ident_new("myfunc"), text1);
+
+   jit_scalar_t result, p0 = { .integer = 5 };
+   fail_unless(jit_fastcall(j, h1, &result, p0, p0));
+
+   ck_assert_int_eq(result.integer, 7);
+
+   jit_free(j);
+}
+END_TEST
+
 Suite *get_jit_tests(void)
 {
    Suite *s = suite_create("jit");
@@ -1196,6 +1218,7 @@ Suite *get_jit_tests(void)
    tcase_add_test(tc, test_process1);
    tcase_add_test(tc, test_value1);
    tcase_add_test(tc, test_ffi1);
+   tcase_add_test(tc, test_assemble1);
    suite_add_tcase(s, tc);
 
    return s;
-- 
2.39.2