From a9ffa6656c802fa95e408275cd3844958eff13c5 Mon Sep 17 00:00:00 2001
From: Hao Sun <hao.sun@arm.com>
Date: Tue, 27 Apr 2021 11:37:09 +0300
Subject: [PATCH] DynASM/arm64: support 64-bit jump table

Similar to DynASM/x86[1], this patch allows the creation of 64-bit jump
tables. A new mapping entry '.addr' is introduced and its parameter can
be either variables or the references of pc/global/local labels.

Example:

```
  |  adr x0, >1
  |  ldr x2, [x0, x1]
  |  br x2

  |.jmp_table
  |.align 8
  |1:
  |  .addr &addr
  |  .addr >2
  |  .addr <1
  |  .addr =>pcexpr
  |  .addr ->label

```

[1]. https://github.com/LuaJIT/LuaJIT/pull/683

Change-Id: I6006afb28b2121052afa75fed474269f2e50ab3c
---
 dynasm/dasm_arm64.h   | 36 +++++++++++++++++++++++++++++-------
 dynasm/dasm_arm64.lua | 24 ++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/dynasm/dasm_arm64.h b/dynasm/dasm_arm64.h
index 5ff4414c..9d223446 100644
--- a/dynasm/dasm_arm64.h
+++ b/dynasm/dasm_arm64.h
@@ -19,10 +19,10 @@
 enum {
   DASM_STOP, DASM_SECTION, DASM_ESC, DASM_REL_EXT,
   /* The following actions need a buffer position. */
-  DASM_ALIGN, DASM_REL_LG, DASM_LABEL_LG,
+  DASM_ALIGN, DASM_REL_LG, DASM_LABEL_LG, DASM_ADDR_LG,
   /* The following actions also have an argument. */
-  DASM_REL_PC, DASM_LABEL_PC,
-  DASM_IMM, DASM_IMM6, DASM_IMM12, DASM_IMM13W, DASM_IMM13X, DASM_IMML,
+  DASM_REL_PC, DASM_LABEL_PC, DASM_ADDR_PC,
+  DASM_IMM, DASM_IMM6, DASM_IMM12, DASM_IMM13W, DASM_IMM13X, DASM_IMML, DASM_IMM_PC,
   DASM_VREG,
   DASM__MAX
 };
@@ -251,14 +251,14 @@ void dasm_put(Dst_DECL, int start, ...)
       case DASM_ESC: p++; ofs += 4; break;
       case DASM_REL_EXT: break;
       case DASM_ALIGN: ofs += (ins & 255); b[pos++] = ofs; break;
-      case DASM_REL_LG:
+      case DASM_REL_LG: case DASM_ADDR_LG:
 	n = (ins & 2047) - 10; pl = D->lglabels + n;
 	/* Bkwd rel or global. */
 	if (n >= 0) { CK(n>=10||*pl<0, RANGE_LG); CKPL(lg, LG); goto putrel; }
 	pl += 10; n = *pl;
 	if (n < 0) n = 0;  /* Start new chain for fwd rel if label exists. */
 	goto linkrel;
-      case DASM_REL_PC:
+      case DASM_REL_PC: case DASM_ADDR_PC:
 	pl = D->pclabels + n; CKPL(pc, PC);
       putrel:
 	n = *pl;
@@ -312,6 +312,12 @@ void dasm_put(Dst_DECL, int start, ...)
 	b[pos++] = m;
 	break;
 	}
+      case DASM_IMM_PC: {
+	int m = va_arg(ap, int);
+	b[pos++] = n;
+	b[pos++] = m;
+	break;
+	}
       case DASM_IMML: {
 #ifdef DASM_CHECKS
 	int scale = (ins & 3);
@@ -378,11 +384,11 @@ int dasm_link(Dst_DECL, size_t *szp)
 	case DASM_ESC: p++; break;
 	case DASM_REL_EXT: break;
 	case DASM_ALIGN: ofs -= (b[pos++] + ofs) & (ins & 255); break;
-	case DASM_REL_LG: case DASM_REL_PC: pos++; break;
+	case DASM_REL_LG: case DASM_ADDR_LG: case DASM_REL_PC: case DASM_ADDR_PC: pos++; break;
 	case DASM_LABEL_LG: case DASM_LABEL_PC: b[pos++] += ofs; break;
 	case DASM_IMM: case DASM_IMM6: case DASM_IMM12: case DASM_IMM13W:
 	case DASM_IMML: case DASM_VREG: pos++; break;
-	case DASM_IMM13X: pos += 2; break;
+	case DASM_IMM13X: case DASM_IMM_PC: pos += 2; break;
 	}
       }
       stop: (void)0;
