4 files changed, 58 insertions, 7 deletions

diff --git a/src/Makefile.dep b/src/Makefile.dep
index 400ef8b0..fda77c83 100644
--- a/src/Makefile.dep
+++ b/src/Makefile.dep
@@ -55,7 +55,7 @@ lj_asm.o: lj_asm.c lj_obj.h lua.h luaconf.h lj_def.h lj_arch.h lj_gc.h \
  lj_buf.h lj_str.h lj_tab.h lj_frame.h lj_bc.h lj_ctype.h lj_ir.h \
  lj_jit.h lj_ircall.h lj_iropt.h lj_mcode.h lj_trace.h lj_dispatch.h \
  lj_traceerr.h lj_snap.h lj_asm.h lj_vm.h lj_target.h lj_target_*.h \
- lj_emit_*.h lj_asm_*.h
+ lj_prng.h lj_emit_*.h lj_asm_*.h
 lj_assert.o: lj_assert.c lj_obj.h lua.h luaconf.h lj_def.h lj_arch.h
 lj_bc.o: lj_bc.c lj_obj.h lua.h luaconf.h lj_def.h lj_arch.h lj_bc.h \
  lj_bcdef.h
diff --git a/src/lj_asm.c b/src/lj_asm.c
index 71079b30..c02a1b9e 100644
--- a/src/lj_asm.c
+++ b/src/lj_asm.c
@@ -29,6 +29,7 @@
 #include "lj_dispatch.h"
 #include "lj_vm.h"
 #include "lj_target.h"
+#include "lj_prng.h"
 
 #ifdef LUA_USE_ASSERT
 #include <stdio.h>
@@ -93,6 +94,12 @@ typedef struct ASMState {
   MCode *flagmcp;       /* Pending opportunity to merge flag setting ins. */
   MCode *realign;       /* Realign loop if not NULL. */
 
+#ifdef LUAJIT_RANDOM_RA
+  /* Randomize register allocation. OK for fuzz testing, not for production. */
+  uint64_t prngbits;
+  PRNGState prngstate;
+#endif
+
 #ifdef RID_NUM_KREF
   intptr_t krefk[RID_NUM_KREF];
 #endif
@@ -173,6 +180,41 @@ IRFLDEF(FLOFS)
   0
 };
 
+#ifdef LUAJIT_RANDOM_RA
+/* Return a fixed number of random bits from the local PRNG state. */
+static uint32_t ra_random_bits(ASMState *as, uint32_t nbits) {
+  uint64_t b = as->prngbits;
+  uint32_t res = (1u << nbits) - 1u;
+  if (b <= res) b = lj_prng_u64(&as->prngstate) | (1ull << 63);
+  res &= (uint32_t)b;
+  as->prngbits = b >> nbits;
+  return res;
+}
+
+/* Pick a random register from a register set. */
+static Reg rset_pickrandom(ASMState *as, RegSet rs)
+{
+  Reg r = rset_pickbot_(rs);
+  rs >>= r;
+  if (rs > 1) {  /* More than one bit set? */
+    while (1) {
+      /* We need to sample max. the GPR or FPR half of the set. */
+      uint32_t d = ra_random_bits(as, RSET_BITS-1);
+      if ((rs >> d) & 1) {
+        r += d;
+        break;
+      }
+    }
+  }
+  return r;
+}
+#define rset_picktop(rs)        rset_pickrandom(as, rs)
+#define rset_pickbot(rs)        rset_pickrandom(as, rs)
+#else
+#define rset_picktop(rs)        rset_picktop_(rs)
+#define rset_pickbot(rs)        rset_pickbot_(rs)
+#endif
+
 /* -- Target-specific instruction emitter --------------------------------- */
 
 #if LJ_TARGET_X86ORX64
@@ -2442,6 +2484,9 @@ void lj_asm_trace(jit_State *J, GCtrace *T)
   as->realign = NULL;
   as->loopinv = 0;
   as->parent = J->parent ? traceref(J, J->parent) : NULL;
+#ifdef LUAJIT_RANDOM_RA
+  (void)lj_prng_u64(&J2G(J)->prng);  /* Ensure PRNG step between traces. */
+#endif
 
   /* Reserve MCode memory. */
   as->mctop = as->mctoporig = lj_mcode_reserve(J, &as->mcbot);
@@ -2483,6 +2528,10 @@ void lj_asm_trace(jit_State *J, GCtrace *T)
 #endif
     as->ir = J->curfinal->ir;  /* Use the copied IR. */
     as->curins = J->cur.nins = as->orignins;
+#ifdef LUAJIT_RANDOM_RA
+    as->prngstate = J2G(J)->prng;  /* Must (re)start from identical state. */
+    as->prngbits = 0;
+#endif
 
     RA_DBG_START();
     RA_DBGX((as, "===== STOP ====="));
diff --git a/src/lj_target.h b/src/lj_target.h
index 2f4d21c1..09d19bd9 100644
--- a/src/lj_target.h
+++ b/src/lj_target.h
@@ -57,8 +57,10 @@ typedef uint32_t RegSP;
 */
 #if LJ_TARGET_PPC || LJ_TARGET_MIPS || LJ_TARGET_ARM64
 typedef uint64_t RegSet;
+#define RSET_BITS               6
 #else
 typedef uint32_t RegSet;
+#define RSET_BITS               5
 #endif
 
 #define RID2RSET(r)             (((RegSet)1) << (r))
@@ -70,11 +72,11 @@ typedef uint32_t RegSet;
 #define rset_clear(rs, r)       (rs &= ~RID2RSET(r))
 #define rset_exclude(rs, r)     (rs & ~RID2RSET(r))
 #if LJ_TARGET_PPC || LJ_TARGET_MIPS || LJ_TARGET_ARM64
-#define rset_picktop(rs)        ((Reg)(__builtin_clzll(rs)^63))
-#define rset_pickbot(rs)        ((Reg)__builtin_ctzll(rs))
+#define rset_picktop_(rs)       ((Reg)(__builtin_clzll(rs)^63))
+#define rset_pickbot_(rs)       ((Reg)__builtin_ctzll(rs))
 #else
-#define rset_picktop(rs)        ((Reg)lj_fls(rs))
-#define rset_pickbot(rs)        ((Reg)lj_ffs(rs))
+#define rset_picktop_(rs)       ((Reg)lj_fls(rs))
+#define rset_pickbot_(rs)       ((Reg)lj_ffs(rs))
 #endif
 
 /* -- Register allocation cost -------------------------------------------- */
diff --git a/src/lj_target_x86.h b/src/lj_target_x86.h
index 7b8d62ad..3482309b 100644
--- a/src/lj_target_x86.h
+++ b/src/lj_target_x86.h
@@ -116,8 +116,8 @@ enum {
 
 #if LJ_64
 /* Prefer the low 8 regs of each type to reduce REX prefixes. */
-#undef rset_picktop
-#define rset_picktop(rs)        (lj_fls(lj_bswap(rs)) ^ 0x18)
+#undef rset_picktop_
+#define rset_picktop_(rs)       (lj_fls(lj_bswap(rs)) ^ 0x18)
 #endif
 
 /* -- Spill slots --------------------------------------------------------- */