From df0a5d738f5a9f15acc31d2bd515f0e391bad006 Mon Sep 17 00:00:00 2001
From: Andy Wingo <wingo@igalia.com>
Date: Mon, 25 Mar 2019 15:08:05 +0100
Subject: [PATCH] Better loading of args to registers; add tests

---
 .gitignore          |   3 +
 jit.h               |  12 +++-
 jit/x86.c           | 168 ++++++++++++++++++++++++++++++++++++++++----
 tests/Makefile      |   2 +-
 tests/test-addr_d.c |  28 ++++++++
 tests/test-addr_f.c |  28 ++++++++
 tests/test-addx.c   |  65 +++++++++++++++++
 7 files changed, 291 insertions(+), 15 deletions(-)
 create mode 100644 tests/test-addr_d.c
 create mode 100644 tests/test-addr_f.c
 create mode 100644 tests/test-addx.c

diff --git a/.gitignore b/.gitignore
index dc7bc2589..ffa28c972 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,6 @@
 /lightning.info
 /tests/test-addr
 /tests/test-addi
+/tests/test-addr_d
+/tests/test-addr_f
+/tests/test-addx
diff --git a/jit.h b/jit.h
index f52263f07..3b14591b0 100644
--- a/jit.h
+++ b/jit.h
@@ -127,7 +127,8 @@ typedef enum jit_arg_abi
   JIT_ARG_ABI_INT64,
   JIT_ARG_ABI_POINTER,
   JIT_ARG_ABI_FLOAT,
-  JIT_ARG_ABI_DOUBLE
+  JIT_ARG_ABI_DOUBLE,
+  JIT_ARG_ABI_INTMAX = CHOOSE_32_64(JIT_ARG_ABI_INT32, JIT_ARG_ABI_INT64)
 } jit_arg_abi_t;
 
 typedef struct jit_arg
@@ -142,6 +143,12 @@ typedef struct jit_arg
   } loc;
 } jit_arg_t;
 
+typedef union jit_anyreg
+{
+  jit_gpr_t gpr;
+  jit_fpr_t fpr;
+} jit_anyreg_t;
+
 JIT_API jit_bool_t init_jit(void);
 
 JIT_API jit_state_t *jit_new_state(void);
@@ -166,6 +173,9 @@ JIT_API void jit_callr(jit_state_t *, jit_gpr_t f,
                        const jit_arg_t args[]);
 JIT_API void jit_receive(jit_state_t*, size_t argc,
                          const jit_arg_abi_t abi[], jit_arg_t args[]);
+JIT_API void jit_load_args(jit_state_t *_jit, size_t argc,
+                           const jit_arg_abi_t abi[], jit_arg_t args[],
+                           const jit_anyreg_t regs[]);
 
 #define JIT_PROTO_0(stem, ret) \
   ret jit_##stem (jit_state_t* _jit)
diff --git a/jit/x86.c b/jit/x86.c
index 1b617372a..08bda792a 100644
--- a/jit/x86.c
+++ b/jit/x86.c
@@ -363,6 +363,35 @@ jit_callr(jit_state_t *_jit, jit_gpr_t f,
   callr(_jit, f);
 }
 
+static jit_bool_t
+is_fpr_arg(jit_arg_abi_t arg)
+{
+  switch (arg)
+    {
+    case JIT_ARG_ABI_UINT8:
+    case JIT_ARG_ABI_INT8:
+    case JIT_ARG_ABI_UINT16:
+    case JIT_ARG_ABI_INT16:
+    case JIT_ARG_ABI_UINT32:
+    case JIT_ARG_ABI_INT32:
+    case JIT_ARG_ABI_UINT64:
+    case JIT_ARG_ABI_INT64:
+    case JIT_ARG_ABI_POINTER:
+      return 0;
+    case JIT_ARG_ABI_FLOAT:
+    case JIT_ARG_ABI_DOUBLE:
+      return 1;
+    default:
+      abort();
+    }
+}
+
+static jit_bool_t
+is_gpr_arg(jit_arg_abi_t arg)
+{
+  return !is_fpr_arg(arg);
+}
+
 void
 jit_receive(jit_state_t *_jit,
             size_t argc, const jit_arg_abi_t abi[], jit_arg_t args[])
