Bug 1642610 part 2 - Optimize Math.random in CacheIR. r=iain

Note that the math-random.js jit-test is a good correctness test for this. It
overrides the RNG's seed, generates a few thousand numbers, and then checks some
of them. I verified this test fails if I add a bug to the Math.random JIT code.

Differential Revision: https://phabricator.services.mozilla.com/D77801

js/src/jit/BaselineCacheIRCompiler.cpp

@@ -1529,6 +1529,24 @@ bool BaselineCacheIRCompiler::emitStringFromCharCodeResult(
   return true;
 }
 
+bool BaselineCacheIRCompiler::emitMathRandomResult(uint32_t rngOffset) {
+  JitSpew(JitSpew_Codegen, "%s", __FUNCTION__);
+
+  AutoOutputRegister output(*this);
+  AutoScratchRegister scratch1(allocator, masm);
+  AutoScratchRegister64 scratch2(allocator, masm);
+  AutoAvailableFloatRegister scratchFloat(*this, FloatReg0);
+
+  Address rngAddr(stubAddress(rngOffset));
+  masm.loadPtr(rngAddr, scratch1);
+
+  masm.randomDouble(scratch1, scratchFloat, scratch2,
+                    output.valueReg().toRegister64());
+
+  masm.boxDouble(scratchFloat, output.valueReg(), scratchFloat);
+  return true;
+}
+
 bool BaselineCacheIRCompiler::emitCallNativeSetter(ObjOperandId objId,
                                                    uint32_t setterOffset,
                                                    ValOperandId rhsId) {

js/src/jit/CacheIR.cpp

@@ -5304,6 +5304,32 @@ AttachDecision CallIRGenerator::tryAttachStringFromCharCode(
   return AttachDecision::Attach;
 }
 
+AttachDecision CallIRGenerator::tryAttachMathRandom(HandleFunction callee) {
+  // Expecting no arguments.
+  if (argc_ != 0) {
+    return AttachDecision::NoAction;
+  }
+
+  MOZ_ASSERT(cx_->realm() == callee->realm(),
+             "Shouldn't inline cross-realm Math.random because per-realm RNG");
+
+  // Initialize the input operand.
+  Int32OperandId argcId(writer.setInputOperandId(0));
+
+  // Guard callee is the 'random' native function.
+  emitNativeCalleeGuard(callee);
+
+  mozilla::non_crypto::XorShift128PlusRNG* rng =
+      &cx_->realm()->getOrCreateRandomNumberGenerator();
+  writer.mathRandomResult(rng);
+
+  writer.returnFromIC();
+  cacheIRStubKind_ = BaselineCacheIRStubKind::Regular;
+
+  trackAttached("MathRandom");
+  return AttachDecision::Attach;
+}
+
 AttachDecision CallIRGenerator::tryAttachMathAbs(HandleFunction callee) {
   // Need one argument.
   if (argc_ != 1) {
@@ -5777,6 +5803,8 @@ AttachDecision CallIRGenerator::tryAttachInlinableNative(
       return tryAttachStringFromCharCode(callee);
 
     // Math natives.
+    case InlinableNative::MathRandom:
+      return tryAttachMathRandom(callee);
     case InlinableNative::MathAbs:
       return tryAttachMathAbs(callee);
     case InlinableNative::MathFloor:

js/src/jit/CacheIR.h

@@ -1535,6 +1535,7 @@ class MOZ_RAII CallIRGenerator : public IRGenerator {
   AttachDecision tryAttachStringCharCodeAt(HandleFunction callee);
   AttachDecision tryAttachStringCharAt(HandleFunction callee);
   AttachDecision tryAttachStringFromCharCode(HandleFunction callee);
+  AttachDecision tryAttachMathRandom(HandleFunction callee);
   AttachDecision tryAttachMathAbs(HandleFunction callee);
   AttachDecision tryAttachMathFloor(HandleFunction callee);
   AttachDecision tryAttachMathCeil(HandleFunction callee);

js/src/jit/CacheIRCompiler.h

@@ -556,6 +556,33 @@ class MOZ_RAII AutoSpectreBoundsScratchRegister {
   operator Register() const { return reg_; }
 };
 
+// Scratch Register64. Implemented with a single AutoScratchRegister on 64-bit
+// platforms and two AutoScratchRegisters on 32-bit platforms.
+class MOZ_RAII AutoScratchRegister64 {
+  AutoScratchRegister reg1_;
+#if JS_BITS_PER_WORD == 32
+  AutoScratchRegister reg2_;
+#endif
+
+ public:
+  AutoScratchRegister64(const AutoScratchRegister64&) = delete;
+  void operator=(const AutoScratchRegister64&) = delete;
+
+#if JS_BITS_PER_WORD == 32
+  AutoScratchRegister64(CacheRegisterAllocator& alloc, MacroAssembler& masm)
+      : reg1_(alloc, masm), reg2_(alloc, masm) {}
+
+  Register64 get() const { return Register64(reg1_, reg2_); }
+#else
+  AutoScratchRegister64(CacheRegisterAllocator& alloc, MacroAssembler& masm)
+      : reg1_(alloc, masm) {}
+
+  Register64 get() const { return Register64(reg1_); }
+#endif
+
+  operator Register64() const { return get(); }
+};
+
 // The FailurePath class stores everything we need to generate a failure path
 // at the end of the IC code. The failure path restores the input registers, if
 // needed, and jumps to the next stub.

js/src/jit/CacheIROps.yaml

@@ -715,6 +715,14 @@
   args:
     input: NumberId
 
+# Because Baseline stub code is shared by all realms in the Zone, this
+# instruction loads a pointer to the RNG from a stub field.
+- name: MathRandomResult
+  shared: false
+  transpile: false
+  args:
+    rng: RawPointerField
+
 - name: MathFloorToInt32Result
   shared: true
   transpile: true

js/src/jit/IonCacheIRCompiler.cpp

@@ -2426,3 +2426,7 @@ bool IonCacheIRCompiler::emitIsArrayResult(ValOperandId inputId) {
 bool IonCacheIRCompiler::emitStringFromCharCodeResult(Int32OperandId codeId) {
   MOZ_CRASH("Call ICs not used in ion");
 }
+
+bool IonCacheIRCompiler::emitMathRandomResult(uint32_t rngOffset) {
+  MOZ_CRASH("Call ICs not used in ion");
+}

js/src/jit/RegisterSets.h

@@ -123,6 +123,9 @@ class ValueOperand {
 
   constexpr Register typeReg() const { return type_; }
   constexpr Register payloadReg() const { return payload_; }
+  constexpr Register64 toRegister64() const {
+    return Register64(typeReg(), payloadReg());
+  }
   constexpr bool aliases(Register reg) const {
     return type_ == reg || payload_ == reg;
   }
@@ -141,6 +144,7 @@ class ValueOperand {
   explicit constexpr ValueOperand(Register value) : value_(value) {}
 
   constexpr Register valueReg() const { return value_; }
+  constexpr Register64 toRegister64() const { return Register64(valueReg()); }
   constexpr bool aliases(Register reg) const { return value_ == reg; }
   constexpr Register payloadOrValueReg() const { return valueReg(); }
   constexpr bool operator==(const ValueOperand& o) const {