Bug 1639153 - Part 3: Implement the algorithm for obtaining tls and use it for wasm signal handling. r=lth

This is the third part of series of patches to Frame without tls pointer.
Here we preserve initial tls in all entry stubs and then use it to find a proper tls instance for a given frame.

To find the TlsData* for specific frame we start from a entry stub's tls
and then track tls through all possible cross-instance calls. This logic
is implemented in GetNearestEffectiveTls procedure.

Then, we use this new procedure to make singal handling free from Frame::tls.

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

Depends on D82888

js/src/jit/MacroAssembler.cpp

@@ -3796,6 +3796,7 @@ CodeOffset MacroAssembler::wasmCallBuiltinInstanceMethod(
            Address(getStackPointer(), WasmCallerTLSOffsetBeforeCall));
   storePtr(WasmTlsReg,
            Address(getStackPointer(), WasmCalleeTLSOffsetBeforeCall));
+
   if (instanceArg.kind() == ABIArg::GPR) {
     loadPtr(Address(WasmTlsReg, offsetof(wasm::TlsData, instance)),
             instanceArg.gpr());
@@ -4656,6 +4657,18 @@ void AutoGenericRegisterScope<RegisterType>::reacquire() {
 template void AutoGenericRegisterScope<Register>::reacquire();
 template void AutoGenericRegisterScope<FloatRegister>::reacquire();
 
+wasm::TlsData* ExtractCallerTlsFromFrameWithTls(wasm::Frame* fp) {
+  return *reinterpret_cast<wasm::TlsData**>(
+      reinterpret_cast<uint8_t*>(fp) + sizeof(wasm::Frame) + ShadowStackSpace +
+      wasm::FrameWithTls::callerTLSOffset());
+}
+
+wasm::TlsData* ExtractCalleeTlsFromFrameWithTls(wasm::Frame* fp) {
+  return *reinterpret_cast<wasm::TlsData**>(
+      reinterpret_cast<uint8_t*>(fp) + sizeof(wasm::Frame) + ShadowStackSpace +
+      wasm::FrameWithTls::calleeTLSOffset());
+}
+
 #endif  // DEBUG
 
 }  // namespace jit

js/src/jit/MacroAssembler.h

@@ -4411,6 +4411,9 @@ class WasmABIArgIter : public ABIArgIterBase<VecT, WasmABIArgGenerator> {
       : ABIArgIterBase<VecT, WasmABIArgGenerator>(types) {}
 };
 
+wasm::TlsData* ExtractCalleeTlsFromFrameWithTls(wasm::Frame* fp);
+wasm::TlsData* ExtractCallerTlsFromFrameWithTls(wasm::Frame* fp);
+
 }  // namespace jit
 }  // namespace js
 

js/src/wasm/WasmFrameIter.cpp

@@ -920,6 +920,34 @@ static bool isSignatureCheckFail(uint32_t offsetInCode,
          (offsetInCode - codeRange->funcCheckedCallEntry()) > SetFP;
 }
 
