1 files changed, 42 insertions, 3 deletions

diff --git a/js/src/wasm/WasmStubs.cpp b/js/src/wasm/WasmStubs.cpp
index dfaa898744..76f015d34b 100644
--- a/js/src/wasm/WasmStubs.cpp
+++ b/js/src/wasm/WasmStubs.cpp
@@ -1862,6 +1862,39 @@ static void FillArgumentArrayForJitExit(MacroAssembler& masm, Register instance,
   GenPrintf(DebugChannel::Import, masm, "\n");
 }
 
+static bool AddStackCheckForImportFunctionEntry(jit::MacroAssembler& masm,
+                                                unsigned reserve,
+                                                const FuncType& funcType,
+                                                StackMaps* stackMaps) {
+  std::pair<CodeOffset, uint32_t> pair =
+      masm.wasmReserveStackChecked(reserve, BytecodeOffset(0));
+
+  // Attempt to create stack maps for masm.wasmReserveStackChecked.
+  ArgTypeVector argTypes(funcType);
+  RegisterOffsets trapExitLayout;
+  size_t trapExitLayoutNumWords;
+  GenerateTrapExitRegisterOffsets(&trapExitLayout, &trapExitLayoutNumWords);
+  CodeOffset trapInsnOffset = pair.first;
+  size_t nBytesReservedBeforeTrap = pair.second;
+  size_t nInboundStackArgBytes = StackArgAreaSizeUnaligned(argTypes);
+  wasm::StackMap* stackMap = nullptr;
+  if (!CreateStackMapForFunctionEntryTrap(
+          argTypes, trapExitLayout, trapExitLayoutNumWords,
+          nBytesReservedBeforeTrap, nInboundStackArgBytes, &stackMap)) {
+    return false;
+  }
+
+  // In debug builds, we'll always have a stack map, even if there are no
+  // refs to track.
+  MOZ_ASSERT(stackMap);
+  if (stackMap &&
+      !stackMaps->add((uint8_t*)(uintptr_t)trapInsnOffset.offset(), stackMap)) {
+    stackMap->destroy();
+    return false;
+  }
+  return true;
+}
+
 // Generate a wrapper function with the standard intra-wasm call ABI which
 // simply calls an import. This wrapper function allows any import to be treated
 // like a normal wasm function for the purposes of exports and table calls. In
@@ -1873,7 +1906,7 @@ static bool GenerateImportFunction(jit::MacroAssembler& masm,
                                    const FuncImport& fi,
                                    const FuncType& funcType,
                                    CallIndirectId callIndirectId,
-                                   FuncOffsets* offsets) {
+                                   FuncOffsets* offsets, StackMaps* stackMaps) {
   AutoCreatedBy acb(masm, "wasm::GenerateImportFunction");
 
   AssertExpectedSP(masm);
@@ -1886,7 +1919,12 @@ static bool GenerateImportFunction(jit::MacroAssembler& masm,
       WasmStackAlignment,
       sizeof(Frame),  // pushed by prologue
       StackArgBytesForWasmABI(funcType) + sizeOfInstanceSlot);
-  masm.wasmReserveStackChecked(framePushed, BytecodeOffset(0));
+
+  if (!AddStackCheckForImportFunctionEntry(masm, framePushed, funcType,
+                                           stackMaps)) {
+    return false;
+  }
+
   MOZ_ASSERT(masm.framePushed() == framePushed);
 
   masm.storePtr(InstanceReg, Address(masm.getStackPointer(),
@@ -1950,7 +1988,8 @@ bool wasm::GenerateImportFunctions(const ModuleEnvironment& env,
     CallIndirectId callIndirectId = CallIndirectId::forFunc(env, funcIndex);
 
     FuncOffsets offsets;
-    if (!GenerateImportFunction(masm, fi, funcType, callIndirectId, &offsets)) {
+    if (!GenerateImportFunction(masm, fi, funcType, callIndirectId, &offsets,
+                                &code->stackMaps)) {
       return false;
     }
     if (!code->codeRanges.emplaceBack(funcIndex, /* bytecodeOffset = */ 0,