fix ARM64 too, theoretically

5 files changed, 87 insertions, 20 deletions

diff --git a/lib/arm64/assemble.h b/lib/arm64/assemble.h
index 8a98b7b..1dca7eb 100644
--- a/lib/arm64/assemble.h
+++ b/lib/arm64/assemble.h
@@ -1,9 +1,15 @@
 #pragma once
 #include "dis.h"
+
+static inline int size_of_MOVi64(uint64_t val) {
+    int num_nybbles = val == 0 ? 1 : ((64 - __builtin_clzll(val) + 15) / 16);
+    return 4 * num_nybbles;
+}
+
 static inline void MOVi64(void **codep, int Rd, uint64_t val) {
     int shift_nybbles = 0;
     do {
-        int k = shift_nybbles != 0 ? 1 : 0;
+        int k = shift_nybbles != 0;
         op32(codep, 0xd2800000 | k << 29 | Rd | (val & 0xffff) << 5 |
                     shift_nybbles << 21);
         shift_nybbles++;
@@ -54,3 +60,7 @@ static inline void BR(void **codep, int reg) {
     op32(codep, 0xd61f0000 | reg << 5);
 }
 
+static inline void Bccrel(void **codep, int cc, int offset) {
+    op32(codep, 0x54000000 | (offset / 4) << 5 | cc);
+}
+
diff --git a/lib/arm64/dis-arm64.inc.h b/lib/arm64/dis-arm64.inc.h
index 14a3a92..04349f2 100644
--- a/lib/arm64/dis-arm64.inc.h
+++ b/lib/arm64/dis-arm64.inc.h
@@ -8,15 +8,32 @@ static INLINE void P(adrplabel_label_unk_Xd_1_ADRP)(tdis_ctx ctx, struct bitslic
 }
 static INLINE void P(am_b_target_addr_B_1_B)(tdis_ctx ctx, struct bitslice addr) {
     return P(branch)(ctx, ctx->pc + sext(bs_get(addr, ctx->op), 26) * 4,
-                     /*cond*/ false);
+                     /*cc*/ 0);
 }
 static INLINE void P(am_bl_target_addr_1_BL)(tdis_ctx ctx, struct bitslice addr) {
     return P(branch)(ctx, ctx->pc + sext(bs_get(addr, ctx->op), 26) * 4,
-                     /*cond*/ false);
+                     /*cc*/ 0);
 }