+TlsData* js::wasm::GetNearestEffectiveTls(Frame* fp) {
+  while (true) {
+    if (fp->callerIsExitOrJitEntryFP()) {
+      // It is a direct call from JIT.
+      MOZ_ASSERT(!LookupCode(fp->returnAddress()));
+      return ExtractCalleeTlsFromFrameWithTls(fp);
+    }
+
+    uint8_t* returnAddress = fp->returnAddress();
+    const CodeRange* codeRange = nullptr;
+    const Code* code = LookupCode(returnAddress, &codeRange);
+    MOZ_ASSERT(codeRange);
+
+    if (codeRange->isEntry()) {
+      return ExtractCalleeTlsFromFrameWithTls(fp);
+    }
+
+    MOZ_ASSERT(codeRange->kind() == CodeRange::Function);
+    MOZ_ASSERT(code);
+    const CallSite* callsite = code->lookupCallSite(returnAddress);
+    if (callsite->mightBeCrossInstance()) {
+      return ExtractCalleeTlsFromFrameWithTls(fp);
+    }
+
+    fp = fp->wasmCaller();
+  }
+}
+
 bool js::wasm::StartUnwinding(const RegisterState& registers,
                               UnwindState* unwindState, bool* unwoundCaller) {
   // Shorthands.

js/src/wasm/WasmFrameIter.h

@@ -233,6 +233,10 @@ void GenerateFunctionPrologue(jit::MacroAssembler& masm,
 void GenerateFunctionEpilogue(jit::MacroAssembler& masm, unsigned framePushed,
                               FuncOffsets* offsets);
 
+// Iterates through frames for either possible cross-instance call or an entry
+// stub to obtain tls that corresponds to the passed fp.
+TlsData* GetNearestEffectiveTls(Frame* fp);
+
 // Describes register state and associated code at a given call frame.
 
 struct UnwindState {

js/src/wasm/WasmSignalHandlers.cpp

@@ -717,7 +717,9 @@ static MOZ_MUST_USE bool HandleTrap(CONTEXT* context,
   // due to this trap occurring in the indirect call prologue, while fp points
   // to the caller's Frame which can be in a different Module. In any case,
   // though, the containing JSContext is the same.
-  Instance* instance = ((Frame*)ContextToFP(context))->instance();
+
+  auto* frame = reinterpret_cast<Frame*>(ContextToFP(context));
+  Instance* instance = GetNearestEffectiveTls(frame)->instance;
   MOZ_RELEASE_ASSERT(&instance->code() == &segment.code() ||
                      trap == Trap::IndirectCallBadSig);
 

js/src/wasm/WasmStubs.cpp

@@ -832,6 +832,9 @@ static bool GenerateInterpEntry(MacroAssembler& masm, const FuncExport& fe,
   masm.movePtr(ImmWord(0), FramePointer);
   masm.loadWasmPinnedRegsFromTls();
 
+  masm.storePtr(WasmTlsReg,
+                Address(masm.getStackPointer(), WasmCalleeTLSOffsetBeforeCall));
+
   // Call into the real function. Note that, due to the throw stub, fp, tls
   // and pinned registers may be clobbered.
   masm.assertStackAlignment(WasmStackAlignment);
@@ -1290,6 +1293,9 @@ static bool GenerateJitEntry(MacroAssembler& masm, size_t funcExportIndex,
   // Setup wasm register state.
   masm.loadWasmPinnedRegsFromTls();
 
+  masm.storePtr(WasmTlsReg,
+                Address(masm.getStackPointer(), WasmCalleeTLSOffsetBeforeCall));
+
   // Call into the real function. Note that, due to the throw stub, fp, tls
   // and pinned registers may be clobbered.
   masm.assertStackAlignment(WasmStackAlignment);
@@ -1604,6 +1610,8 @@ void wasm::GenerateDirectCallFromJit(MacroAssembler& masm, const FuncExport& fe,
 
   // Load tls; from now on, WasmTlsReg is live.
   masm.movePtr(ImmPtr(inst.tlsData()), WasmTlsReg);
+  masm.storePtr(WasmTlsReg,
+                Address(masm.getStackPointer(), WasmCalleeTLSOffsetBeforeCall));
   masm.loadWasmPinnedRegsFromTls();
 
   // Actual call.

js/src/wasm/WasmTypes.h

@@ -2585,6 +2585,7 @@ class CallSiteDesc {
   }
   uint32_t lineOrBytecode() const { return lineOrBytecode_; }
   Kind kind() const { return Kind(kind_); }
+  bool mightBeCrossInstance() const { return kind() == CallSiteDesc::Dynamic; }
 };
 
 class CallSite : public CallSiteDesc {
@@ -3304,10 +3305,13 @@ class Frame {
 static_assert(!std::is_polymorphic_v<Frame>, "Frame doesn't need a vtable.");
 
 class FrameWithTls : public Frame {
- public:
   TlsData* calleeTls_;
   TlsData* callerTls_;
 
+ public:
+  TlsData* calleeTls() { return calleeTls_; }
+  TlsData* callerTls() { return callerTls_; }
+
   constexpr static uint32_t sizeWithoutFrame() {
     return sizeof(wasm::FrameWithTls) - sizeof(wasm::Frame);
   }