@@ -421,6 +427,7 @@ int dasm_encode(Dst_DECL, void *buffer)
       while (1) {
 	unsigned int ins = *p++;
 	unsigned int action = (ins >> 16);
+	unsigned long long addr = 0;
 	int n = (action >= DASM_ALIGN && action < DASM__MAX) ? *b++ : 0;
 	switch (action) {
 	case DASM_STOP: case DASM_SECTION: goto stop;
@@ -461,6 +468,17 @@ int dasm_encode(Dst_DECL, void *buffer)
 	  ins &= 2047; if (ins >= 20) D->globals[ins-10] = (void *)(base + n);
 	  break;
 	case DASM_LABEL_PC: break;
+	case DASM_ADDR_LG:
+	  if (n < 0) {
+	    addr = (unsigned long long)D->globals[-n];
+	    goto patchaddr;
+	  }
+	case DASM_ADDR_PC:
+	  addr = (unsigned long long)(*DASM_POS2PTR(D, n) + base);
+	patchaddr:
+	  cp[-2] = (unsigned int)(addr);
+	  cp[-1] = (unsigned int)(addr >> 32);
+	  break;
 	case DASM_IMM:
 	  cp[-1] |= (n & ((1<<((ins>>5)&31))-1)) << (ins&31);
 	  break;
@@ -473,6 +491,10 @@ int dasm_encode(Dst_DECL, void *buffer)
 	case DASM_IMM13W:
 	  cp[-1] |= (dasm_imm13(n, n) << 10);
 	  break;
+	case DASM_IMM_PC:
+	  cp[-2] = (unsigned int)(n);
+	  cp[-1] = (unsigned int)(*b++);
+	  break;
 	case DASM_IMM13X:
 	  cp[-1] |= (dasm_imm13(n, *b++) << 10);
 	  break;
diff --git a/dynasm/dasm_arm64.lua b/dynasm/dasm_arm64.lua
index 82412a05..932edb5e 100644
--- a/dynasm/dasm_arm64.lua
+++ b/dynasm/dasm_arm64.lua
@@ -38,8 +38,8 @@ local wline, werror, wfatal, wwarn
 -- CHECK: Keep this in sync with the C code!
 local action_names = {
   "STOP", "SECTION", "ESC", "REL_EXT",
-  "ALIGN", "REL_LG", "LABEL_LG",
-  "REL_PC", "LABEL_PC", "IMM", "IMM6", "IMM12", "IMM13W", "IMM13X", "IMML",
+  "ALIGN", "REL_LG", "LABEL_LG", "ADDR_LG",
+  "REL_PC", "LABEL_PC", "ADDR_PC", "IMM", "IMM6", "IMM12", "IMM13W", "IMM13X", "IMML", "IMM_PC",
   "VREG",
 }
 
@@ -1068,6 +1068,26 @@ map_op[".long_*"] = function(params)
   end
 end
 
+-- Pseudo-opcodes for jump table entry.
+map_op[".addr_1"] = function(params)
+  if not params then return "&addr | >label | <label | ->label | =>label" end
+  if secpos+1 > maxsecpos then wflush() end
+  local prefix = sub(params[1], 1, 1)
+  if prefix == "&" then
+    wputw(0)
+    wputw(0)
+    local imm = match(params[1], "^&(.*)$")
+    waction("IMM_PC", 0, format("(unsigned int)(unsigned long long)(%s)", imm))
+    actargs[#actargs+1] = format("(unsigned int)((unsigned long long)(%s)>>32)", imm)
+  else
+    local mode, n, s = parse_label(params[1], false)
+    if mode == "EXT" then werror("not support extern label reference") end
+    wputw(0)
+    wputw(0)
+    waction("ADDR_"..mode, n, s, 1)
+  end
+end
+
 -- Alignment pseudo-opcode.
 map_op[".align_1"] = function(params)
   if not params then return "numpow2" end