[wasm][RyuJIT] Restore __stack_pointer in reverse P/Invoke epilogs

src/coreclr/jit/codegenwasm.cpp

@@ -474,7 +474,27 @@ void CodeGen::genFnEpilog(BasicBlock* block)
         return;
     }
 
-    // TODO-WASM: shadow stack maintenance
+    // Restore __stack_pointer to the value it had on entry
+    if (m_compiler->opts.IsReversePInvoke())
+    {
+        assert(m_compiler->funCurrentFuncIdx() == ROOT_FUNC_IDX);
+        regNumber fpReg = GetFramePointerReg(ROOT_FUNC_IDX);
+        regNumber spReg = GetStackPointerReg(ROOT_FUNC_IDX);
+        assert(spReg != REG_NA);
+
+        // The FP local is fixed after the prolog; the SP local is only moved by localloc, which
+        // requires a frame pointer.
+        regNumber frameBaseReg = (fpReg != REG_NA) ? fpReg : spReg;
+        GetEmitter()->emitIns_I(INS_local_get, EA_PTRSIZE, WasmRegToIndex(frameBaseReg));
+        if (genTotalFrameSize() != 0)
+        {
+            GetEmitter()->emitIns_I(INS_I_const, EA_PTRSIZE, genTotalFrameSize());
+            GetEmitter()->emitIns(INS_I_add);
+        }
+        GetEmitter()->emitIns_I(INS_global_set, EA_HANDLE_CNS_RELOC,
+                                (cnsval_ssize_t)(size_t)m_compiler->eeGetWasmWellKnownGlobals()->stackPointer);
+    }
+
     // TODO-WASM: we need to handle the end-of-function case if we reach the end of a codegen for a function
     // and do NOT have an epilog. In those cases we currently will not emit an end instruction.
     if (block->IsLast() || m_compiler->bbIsFuncletBeg(block->Next()))

src/tests/readytorun/wasm/WasmInterpreterTransitions/WasmInterpreterTransitions.cs

@@ -122,6 +122,8 @@ public static void TestEntryPoint()
                     delegate* unmanaged<IntPtr, IntPtr, int> fp = &StreamLengthProxy;
                     Assert.Equal(C, fp(IntPtr.Zero, ctx));
                 }
+
+                Assert.Equal(42, R2RCallsNestingUco());
             }
             finally
             {
@@ -510,6 +512,26 @@ private static T GetUserData<T>(IntPtr ptr, out GCHandle handle) where T : class
         return (T)handle.Target!;
     }
 
+    // R2R caller -> UCO -> UCO. Each calli is an inlined PInvoke; the inner one lowers the
+    // __stack_pointer global, and the outer one's epilog asserts that the global is back at the
+    // caller's SP, so this fails unless the UCO epilog restores the global before returning.
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static unsafe int R2RCallsNestingUco()
+    {
+        delegate* unmanaged<int> fp = &UcoWithInlinedPInvoke;
+        return fp();
+    }
+
+    [UnmanagedCallersOnly]
+    private static unsafe int UcoWithInlinedPInvoke()
+    {
+        delegate* unmanaged<int> fp = &UcoLeaf;
+        return fp() + 1;
+    }
+
+    [UnmanagedCallersOnly]
+    private static int UcoLeaf() => 41;
+
     [UnmanagedCallersOnly]
     private static int StreamLengthProxy(IntPtr s, IntPtr context)
     {