@@ -400,16 +429,7 @@ jit_receive(jit_state_t *_jit,
 #endif
 
   for (size_t i = 0; i < argc; i++) {
-    switch (abi[i]) {
-    case JIT_ARG_ABI_UINT8:
-    case JIT_ARG_ABI_INT8:
-    case JIT_ARG_ABI_UINT16:
-    case JIT_ARG_ABI_INT16:
-    case JIT_ARG_ABI_UINT32:
-    case JIT_ARG_ABI_INT32:
-    case JIT_ARG_ABI_UINT64:
-    case JIT_ARG_ABI_INT64:
-    case JIT_ARG_ABI_POINTER:
+    if (is_gpr_arg(abi[i])) {
       if (gpr_arg_idx < gpr_arg_count) {
         args[i].kind = JIT_ARG_LOC_GPR;
         args[i].loc.gpr = gpr_args[gpr_arg_idx];
@@ -417,9 +437,8 @@ jit_receive(jit_state_t *_jit,
       } else {
         abort();
       }
-      break;
-    case JIT_ARG_ABI_FLOAT:
-    case JIT_ARG_ABI_DOUBLE:
+    } else {
+      ASSERT(is_fpr_arg(abi[i]));
       if (fpr_arg_idx < fpr_arg_count) {
         args[i].kind = JIT_ARG_LOC_FPR;
         args[i].loc.fpr = fpr_args[fpr_arg_idx];
@@ -427,7 +446,130 @@ jit_receive(jit_state_t *_jit,
       } else {
         abort();
       }
+    }
+  }
+}
+
+void
+jit_load_args(jit_state_t *_jit, size_t argc,
+              const jit_arg_abi_t abi[], jit_arg_t args[],
+              const jit_anyreg_t regs[])
+{
+  /* First shuffle the arguments in registers into position.  */
+  for (size_t i = 0; i < argc; i++) {
+    const jit_arg_t arg = args[i];
+    const jit_anyreg_t reg = regs[i];
+    switch (arg.kind) {
+    case JIT_ARG_LOC_IMM:
+      abort();
+    case JIT_ARG_LOC_GPR:
+      {
+        if (arg.loc.gpr != reg.gpr)
+          /* Arg in a reg but it's not the right one.  See if this reg
+             holds some other arg, and swap if so.  */
+          for (size_t j=i+1; j<argc; j++)
+            if (args[j].kind == JIT_ARG_LOC_GPR && args[j].loc.gpr == reg.gpr)
+              {
+                xchgr(_jit, rn(arg.loc.gpr), rn(reg.gpr));
+                args[j].loc.gpr = arg.loc.gpr;
+                args[i].loc.gpr = reg.gpr;
+                break;
+              }
+        if (arg.loc.gpr != reg.gpr)
+          /* Arg in reg, but it's not the right one, and the desired reg
+             is free.  */
+          {
+            movr(_jit, rn(reg.gpr), rn(arg.loc.gpr));
+            args[i].loc.gpr = reg.gpr;
+          }
+      }
+    case JIT_ARG_LOC_FPR:
+      {
+        if (arg.loc.fpr != reg.fpr)
+          /* Arg in a reg but it's not the right one.  See if this reg
+             holds some other arg, and swap if so.  */
+          for (size_t j=i+1; j<argc; j++)
+            if (args[j].kind == JIT_ARG_LOC_FPR && args[j].loc.fpr == reg.fpr)
+              {
+                jit_fpr_t tmp = get_temp_xpr(_jit);
+                movr_d (_jit, tmp, rn(arg.loc.fpr));
+                movr_d (_jit, rn(arg.loc.fpr), rn(reg.fpr));
+                movr_d (_jit, rn(reg.fpr), tmp);
+                unget_temp_xpr(_jit);
+                args[j].loc.fpr = arg.loc.fpr;
+                args[i].loc.fpr = reg.fpr;
+                break;
+              }
+        if (arg.loc.fpr != reg.fpr)
+          /* Arg in reg, but it's not the right one, and the desired reg
+             is free.  */
+          {
+            movr_d(_jit, rn(reg.fpr), rn(arg.loc.fpr));
+            args[i].loc.fpr = reg.fpr;
+          }
+      }
+    case JIT_ARG_LOC_MEM:
+      /* Load spilled arguments once we're done with registers.  */
       break;
+    default:
+      abort();
+    }
+  }
+
+  /* Load spilled arguments from memory into registers.  */
+  for (size_t i = 0; i < argc; i++) {
+    const jit_arg_t arg = args[i];
+    const jit_anyreg_t reg = regs[i];
+    switch (arg.kind) {
+    case JIT_ARG_LOC_IMM:
+      abort();
+    case JIT_ARG_LOC_GPR:
+    case JIT_ARG_LOC_FPR:
+      break;
+    case JIT_ARG_LOC_MEM:
+      {
+        jit_gpr_t base = arg.loc.mem.base;
+        ptrdiff_t offset = arg.loc.mem.offset;
+        switch (abi[i]) {
+        case JIT_ARG_ABI_UINT8:
+          jit_ldxi_uc(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_INT8:
+          jit_ldxi_c(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_UINT16:
+          jit_ldxi_us(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_INT16:
+          jit_ldxi_s(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_UINT32:
+          jit_ldxi_ui(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_INT32:
+          jit_ldxi_i(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_UINT64:
+          jit_ldxi_l(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_INT64:
+          jit_ldxi_l(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_POINTER:
+          jit_ldxi_l(_jit, reg.gpr, base, offset);
+          break;
+        case JIT_ARG_ABI_FLOAT:
+          jit_ldxi_f(_jit, reg.fpr, base, offset);
+          break;
+        case JIT_ARG_ABI_DOUBLE:
+          jit_ldxi_d(_jit, reg.fpr, base, offset);
+          break;
+        default:
+          abort();
+        }
+      }
+    default:
+      abort();
     }
   }
 }
diff --git a/tests/Makefile b/tests/Makefile
index ead09bdab..2dbad2f9a 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -1,4 +1,4 @@
-TESTS = addr addi
+TESTS=$(sort $(patsubst test-%.c,%,$(wildcard test-*.c)))
 
 CC = gcc
 CFLAGS = -Wall -O0 -g
diff --git a/tests/test-addr_d.c b/tests/test-addr_d.c
new file mode 100644
index 000000000..cae33cd89
--- /dev/null
+++ b/tests/test-addr_d.c
@@ -0,0 +1,28 @@
+#include "test.h"
+
+static void
+run_test(jit_state_t *j, uint8_t *arena_base, size_t arena_size)
+{
+  jit_begin(j, arena_base, arena_size);
+
+  jit_arg_abi_t abi[] = { JIT_ARG_ABI_DOUBLE, JIT_ARG_ABI_DOUBLE };
+  jit_arg_t args[2];
+  jit_receive(j, 2, abi, args);
+  ASSERT(args[0].kind == JIT_ARG_LOC_FPR);
+  ASSERT(args[1].kind == JIT_ARG_LOC_FPR);
+  jit_addr_d(j, JIT_F0, args[0].loc.fpr, args[1].loc.fpr);
+  jit_retr_d(j, JIT_F0);
+
+  size_t size = 0;
+  void* ret = jit_end(j, &size);
+
+  double (*f)(double, double) = ret;
+  ASSERT(f(42., 69.) == 111.);
+  ASSERT(f(42.5, 69.5) == 112.);
+}
+
+int
+main (int argc, char *argv[])
+{
+  return main_helper(argc, argv, run_test);
+}
diff --git a/tests/test-addr_f.c b/tests/test-addr_f.c
new file mode 100644
index 000000000..c85394a8f
--- /dev/null
+++ b/tests/test-addr_f.c
@@ -0,0 +1,28 @@
+#include "test.h"
+
+static void
+run_test(jit_state_t *j, uint8_t *arena_base, size_t arena_size)
+{
+  jit_begin(j, arena_base, arena_size);
+
+  jit_arg_abi_t abi[] = { JIT_ARG_ABI_FLOAT, JIT_ARG_ABI_FLOAT };
+  jit_arg_t args[2];
+  jit_receive(j, 2, abi, args);
+  ASSERT(args[0].kind == JIT_ARG_LOC_FPR);
+  ASSERT(args[1].kind == JIT_ARG_LOC_FPR);
+  jit_addr_f(j, JIT_F0, args[0].loc.fpr, args[1].loc.fpr);
+  jit_retr_f(j, JIT_F0);
+
+  size_t size = 0;
+  void* ret = jit_end(j, &size);
+
+  float (*f)(float, float) = ret;
+  ASSERT(f(42.f, 69.f) == 111.f);
+  ASSERT(f(42.5f, 69.5f) == 112.f);
+}
+
+int
+main (int argc, char *argv[])
+{
+  return main_helper(argc, argv, run_test);
+}
diff --git a/tests/test-addx.c b/tests/test-addx.c
new file mode 100644
index 000000000..664c13998
--- /dev/null
+++ b/tests/test-addx.c
@@ -0,0 +1,65 @@
+#include "test.h"
+
+static void
+run_test(jit_state_t *j, uint8_t *arena_base, size_t arena_size)
+{
+  jit_begin(j, arena_base, arena_size);
+
+  const jit_arg_abi_t abi[] = { JIT_ARG_ABI_INTMAX, JIT_ARG_ABI_INTMAX };
+  jit_arg_t args[2];
+  const jit_anyreg_t regs[] = { { .gpr=JIT_R0 }, { .gpr=JIT_R1 }};
+
+  jit_receive(j, 2, abi, args);
+  jit_load_args(j, 2, abi, args, regs);
+  jit_movi(j, JIT_R2, 0);
+  jit_addcr(j, JIT_R0, JIT_R0, JIT_R1);
+  jit_addxi(j, JIT_R2, JIT_R2, 0);
+  jit_retr(j, JIT_R2);
+
+  size_t size = 0;
+  void* ret = jit_end(j, &size);
+
+  intmax_t (*f)(intmax_t, intmax_t) = ret;
+
+  ASSERT(f(0, 0) == 0);
+
+#if __WORDSIZE == 32
+  /* carry */
+  ASSERT(f(0xffffffff, 0xffffffff) == 1);
+  /* overflow */
+  ASSERT(f(0x7fffffff, 1) == 0);
+  /* overflow */
+  ASSERT(f(0x7fffffff, 0x7fffffff) == 0);
+  /* carry */
+  ASSERT(f(0x7fffffff, 0x80000000) == 0);
+  /* carry+overflow */
+  ASSERT(f(0x80000000, 0x80000000) == 1);
+#else
+  /* nothing */
+  ASSERT(f(0xffffffff, 0xffffffff) == 0);
+  /* nothing */
+  ASSERT(f(0x7fffffff, 1) == 0);
+  /* nothing */
+  ASSERT(f(0x7fffffff, 0x7fffffff) == 0);
+  /* nothing */
+  ASSERT(f(0x7fffffff, 0x80000000) == 0);
+  /* nothing */
+  ASSERT(f(0x80000000, 0x80000000) == 0);
+  /* carry */
+  ASSERT(f(0xffffffffffffffff, 0xffffffffffffffff) == 1);
+  /* overflow */
+  ASSERT(f(0x7fffffffffffffff, 1) == 0);
+  /* overflow */
+  ASSERT(f(0x7fffffffffffffff, 0x7fffffffffffffff) == 0);
+  /* overflow */
+  ASSERT(f(0x7fffffffffffffff, 0x8000000000000000) == 0);
+  /* carry+overflow */
+  ASSERT(f(0x8000000000000000, 0x8000000000000000) == 1);
+#endif
+}
+
+int
+main (int argc, char *argv[])
+{
+  return main_helper(argc, argv, run_test);
+}