Bug 1646035 - Fix return type for Uint32Array loads in Warp. r=evilpie

As long as only int32 values are read, use an int32-only stub. If the Uint32
value doesn't fit in an int32 support loading doubles.

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

js/src/jit/CacheIR.cpp

@@ -2300,6 +2300,24 @@ static Scalar::Type TypedThingElementType(JSObject* obj) {
                                      : PrimitiveArrayTypedObjectType(obj);
 }
 
+// For Uint32Array we let the stub return a double only if the current result is
+// a double, to allow better codegen in Warp.
+static bool AllowDoubleForUint32Array(TypedArrayObject* tarr, uint32_t index) {
+  if (TypedThingElementType(tarr) != Scalar::Type::Uint32) {
+    // Return value is only relevant for Uint32Array.
+    return false;
+  }
+
+  if (index >= tarr->length()) {
+    return false;
+  }
+
+  Value res;
+  MOZ_ALWAYS_TRUE(tarr->getElementPure(index, &res));
+  MOZ_ASSERT(res.isNumber());
+  return res.isDouble();
+}
+
 AttachDecision GetPropIRGenerator::tryAttachTypedElement(
     HandleObject obj, ObjOperandId objId, uint32_t index,
     Int32OperandId indexId) {
@@ -2324,9 +2342,11 @@ AttachDecision GetPropIRGenerator::tryAttachTypedElement(
   // Don't handle out-of-bounds accesses here because we have to ensure the
   // |undefined| type is monitored. See also tryAttachTypedArrayNonInt32Index.
   if (layout == TypedThingLayout::TypedArray) {
-    writer.loadTypedArrayElementResult(objId, indexId,
-                                       TypedThingElementType(obj),
-                                       /* handleOOB = */ false);
+    TypedArrayObject* tarr = &obj->as<TypedArrayObject>();
+    bool allowDoubleForUint32 = AllowDoubleForUint32Array(tarr, index);
+    writer.loadTypedArrayElementResult(
+        objId, indexId, TypedThingElementType(obj),
+        /* handleOOB = */ false, allowDoubleForUint32);
   } else {
     writer.loadTypedObjectElementResult(objId, indexId, layout,
                                         TypedThingElementType(obj));
@@ -2354,13 +2374,27 @@ AttachDecision GetPropIRGenerator::tryAttachTypedArrayNonInt32Index(
     return AttachDecision::NoAction;
   }
 
+  TypedArrayObject* tarr = &obj->as<TypedArrayObject>();
+
+  // Try to convert the number to a typed array index. Use NumberEqualsInt32
+  // because ToPropertyKey(-0) is 0. If the number is not representable as an
+  // int32 the result will be |undefined| so we leave |allowDoubleForUint32| as
+  // false.
+  bool allowDoubleForUint32 = false;
+  int32_t indexInt32;
+  if (mozilla::NumberEqualsInt32(idVal_.toNumber(), &indexInt32)) {
+    uint32_t index = uint32_t(indexInt32);
+    allowDoubleForUint32 = AllowDoubleForUint32Array(tarr, index);
+  }
+
   ValOperandId keyId = getElemKeyValueId();
   Int32OperandId indexId = writer.guardToTypedArrayIndex(keyId);
 
-  writer.guardShapeForClass(objId, obj->as<TypedArrayObject>().shape());
+  writer.guardShapeForClass(objId, tarr->shape());
 
   writer.loadTypedArrayElementResult(objId, indexId, TypedThingElementType(obj),
-                                     /* handleOOB = */ true);
+                                     /* handleOOB = */ true,
+                                     allowDoubleForUint32);
 
   // Always monitor the result when out-of-bounds accesses are expected.
   writer.typeMonitorResult();

js/src/jit/CacheIRCompiler.cpp

@@ -4156,11 +4156,9 @@ static void EmitAllocateBigInt(MacroAssembler& masm, Register result,
   masm.bind(&done);
 }
 
-bool CacheIRCompiler::emitLoadTypedElementResult(ObjOperandId objId,
-                                                 Int32OperandId indexId,
-                                                 TypedThingLayout layout,
-                                                 Scalar::Type elementType,
-                                                 bool handleOOB) {
+bool CacheIRCompiler::emitLoadTypedElementResult(
+    ObjOperandId objId, Int32OperandId indexId, TypedThingLayout layout,
+    Scalar::Type elementType, bool handleOOB, bool allowDoubleForUint32) {
   JitSpew(JitSpew_Codegen, "%s", __FUNCTION__);
   AutoOutputRegister output(*this);
   Register obj = allocator.useRegister(masm, objId);
@@ -4249,8 +4247,9 @@ bool CacheIRCompiler::emitLoadTypedElementResult(ObjOperandId objId,
 
       masm.tagValue(JSVAL_TYPE_BIGINT, *bigInt, output.valueReg());
     } else {
+      bool allowDouble = *allowDoubleResult_ && allowDoubleForUint32;
       masm.loadFromTypedArray(elementType, source, output.valueReg(),
-                              *allowDoubleResult_, scratch1, failure->label());
+                              allowDouble, scratch1, failure->label());
     }
   } else {
     bool needGpr =
@@ -4285,19 +4284,20 @@ bool CacheIRCompiler::emitLoadTypedElementResult(ObjOperandId objId,
   return true;
 }
 
-bool CacheIRCompiler::emitLoadTypedArrayElementResult(ObjOperandId objId,
-                                                      Int32OperandId indexId,
-                                                      Scalar::Type elementType,
-                                                      bool handleOOB) {
-  return emitLoadTypedElementResult(
-      objId, indexId, TypedThingLayout::TypedArray, elementType, handleOOB);
+bool CacheIRCompiler::emitLoadTypedArrayElementResult(
+    ObjOperandId objId, Int32OperandId indexId, Scalar::Type elementType,
+    bool handleOOB, bool allowDoubleForUint32) {
+  return emitLoadTypedElementResult(objId, indexId,
+                                    TypedThingLayout::TypedArray, elementType,
+                                    handleOOB, allowDoubleForUint32);
 }
 
 bool CacheIRCompiler::emitLoadTypedObjectElementResult(
     ObjOperandId objId, Int32OperandId indexId, TypedThingLayout layout,
     Scalar::Type elementType) {
   return emitLoadTypedElementResult(objId, indexId, layout, elementType,
-                                    /* handleOOB = */ false);
+                                    /* handleOOB = */ false,
+                                    /* allowDoubleForUint32 = */ true);
 }
 
 bool CacheIRCompiler::emitStoreTypedObjectScalarProperty(

js/src/jit/CacheIRCompiler.h

@@ -756,7 +756,8 @@ class MOZ_RAII CacheIRCompiler {
 
   bool emitLoadTypedElementResult(ObjOperandId objId, Int32OperandId indexId,
                                   TypedThingLayout layout,
-                                  Scalar::Type elementType, bool handleOOB);
+                                  Scalar::Type elementType, bool handleOOB,
+                                  bool allowDoubleForUint32);
 
   bool emitStoreTypedElement(ObjOperandId objId, TypedThingLayout layout,
                              Scalar::Type elementType, Int32OperandId indexId,

js/src/jit/CacheIROps.yaml

@@ -1201,6 +1201,7 @@
     index: Int32Id
     elementType: ScalarTypeImm
     handleOOB: BoolImm
+    allowDoubleForUint32: BoolImm
 
 - name: LoadTypedObjectElementResult
   shared: true

js/src/jit/WarpCacheIRTranspiler.cpp

@@ -756,14 +756,13 @@ bool WarpCacheIRTranspiler::emitLoadDenseElementHoleResult(
 
 bool WarpCacheIRTranspiler::emitLoadTypedArrayElementResult(
     ObjOperandId objId, Int32OperandId indexId, Scalar::Type elementType,
-    bool handleOOB) {
+    bool handleOOB, bool allowDoubleForUint32) {
   MDefinition* obj = getOperand(objId);
   MDefinition* index = getOperand(indexId);
 
   if (handleOOB) {
-    bool allowDouble = true;
-    auto* load = MLoadTypedArrayElementHole::New(alloc(), obj, index,
-                                                 elementType, allowDouble);
+    auto* load = MLoadTypedArrayElementHole::New(
+        alloc(), obj, index, elementType, allowDoubleForUint32);
     add(load);
 
     pushResult(load);
@@ -779,8 +778,8 @@ bool WarpCacheIRTranspiler::emitLoadTypedArrayElementResult(
   add(elements);
 
   auto* load = MLoadUnboxedScalar::New(alloc(), elements, index, elementType);
-  // TODO: Uint32 always loaded as double.
-  load->setResultType(MIRTypeForArrayBufferViewRead(elementType, true));
+  load->setResultType(
+      MIRTypeForArrayBufferViewRead(elementType, allowDoubleForUint32));
   add(load);
 
   pushResult(load);