Bug 1871951: Assign stack arg padding to callee's stackmap. r=jseward

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

js/src/jit/CodeGenerator.cpp

@@ -8991,7 +8991,8 @@ void CodeGenerator::visitWasmCall(LWasmCall* lir) {
   // Now that all the outbound in-memory args are on the stack, note the
   // required lower boundary point of the associated StackMap.
   uint32_t framePushedAtStackMapBase =
-      masm.framePushed() - callBase->stackArgAreaSizeUnaligned();
+      masm.framePushed() -
+      wasm::AlignStackArgAreaSize(callBase->stackArgAreaSizeUnaligned());
   lir->safepoint()->setFramePushedAtStackMapBase(framePushedAtStackMapBase);
   MOZ_ASSERT(lir->safepoint()->wasmSafepointKind() ==
              WasmSafepointKind::LirCall);
@@ -14697,10 +14698,14 @@ static bool CreateStackMapFromLSafepoint(LSafepoint& safepoint,
   // This is the number of bytes in the general spill area, below the Frame.
   const size_t nBodyBytes = safepoint.framePushedAtStackMapBase();
 
+  // The stack map owns any alignment padding around inbound stack args.
+  const size_t nInboundStackArgBytesAligned =
+      wasm::AlignStackArgAreaSize(nInboundStackArgBytes);
+
   // This is the number of bytes in the general spill area, the Frame, and the
   // incoming args, but not including any register dump area.
   const size_t nNonRegisterBytes =
-      nBodyBytes + nFrameBytes + nInboundStackArgBytes;
+      nBodyBytes + nFrameBytes + nInboundStackArgBytesAligned;
   MOZ_ASSERT(nNonRegisterBytes % sizeof(void*) == 0);
 
   // This is the number of bytes in the register dump area, if any, below the
@@ -14826,7 +14831,7 @@ static bool CreateStackMapFromLSafepoint(LSafepoint& safepoint,
   // Record in the map, how far down from the highest address the Frame* is.
   // Take the opportunity to check that we haven't marked any part of the
   // Frame itself as a pointer.
-  stackMap->setFrameOffsetFromTop((nInboundStackArgBytes + nFrameBytes) /
+  stackMap->setFrameOffsetFromTop((nInboundStackArgBytesAligned + nFrameBytes) /
                                   sizeof(void*));
 #ifdef DEBUG
   for (uint32_t i = 0; i < nFrameBytes / sizeof(void*); i++) {

js/src/wasm/WasmBCFrame.cpp

@@ -224,18 +224,16 @@ bool StackMapGenerator::createStackMap(
   // function's prologue.  We now need to calculate how far the machine's
   // stack pointer is below where it was at the start of the body.  But we
   // must take care not to include any words pushed as arguments to an
-  // upcoming function call, since those words "belong" to the stackmap of
-  // the callee, not to the stackmap of this function.  Note however that
-  // any alignment padding pushed prior to pushing the args *does* belong to
-  // this function.
+  // upcoming function call, since those words belong to the stackmap of
+  // the callee, not to the stackmap of this function.  Any alignment padding
+  // for the args also belongs to the callee.
   //
-  // That padding is taken into account at the point where
-  // framePushedExcludingOutboundCallArgs is set, viz, in startCallArgs(),
-  // and comprises two components:
+  // The only padding belonging to the stackmap of this function is that
+  // required to align the upcoming frame. This is accounted for where
+  // framePushedExcludingOutboundCallArgs is set, in startCallArgs(), and is
+  // comprised of just one component:
   //
   // * call->frameAlignAdjustment
-  // * the padding applied to the stack arg area itself.  That is:
-  //   StackArgAreaSize(argTys) - StackArgAreaSizeUnpadded(argTys)
   Maybe<uint32_t> framePushedExcludingArgs;
   if (framePushedAtEntryToBody.isNothing()) {
     // Still in the prologue.  framePushedExcludingArgs remains Nothing.
@@ -355,11 +353,19 @@ bool StackMapGenerator::createStackMap(
     augmentedMst.setGCPointer(offsFromMapLowest / sizeof(void*));
   }
 
+  MOZ_ASSERT(numStackArgBytes % sizeof(void*) == 0);
+  const size_t numStackArgWords = numStackArgBytes / sizeof(void*);
+  const size_t numStackArgPaddingBytes =
+      AlignStackArgAreaSize(numStackArgBytes) - numStackArgBytes;
+  const size_t numStackArgPaddingWords =
+      numStackArgPaddingBytes / sizeof(void*);
+
   // Create the final StackMap.  The initial map is zeroed out, so there's
   // no need to write zero bits in it.
   const uint32_t extraWords = extras.length();
   const uint32_t augmentedMstWords = augmentedMst.length();
-  const uint32_t numMappedWords = extraWords + augmentedMstWords;
+  const uint32_t numMappedWords =
+      numStackArgPaddingWords + extraWords + augmentedMstWords;
   StackMap* stackMap = StackMap::create(numMappedWords);
   if (!stackMap) {
     return false;
@@ -398,7 +404,7 @@ bool StackMapGenerator::createStackMap(
   // Record in the map, how far down from the highest address the Frame* is.
   // Take the opportunity to check that we haven't marked any part of the
   // Frame itself as a pointer.
-  stackMap->setFrameOffsetFromTop(numStackArgWords +
+  stackMap->setFrameOffsetFromTop(numStackArgPaddingWords + numStackArgWords +
                                   sizeof(Frame) / sizeof(void*));
 #ifdef DEBUG
   for (uint32_t i = 0; i < sizeof(Frame) / sizeof(void*); i++) {

js/src/wasm/WasmBCFrame.h

@@ -1318,8 +1318,8 @@ struct StackMapGenerator {
  public:
   // --- These are constant once we've completed beginFunction() ---
 
-  // The number of words of arguments passed to this function in memory.
-  size_t numStackArgWords;
+  // The number of bytes of arguments passed to this function in memory.
+  size_t numStackArgBytes;
 
   MachineStackTracker machineStackTracker;  // tracks machine stack pointerness
 
@@ -1336,14 +1336,14 @@ struct StackMapGenerator {
   // memory.  That is, for an upcoming function call, this will hold
   //
   //   masm.framePushed() at the call instruction -
-  //      StackArgAreaSizeUnaligned(argumentTypes)
+  //      StackArgAreaSizeAligned(argumentTypes)
   //
   // This value denotes the lowest-addressed stack word covered by the current
   // function's stackmap.  Words below this point form the highest-addressed
   // area of the callee's stackmap.  Note that all alignment padding above the
-  // arguments-in-memory themselves belongs to the caller's stackmap, which
-  // is why this is defined in terms of StackArgAreaSizeUnaligned() rather than
-  // StackArgAreaSizeAligned().
+  // arguments-in-memory themselves belongs to the callee's stackmap, as return
+  // calls will replace the function arguments with a new set of arguments which
+  // may have different alignment.
   //
   // When not inside a function call setup/teardown sequence, it is Nothing.
   // It can make Nothing-to/from-Some transitions arbitrarily as we progress
@@ -1366,7 +1366,7 @@ struct StackMapGenerator {
         trapExitLayoutNumWords_(trapExitLayoutNumWords),
         stackMaps_(stackMaps),
         masm_(masm),
-        numStackArgWords(0),
+        numStackArgBytes(0),
         memRefsOnStk(0) {}
 
   // At the beginning of a function, we may have live roots in registers (as

js/src/wasm/WasmBaselineCompile.cpp

@@ -381,11 +381,11 @@ bool BaseCompiler::beginFunction() {
   ArgTypeVector args(funcType());
   size_t inboundStackArgBytes = StackArgAreaSizeUnaligned(args);
   MOZ_ASSERT(inboundStackArgBytes % sizeof(void*) == 0);
-  stackMapGenerator_.numStackArgWords = inboundStackArgBytes / sizeof(void*);
+  stackMapGenerator_.numStackArgBytes = inboundStackArgBytes;
 
   MOZ_ASSERT(stackMapGenerator_.machineStackTracker.length() == 0);
   if (!stackMapGenerator_.machineStackTracker.pushNonGCPointers(
-          stackMapGenerator_.numStackArgWords)) {
+          stackMapGenerator_.numStackArgBytes / sizeof(void*))) {
     return false;
   }
 
@@ -1389,18 +1389,15 @@ void BaseCompiler::startCallArgs(size_t stackArgAreaSizeUnaligned,
   MOZ_ASSERT(stackArgAreaSizeUnaligned <= stackArgAreaSizeAligned);
 
   // Record the masm.framePushed() value at this point, before we push args
-  // for the call, but including the alignment space placed above the args.
-  // This defines the lower limit of the stackmap that will be created for
-  // this call.
+  // for the call and any required alignment space. This defines the lower limit
+  // of the stackmap that will be created for this call.
   MOZ_ASSERT(
       stackMapGenerator_.framePushedExcludingOutboundCallArgs.isNothing());
   stackMapGenerator_.framePushedExcludingOutboundCallArgs.emplace(
       // However much we've pushed so far
       masm.framePushed() +
       // Extra space we'll push to get the frame aligned
-      call->frameAlignAdjustment +
-      // Extra space we'll push to get the outbound arg area aligned
-      (stackArgAreaSizeAligned - stackArgAreaSizeUnaligned));
+      call->frameAlignAdjustment);
 
   call->stackArgAreaSize = stackArgAreaSizeAligned;
 

js/src/wasm/WasmGC.cpp

@@ -79,10 +79,14 @@ bool wasm::CreateStackMapForFunctionEntryTrap(
   // The size of the register dump (trap) area.
   const size_t trapExitLayoutBytes = trapExitLayoutWords * sizeof(void*);
 
+  // The stack map owns any alignment padding for incoming stack args.
+  const size_t nInboundStackArgBytesAligned =
+      AlignStackArgAreaSize(nInboundStackArgBytes);
+
   // This is the total number of bytes covered by the map.
   const DebugOnly<size_t> nTotalBytes = trapExitLayoutBytes +
                                         nBytesReservedBeforeTrap + nFrameBytes +
-                                        nInboundStackArgBytes;
+                                        nInboundStackArgBytesAligned;
 
   // Create the stackmap initially in this vector.  Since most frames will
   // contain 128 or fewer words, heap allocation is avoided in the majority of
@@ -121,8 +125,8 @@ bool wasm::CreateStackMapForFunctionEntryTrap(
   }
 
   // INBOUND ARG AREA
-  MOZ_ASSERT(nInboundStackArgBytes % sizeof(void*) == 0);
-  const size_t numStackArgWords = nInboundStackArgBytes / sizeof(void*);
+  MOZ_ASSERT(nInboundStackArgBytesAligned % sizeof(void*) == 0);
+  const size_t numStackArgWords = nInboundStackArgBytesAligned / sizeof(void*);
 
   const size_t wordsSoFar = vec.length();
   if (!vec.appendN(false, numStackArgWords)) {