-static INLINE void P(am_brcond_target_B_5_Bcc)(tdis_ctx ctx, struct bitslice target) {
-    return P(branch)(ctx, ctx->pc + sext(bs_get(target, ctx->op), 19) * 4,
-                     /*cond*/ true);
+static INLINE void P(ccode_cond_am_brcond_target_B_1_Bcc)(tdis_ctx ctx, struct bitslice cond, struct bitslice target) {
+    int bits = bs_get(cond, ctx->op);
+    /* Bcc with AL/NV (which is actually just another AL) is useless but possible. */
+    int cc = bits >= 0xe ? 0 : (CC_ARMCC | bits);
+    return P(branch)(ctx, ctx->pc + sext(bs_get(target, ctx->op), 19) * 4, cc);
+}
+static INLINE void P(am_tbrcond_target_B_4_TBNZW)(tdis_ctx ctx, struct bitslice target) {
+    P(branch)(ctx, ctx->pc + sext(bs_get(target, ctx->op), 14) * 4, CC_XBXZ);
+    if (TDIS_CTX_MODIFY(ctx)) {
+        /* ditto CBNZ on ARM */
+        int new_target = (TDIS_CTX_NEWVAL(ctx, 0) - ctx->pc) / 4;
+        unsigned new = bs_set(target, new_target, ctx->op);
+        if (TDIS_CTX_NEWVAL(ctx, 1))
+            new ^= 1 << 24;
+        TDIS_CTX_SET_NEWOP(ctx, new);
+    }
+}
+static INLINE void P(am_brcond_target_B_4_CBNZW)(tdis_ctx ctx, struct bitslice target) {
+    /* both have the same bit to control Z/NZ */
+    return P(am_tbrcond_target_B_4_TBNZW)(ctx, target);
 }
 static INLINE void P(am_ldrlit_label_unk_Rt_6_LDRDl)(tdis_ctx ctx, struct bitslice Rt, struct bitslice label) {
     enum pcrel_load_mode mode;
@@ -38,10 +55,6 @@ static INLINE void P(am_ldrlit_label_unk_Rt_6_LDRDl)(tdis_ctx ctx, struct bitsli
     return P(pcrel)(ctx, ctx->pc + sext(bs_get(label, ctx->op), 19) * 4,
                     bs_get(Rt, ctx->op), mode);
 }
-static INLINE void P(am_tbrcond_target_B_4_TBNZW)(tdis_ctx ctx, struct bitslice target) {
-    return P(branch)(ctx, ctx->pc + sext(bs_get(target, ctx->op), 14) * 4,
-                     /*cond*/ true);
-}
 static INLINE void P(GPR64_Rn_1_RET)(tdis_ctx ctx, UNUSED struct bitslice Rn) {
     return P(ret)(ctx);
 }
diff --git a/lib/arm64/jump-patch.h b/lib/arm64/jump-patch.h
index cc94f90..3d3d653 100644
--- a/lib/arm64/jump-patch.h
+++ b/lib/arm64/jump-patch.h
@@ -14,8 +14,8 @@ static inline int jump_patch_size(uintptr_t pc, uintptr_t dpc,
 }
 
 static inline void make_jump_patch(void **codep, uintptr_t pc, uintptr_t dpc,
-                                   UNUSED struct arch_dis_ctx arch) {
-    int reg = 15;
+                                   struct arch_dis_ctx arch) {
+    int reg = arm64_get_unwritten_temp_reg(&arch);
     intptr_t diff = (dpc & ~0xfff) - (pc & ~0xfff);
     if (!(diff >= -0x100000000 && diff < 0x100000000))
         MOVi64(codep, reg, dpc);
diff --git a/lib/arm64/misc.h b/lib/arm64/misc.h
index 84bd638..c21bc0b 100644
--- a/lib/arm64/misc.h
+++ b/lib/arm64/misc.h
@@ -5,5 +5,32 @@
 #define TARGET_TRANSFORM_DIS_HEADER "arm64/transform-dis-arm64.inc.h"
 #define MIN_INSN_SIZE 4
 #define TD_MAX_REWRITTEN_SIZE (7 * 2 * 4) /* also conservative */
-struct arch_dis_ctx {};
-static inline void arch_dis_ctx_init(UNUSED struct arch_dis_ctx *ctx) {}
+
+struct arch_dis_ctx {
+    /* For transform_dis only - used to get temporary registers.  We assume
+     * that we can use any caller-saved or IP register which was not written,
+     * so r9-r18.
+     * This is a massive overestimate: we just OR in each instruction's bits
+     * 4:0 (Rd for data, Rt for loads, most common), 14:10 (Rt2 for load-pair
+     * instructions), and 20:16 (Rs for store-exclusive insturctions).  It
+     * would be easy to restrict the latter two to the few instructions that
+     * actually use them, but with 10 available registers, and a patch of at
+     * most 3 instructions (and none of the instructions that require a temp
+     * use Rt2/Rs or could read their Rd, so the third doesn't count), we won't
+     * run out even with the dumbest possible thing. */
+    uint32_t regs_possibly_written;
+};
+
+static inline void arch_dis_ctx_init(struct arch_dis_ctx *ctx) {
+    ctx->regs_possibly_written = 0;
+}
+
+static inline int arm64_get_unwritten_temp_reg(struct arch_dis_ctx *ctx) {
+    uint32_t avail = ~ctx->regs_possibly_written & ((1 << 19) - (1 << 9));
+    if (!avail)
+        __builtin_abort();
+    return 31 - __builtin_clz(avail);
+}
+
+#define CC_ARMCC         (CC_CONDITIONAL | 0x200)
+#define CC_XBXZ          (CC_CONDITIONAL | 0x400)
diff --git a/lib/arm64/transform-dis-arm64.inc.h b/lib/arm64/transform-dis-arm64.inc.h
index 97a4aca..a98932d 100644
--- a/lib/arm64/transform-dis-arm64.inc.h
+++ b/lib/arm64/transform-dis-arm64.inc.h
@@ -6,11 +6,9 @@ void transform_dis_pcrel(struct transform_dis_ctx *ctx, uintptr_t dpc, unsigned
     ctx->write_newop_here = NULL;
     void **codep = ctx->rewritten_ptr_ptr;
     if (load_mode >= PLM_U32_SIMD) {
-        /* use x0 as scratch */
-        op32(codep, 0xf81f0fe0); /* str x0, [sp, #-0x10]! */
+        int reg = arm64_get_unwritten_temp_reg(&ctx->arch);
         MOVi64(codep, 0, dpc);
         LDRxi(codep, reg, 0, 0, true, load_mode);
-        op32(codep, 0xf84107e0); /* ldr x0, [sp], #0x10 */
     } else {
         MOVi64(codep, reg, dpc);
         LDRxi(codep, reg, reg, 0, true, load_mode);
@@ -26,9 +24,28 @@ void transform_dis_branch(struct transform_dis_ctx *ctx, uintptr_t dpc, int cc)
         ctx->err = SUBSTITUTE_ERR_FUNC_BAD_INSN_AT_START;
         return;
     }
-    /* TODO */
-    (void) cc;
+    ctx->write_newop_here = NULL;
+    int mov_br_size = size_of_MOVi64(dpc) + 4;
+
+    void **codep = ctx->rewritten_ptr_ptr;
+    if ((cc & CC_ARMCC) == CC_ARMCC) {
+        int icc = (cc & 0xf) ^ 1;
+        Bccrel(codep, icc, 4 + mov_br_size);
+    } else if ((cc & CC_XBXZ) == CC_XBXZ) {
+        ctx->modify = true;
+        ctx->newval[0] = ctx->pc + 4 + mov_br_size;
+        ctx->newval[1] = 1; /* do invert */
+        ctx->write_newop_here = *codep; *codep += 4;
+    }
+    int reg = arm64_get_unwritten_temp_reg(&ctx->arch);
+    MOVi64(codep, reg, dpc);
+    BR(codep, reg);
 }
 
 static void transform_dis_pre_dis(UNUSED struct transform_dis_ctx *ctx) {}
-static void transform_dis_post_dis(UNUSED struct transform_dis_ctx *ctx) {}
+static void transform_dis_post_dis(struct transform_dis_ctx *ctx) {
+    uint32_t op = ctx->op;
+    ctx->arch.regs_possibly_written |= op & 31;
+    ctx->arch.regs_possibly_written |= op >> 10 & 31;
+    ctx->arch.regs_possibly_written |= op >> 16 & 31;
+}