Add jmpi_with_link instruction

The existing calli / callr interface is for ABI calls.  Sometimes though
you want to call some of your own code, just to get the current return
address.  ARM's branch-and-link instructions are ideal for this but they
don't exist on x86; there we emulate them by adding corresponding
pop_link_register / push_link_register instructions that are no-ops on
ARM.

* lightening.h (FOR_EACH_INSTRUCTION): Add jit_jmpi_with_link,
  pop_link_register, push_link_register.
* lightening/arm-cpu.c:
* lightening/x86-cpu.c:
* lightening/aarch64-cpu.c (jmpi_with_link, push_link_register)
  (pop_link_register): Add implementations.
* lightening/arm.h:
* lightening/aarch64.h:
* lightening/x86.h (JIT_LR): New definition.
* tests/link-register.c: New test.

lightening.h

@@ -584,6 +584,10 @@ jit_load_args_3(jit_state_t *_jit, jit_operand_t a, jit_operand_t b,
           M(_G___, jmpr)		\
           M(_p___, jmpi)		\
           M(R____, jmp)			\
+					\
+          M(_p___, jmpi_with_link)	\
+          M(_____, pop_link_register)	\
+          M(_____, push_link_register)	\
   					\
           M(_____, ret)			\
           M(_G___, retr)		\

lightening/aarch64-cpu.c

@@ -2425,6 +2425,22 @@ calli(jit_state_t *_jit, jit_word_t i0)
   return jit_patch_there(_jit, call(_jit), (void*)i0);
 }
 
+static void
+jmpi_with_link(jit_state_t *_jit, jit_word_t i0)
+{
+  return calli(_jit, i0);
+}
+
+static void
+push_link_register(jit_state_t *_jit)
+{
+}
+
+static void
+pop_link_register(jit_state_t *_jit)
+{
+}
+
 static void
 ret(jit_state_t *_jit)
 {

lightening/aarch64.h

@@ -123,6 +123,7 @@
 #define JIT_PLATFORM_CALLEE_SAVE_GPRS _X29, _X30
 
 // x31 is stack pointer.
+#define JIT_LR    _X30
 #define JIT_SP    _X31
 
 #define JIT_F0  _D0

lightening/arm-cpu.c

@@ -2911,6 +2911,22 @@ calli(jit_state_t *_jit, jit_word_t i0)
     jit_patch_there(_jit, T2_BLXI(_jit), (void*)i0);
 }
 
+static void
+jmpi_with_link(jit_state_t *_jit, jit_word_t i0)
+{
+  jit_patch_there(_jit, T2_BLI(_jit), (void*)i0);
+}
+
+static void
+push_link_register(jit_state_t *_jit)
+{
+}
+
+static void
+pop_link_register(jit_state_t *_jit)
+{
+}
+
 static void
 ret(jit_state_t *_jit)
 {

lightening/arm.h

@@ -105,6 +105,7 @@
 #define JIT_V5    _R10
 #define JIT_V6    _R11
 
+#define JIT_LR _R14
 #define JIT_SP _R13
 #define _LR _R14
 #define _PC _R15

lightening/x86-cpu.c

@@ -2552,6 +2552,24 @@ calli(jit_state_t *_jit, jit_word_t i0)
     }
 }
 
+static void
+jmpi_with_link(jit_state_t *_jit, jit_word_t i0)
+{
+  return calli(_jit, i0);
+}
+
+static void
+pop_link_register(jit_state_t *_jit)
+{
+  popr(_jit, jit_gpr_regno (JIT_LR));
+}
+
+static void
+push_link_register(jit_state_t *_jit)
+{
+  pushr(_jit, jit_gpr_regno (JIT_LR));
+}
+
 static void
 jmpr(jit_state_t *_jit, int32_t r0)
 {

lightening/x86.h

@@ -74,6 +74,7 @@
 #endif
 
 #define JIT_SP     _RSP
+#define JIT_LR     JIT_TMP0
 #if __X32
 #  define JIT_R0   _RAX
 #  define JIT_R1   _RCX

tests/link-register.c

@@ -0,0 +1,35 @@
+#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);
+  size_t align = jit_enter_jit_abi(j, 0, 0, 0);
+  jit_load_args_1(j, jit_operand_gpr (JIT_OPERAND_ABI_POINTER, JIT_R0));
+
+  jit_reloc_t call_tramp = jit_jmp (j);
+
+  void *tramp = jit_address (j);
+  jit_pop_link_register (j);
+  jit_movr (j, JIT_R0, JIT_LR);
+  jit_leave_jit_abi(j, 0, 0, align);
+  jit_retr (j, JIT_R0);
+
+  jit_patch_here (j, call_tramp);
+  jit_jmpi_with_link (j, tramp);
+
+  void *expected_link = jit_address_to_function_pointer (jit_address (j));
+
+  size_t size = 0;
+  void* ret = jit_end(j, &size);
+
+  void* (*f)(void) = ret;
+
+  ASSERT(f() == expected_link);
+}
+
+int
+main (int argc, char *argv[])
+{
+  return main_helper(argc, argv, run_test);
+}