2 files changed, 41 insertions, 18 deletions

diff --git a/dynasm/dasm_arm64.h b/dynasm/dasm_arm64.h
index 577b54bc..b2251532 100644
--- a/dynasm/dasm_arm64.h
+++ b/dynasm/dasm_arm64.h
@@ -23,6 +23,7 @@ enum {
   /* 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_VREG,
   DASM__MAX
 };
 
@@ -39,6 +40,7 @@ enum {
 #define DASM_S_RANGE_LG         0x13000000
 #define DASM_S_RANGE_PC         0x14000000
 #define DASM_S_RANGE_REL        0x15000000
+#define DASM_S_RANGE_VREG       0x16000000
 #define DASM_S_UNDEF_LG         0x21000000
 #define DASM_S_UNDEF_PC         0x22000000
 
@@ -312,13 +314,17 @@ void dasm_put(Dst_DECL, int start, ...)
         }
       case DASM_IMML: {
 #ifdef DASM_CHECKS
-        int scale = (p[-2] >> 30);
+        int scale = (ins & 3);
         CK((!(n & ((1<<scale)-1)) && (unsigned int)(n>>scale) < 4096) ||
            (unsigned int)(n+256) < 512, RANGE_I);
 #endif
         b[pos++] = n;
         break;
         }
+      case DASM_VREG:
+        CK(n < 32, RANGE_VREG);
+        b[pos++] = n;
+        break;
       }
     }
   }
@@ -375,7 +381,7 @@ int dasm_link(Dst_DECL, size_t *szp)
         case DASM_REL_LG: case DASM_REL_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: pos++; break;
+        case DASM_IMML: case DASM_VREG: pos++; break;
         case DASM_IMM13X: pos += 2; break;
         }
       }
@@ -468,11 +474,14 @@ int dasm_encode(Dst_DECL, void *buffer)
           cp[-1] |= (dasm_imm13(n, *b++) << 10);
           break;
         case DASM_IMML: {
-          int scale = (p[-2] >> 30);
+          int scale = (ins & 3);
           cp[-1] |= (!(n & ((1<<scale)-1)) && (unsigned int)(n>>scale) < 4096) ?
             ((n << (10-scale)) | 0x01000000) : ((n & 511) << 12);
           break;
           }
+        case DASM_VREG:
+          cp[-1] |= (n & 0x1f) << (ins & 0x1f);
+          break;
         default: *cp++ = ins; break;
         }
       }
diff --git a/dynasm/dasm_arm64.lua b/dynasm/dasm_arm64.lua
index 861ff2b2..82412a05 100644
--- a/dynasm/dasm_arm64.lua
+++ b/dynasm/dasm_arm64.lua
@@ -40,6 +40,7 @@ local action_names = {
   "STOP", "SECTION", "ESC", "REL_EXT",
   "ALIGN", "REL_LG", "LABEL_LG",
   "REL_PC", "LABEL_PC", "IMM", "IMM6", "IMM12", "IMM13W", "IMM13X", "IMML",
+  "VREG",
 }
 
 -- Maximum number of section buffer positions for dasm_put().
@@ -246,9 +247,12 @@ local map_cond = {
 
 local parse_reg_type
 
-local function parse_reg(expr)
+local function parse_reg(expr, shift)
   if not expr then werror("expected register name") end
   local tname, ovreg = match(expr, "^([%w_]+):(@?%l%d+)$")
+  if not tname then
+    tname, ovreg = match(expr, "^([%w_]+):(R[xwqdshb]%b())$")
+  end
   local tp = map_type[tname or expr]
   if tp then
     local reg = ovreg or tp.reg
@@ -266,18 +270,28 @@ local function parse_reg(expr)
       elseif parse_reg_type ~= rt then
         werror("register size mismatch")
       end
-      return r, tp
+      return shl(r, shift), tp
+    end
+  end
+  local vrt, vreg = match(expr, "^R([xwqdshb])(%b())$")
+  if vreg then
+    if not parse_reg_type then
+      parse_reg_type = vrt
+    elseif parse_reg_type ~= vrt then
+      werror("register size mismatch")
     end
+    if shift then waction("VREG", shift, vreg) end
+    return 0
   end
   werror("bad register name `"..expr.."'")
 end
 
 local function parse_reg_base(expr)
   if expr == "sp" then return 0x3e0 end
-  local base, tp = parse_reg(expr)
+  local base, tp = parse_reg(expr, 5)
   if parse_reg_type ~= "x" then werror("bad register type") end
   parse_reg_type = false
-  return shl(base, 5), tp
+  return base, tp
 end
 
 local parse_ctx = {}
@@ -403,7 +417,7 @@ local function parse_imm_load(imm, scale)
     end
     werror("out of range immediate `"..imm.."'")
   else
-    waction("IMML", 0, imm)
+    waction("IMML", scale, imm)
     return 0
   end
 end
@@ -462,6 +476,7 @@ end
 
 local function parse_load(params, nparams, n, op)
   if params[n+2] then werror("too many operands") end
+  local scale = shr(op, 30)
   local pn, p2 = params[n], params[n+1]
   local p1, wb = match(pn, "^%[%s*(.-)%s*%](!?)$")
   if not p1 then
@@ -470,14 +485,13 @@ local function parse_load(params, nparams, n, op)
       if reg and tailr ~= "" then
         local base, tp = parse_reg_base(reg)
         if tp then
-          waction("IMML", 0, format(tp.ctypefmt, tailr))
+          waction("IMML", scale, format(tp.ctypefmt, tailr))
           return op + base
         end
       end
     end
     werror("expected address operand")
   end
-  local scale = shr(op, 30)
   if p2 then
     if wb == "!" then werror("bad use of '!'") end
     op = op + parse_reg_base(p1) + parse_imm(p2, 9, 12, 0, true) + 0x400
@@ -494,7 +508,7 @@ local function parse_load(params, nparams, n, op)
         op = op + parse_imm_load(imm, scale)
       else
         local p2b, p3b, p3s = match(p2a, "^,%s*([^,%s]*)%s*,?%s*(%S*)%s*(.*)$")
-        op = op + shl(parse_reg(p2b), 16) + 0x00200800
+        op = op + parse_reg(p2b, 16) + 0x00200800
         if parse_reg_type ~= "x" and parse_reg_type ~= "w" then
           werror("bad index register type")
         end
@@ -620,7 +634,7 @@ local function alias_bfx(p)
 end
 
 local function alias_bfiz(p)
-  parse_reg(p[1])
+  parse_reg(p[1], 0)
   if parse_reg_type == "w" then
     p[3] = "#-("..p[3]:sub(2)..")%32"
     p[4] = "#("..p[4]:sub(2)..")-1"
@@ -631,7 +645,7 @@ local function alias_bfiz(p)
 end
 
 local alias_lslimm = op_alias("ubfm_4", function(p)
-  parse_reg(p[1])
+  parse_reg(p[1], 0)
   local sh = p[3]:sub(2)
   if parse_reg_type == "w" then
     p[3] = "#-("..sh..")%32"
@@ -891,15 +905,15 @@ local function parse_template(params, template, nparams, pos)
   for p in gmatch(sub(template, 9), ".") do
     local q = params[n]
     if p == "D" then
-      op = op + parse_reg(q); n = n + 1
+      op = op + parse_reg(q, 0); n = n + 1
     elseif p == "N" then
-      op = op + shl(parse_reg(q), 5); n = n + 1
+      op = op + parse_reg(q, 5); n = n + 1
     elseif p == "M" then
-      op = op + shl(parse_reg(q), 16); n = n + 1
+      op = op + parse_reg(q, 16); n = n + 1
     elseif p == "A" then
-      op = op + shl(parse_reg(q), 10); n = n + 1
+      op = op + parse_reg(q, 10); n = n + 1
     elseif p == "m" then
-      op = op + shl(parse_reg(params[n-1]), 16)
+      op = op + parse_reg(params[n-1], 16)
 
     elseif p == "p" then
       if q == "sp" then params[n] = "@x31" end