Enhance MacroAssembler::probe() to support an initializeStackFunction callback.
authormark.lam@apple.com <mark.lam@apple.com@268f45cc-cd09-0410-ab3c-d52691b4dbfc>
Wed, 16 Aug 2017 20:38:57 +0000 (20:38 +0000)
committermark.lam@apple.com <mark.lam@apple.com@268f45cc-cd09-0410-ab3c-d52691b4dbfc>
Wed, 16 Aug 2017 20:38:57 +0000 (20:38 +0000)
https://bugs.webkit.org/show_bug.cgi?id=175617
<rdar://problem/33912104>

Reviewed by JF Bastien.

This patch adds a new feature to MacroAssembler::probe() where the probe function
can provide a ProbeFunction callback to fill in stack values after the stack
pointer has been adjusted.  The probe function can use this feature as follows:

1. Set the new sp value in the ProbeContext's CPUState.

2. Set the ProbeContext's initializeStackFunction to a ProbeFunction callback
   which will do the work of filling in the stack values after the probe
   trampoline has adjusted the machine stack pointer.

3. Set the ProbeContext's initializeStackArgs to any value that the client wants
   to pass to the initializeStackFunction callback.

4. Return from the probe function.

Upon returning from the probe function, the probe trampoline will adjust the
the stack pointer based on the sp value in CPUState.  If initializeStackFunction
is not set, the probe trampoline will restore registers and return to its caller.

If initializeStackFunction is set, the trampoline will move the ProbeContext
beyond the range of the stack pointer i.e. it will place the new ProbeContext at
an address lower than where CPUState.sp() points.  This ensures that the
ProbeContext will not be trashed by the initializeStackFunction when it writes to
the stack.  Then, the trampoline will call back to the initializeStackFunction
ProbeFunction to let it fill in the stack values as desired.  The
initializeStackFunction ProbeFunction will be passed the moved ProbeContext at
the new location.

initializeStackFunction may now write to the stack at addresses greater or
equal to CPUState.sp(), but not below that.  initializeStackFunction is also
not allowed to change CPUState.sp().  If the initializeStackFunction does not
abide by these rules, then behavior is undefined, and bad things may happen.

For future reference, some implementation details that this patch needed to
be mindful of:

1. When the probe trampoline allocates stack space for the ProbeContext, it
   should include OUT_SIZE as well.  This ensures that it doesn't have to move
   the ProbeContext on exit if the probe function didn't change the sp.

2. If the trampoline has to move the ProbeContext, it needs to point the machine
   sp to new ProbeContext first before copying over the ProbeContext data.  This
   protects the new ProbeContext from possibly being trashed by interrupts.

3. When computing the new address of ProbeContext to move to, we need to make
   sure that it is properly aligned in accordance with stack ABI requirements
   (just like we did when we allocated the ProbeContext on entry to the
   probe trampoline).

4. When copying the ProbeContext to its new location, the trampoline should
   always copy words from low addresses to high addresses.  This is because if
   we're moving the ProbeContext, we'll always be moving it to a lower address.

* assembler/MacroAssembler.h:
* assembler/MacroAssemblerARM.cpp:
* assembler/MacroAssemblerARM64.cpp:
* assembler/MacroAssemblerARMv7.cpp:
* assembler/MacroAssemblerX86Common.cpp:
* assembler/testmasm.cpp:
(JSC::testProbePreservesGPRS):
(JSC::testProbeModifiesStackPointer):
(JSC::fillStack):
(JSC::testProbeModifiesStackWithCallback):
(JSC::run):

git-svn-id: https://svn.webkit.org/repository/webkit/trunk@220807 268f45cc-cd09-0410-ab3c-d52691b4dbfc

Source/JavaScriptCore/ChangeLog
Source/JavaScriptCore/assembler/MacroAssembler.h
Source/JavaScriptCore/assembler/MacroAssemblerARM.cpp
Source/JavaScriptCore/assembler/MacroAssemblerARM64.cpp
Source/JavaScriptCore/assembler/MacroAssemblerARMv7.cpp
Source/JavaScriptCore/assembler/MacroAssemblerX86Common.cpp
Source/JavaScriptCore/assembler/testmasm.cpp

index 6f05787..03e9d3d 100644 (file)
@@ -1,3 +1,76 @@
+2017-08-16  Mark Lam  <mark.lam@apple.com>
+
+        Enhance MacroAssembler::probe() to support an initializeStackFunction callback.
+        https://bugs.webkit.org/show_bug.cgi?id=175617
+        <rdar://problem/33912104>
+
+        Reviewed by JF Bastien.
+
+        This patch adds a new feature to MacroAssembler::probe() where the probe function
+        can provide a ProbeFunction callback to fill in stack values after the stack
+        pointer has been adjusted.  The probe function can use this feature as follows:
+
+        1. Set the new sp value in the ProbeContext's CPUState.
+
+        2. Set the ProbeContext's initializeStackFunction to a ProbeFunction callback
+           which will do the work of filling in the stack values after the probe
+           trampoline has adjusted the machine stack pointer.
+
+        3. Set the ProbeContext's initializeStackArgs to any value that the client wants
+           to pass to the initializeStackFunction callback.
+
+        4. Return from the probe function.
+
+        Upon returning from the probe function, the probe trampoline will adjust the
+        the stack pointer based on the sp value in CPUState.  If initializeStackFunction
+        is not set, the probe trampoline will restore registers and return to its caller.
+
+        If initializeStackFunction is set, the trampoline will move the ProbeContext
+        beyond the range of the stack pointer i.e. it will place the new ProbeContext at
+        an address lower than where CPUState.sp() points.  This ensures that the
+        ProbeContext will not be trashed by the initializeStackFunction when it writes to
+        the stack.  Then, the trampoline will call back to the initializeStackFunction
+        ProbeFunction to let it fill in the stack values as desired.  The
+        initializeStackFunction ProbeFunction will be passed the moved ProbeContext at
+        the new location.
+
+        initializeStackFunction may now write to the stack at addresses greater or
+        equal to CPUState.sp(), but not below that.  initializeStackFunction is also
+        not allowed to change CPUState.sp().  If the initializeStackFunction does not
+        abide by these rules, then behavior is undefined, and bad things may happen.
+
+        For future reference, some implementation details that this patch needed to
+        be mindful of:
+
+        1. When the probe trampoline allocates stack space for the ProbeContext, it
+           should include OUT_SIZE as well.  This ensures that it doesn't have to move
+           the ProbeContext on exit if the probe function didn't change the sp.
+
+        2. If the trampoline has to move the ProbeContext, it needs to point the machine
+           sp to new ProbeContext first before copying over the ProbeContext data.  This
+           protects the new ProbeContext from possibly being trashed by interrupts.
+
+        3. When computing the new address of ProbeContext to move to, we need to make
+           sure that it is properly aligned in accordance with stack ABI requirements
+           (just like we did when we allocated the ProbeContext on entry to the
+           probe trampoline).
+
+        4. When copying the ProbeContext to its new location, the trampoline should
+           always copy words from low addresses to high addresses.  This is because if
+           we're moving the ProbeContext, we'll always be moving it to a lower address.
+
+        * assembler/MacroAssembler.h:
+        * assembler/MacroAssemblerARM.cpp:
+        * assembler/MacroAssemblerARM64.cpp:
+        * assembler/MacroAssemblerARMv7.cpp:
+        * assembler/MacroAssemblerX86Common.cpp:
+        * assembler/testmasm.cpp:
+        (JSC::testProbePreservesGPRS):
+        (JSC::testProbeModifiesStackPointer):
+        (JSC::fillStack):
+        (JSC::testProbeModifiesStackWithCallback):
+        (JSC::run):
+
 2017-08-16  Csaba Osztrogon√°c  <ossy@webkit.org>
 
         Fix JSCOnly ARM buildbots after r220047 and r220184
index 65977f0..19878d1 100644 (file)
@@ -1842,6 +1842,37 @@ public:
     // The ProbeContext is stack allocated and is only valid for the duration
     // of the call to the user probe function.
     //
+    // The probe function may choose to move the stack pointer (in any direction).
+    // To do this, the probe function needs to set the new sp value in the CPUState.
+    //
+    // The probe function may also choose to fill stack space with some values.
+    // To do this, the probe function must first:
+    // 1. Set the new sp value in the ProbeContext's CPUState.
+    // 2. Set the ProbeContext's initializeStackFunction to a ProbeFunction callback
+    //    which will do the work of filling in the stack values after the probe
+    //    trampoline has adjusted the machine stack pointer.
+    // 3. Set the ProbeContext's initializeStackArgs to any value that the client wants
+    //    to pass to the initializeStackFunction callback.
+    // 4. Return from the probe function.
+    //
+    // Upon returning from the probe function, the probe trampoline will adjust the
+    // the stack pointer based on the sp value in CPUState. If initializeStackFunction
+    // is not set, the probe trampoline will restore registers and return to its caller.
+    //
+    // If initializeStackFunction is set, the trampoline will move the ProbeContext
+    // beyond the range of the stack pointer i.e. it will place the new ProbeContext at
+    // an address lower than where CPUState.sp() points. This ensures that the
+    // ProbeContext will not be trashed by the initializeStackFunction when it writes to
+    // the stack. Then, the trampoline will call back to the initializeStackFunction
+    // ProbeFunction to let it fill in the stack values as desired. The
+    // initializeStackFunction ProbeFunction will be passed the moved ProbeContext at
+    // the new location.
+    //
+    // initializeStackFunction may now write to the stack at addresses greater or
+    // equal to CPUState.sp(), but not below that. initializeStackFunction is also
+    // not allowed to change CPUState.sp(). If the initializeStackFunction does not
+    // abide by these rules, then behavior is undefined, and bad things may happen.
+    //
     // Note: this version of probe() should be implemented by the target specific
     // MacroAssembler.
     void probe(ProbeFunction, void* arg);
@@ -1995,6 +2026,8 @@ struct ProbeContext {
 
     ProbeFunction probeFunction;
     void* arg;
+    ProbeFunction initializeStackFunction;
+    void* initializeStackArg;
     CPUState cpu;
 
     // Convenience methods:
index 7e206ac..efac95b 100644 (file)
@@ -99,15 +99,17 @@ void MacroAssemblerARM::load32WithUnalignedHalfWords(BaseIndex address, Register
 extern "C" void ctiMasmProbeTrampoline();
 
 #if COMPILER(GCC_OR_CLANG)
-    
+
 // The following are offsets for ProbeContext fields accessed
 // by the ctiMasmProbeTrampoline stub.
 
 #define PTR_SIZE 4
 #define PROBE_PROBE_FUNCTION_OFFSET (0 * PTR_SIZE)
 #define PROBE_ARG_OFFSET (1 * PTR_SIZE)
+#define PROBE_INIT_STACK_FUNCTION_OFFSET (2 * PTR_SIZE)
+#define PROBE_INIT_STACK_ARG_OFFSET (3 * PTR_SIZE)
 
-#define PROBE_FIRST_GPREG_OFFSET (2 * PTR_SIZE)
+#define PROBE_FIRST_GPREG_OFFSET (4 * PTR_SIZE)
 
 #define GPREG_SIZE 4
 #define PROBE_CPU_R0_OFFSET (PROBE_FIRST_GPREG_OFFSET + (0 * GPREG_SIZE))
@@ -167,13 +169,16 @@ extern "C" void ctiMasmProbeTrampoline();
 #define PROBE_CPU_D31_OFFSET (PROBE_FIRST_FPREG_OFFSET + (31 * FPREG_SIZE))
 
 #define PROBE_SIZE (PROBE_FIRST_FPREG_OFFSET + (32 * FPREG_SIZE))
-#define PROBE_ALIGNED_SIZE (PROBE_SIZE)
+
+#define OUT_SIZE GPREG_SIZE
 
 // These ASSERTs remind you that if you change the layout of ProbeContext,
 // you need to change ctiMasmProbeTrampoline offsets above to match.
 #define PROBE_OFFSETOF(x) offsetof(struct ProbeContext, x)
 COMPILE_ASSERT(PROBE_OFFSETOF(probeFunction) == PROBE_PROBE_FUNCTION_OFFSET, ProbeContext_probeFunction_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(arg) == PROBE_ARG_OFFSET, ProbeContext_arg_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackFunction) == PROBE_INIT_STACK_FUNCTION_OFFSET, ProbeContext_initializeStackFunction_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackArg) == PROBE_INIT_STACK_ARG_OFFSET, ProbeContext_initializeStackArg_offset_matches_ctiMasmProbeTrampoline);
 
 COMPILE_ASSERT(!(PROBE_CPU_R0_OFFSET & 0x3), ProbeContext_cpu_r0_offset_should_be_4_byte_aligned);
 
@@ -197,7 +202,7 @@ COMPILE_ASSERT(PROBE_OFFSETOF(cpu.gprs[ARMRegisters::pc]) == PROBE_CPU_PC_OFFSET
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.sprs[ARMRegisters::apsr]) == PROBE_CPU_APSR_OFFSET, ProbeContext_cpu_apsr_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.sprs[ARMRegisters::fpscr]) == PROBE_CPU_FPSCR_OFFSET, ProbeContext_cpu_fpscr_offset_matches_ctiMasmProbeTrampoline);
 
-COMPILE_ASSERT(!(PROBE_CPU_D0_OFFSET & 0xf), ProbeContext_cpu_d0_offset_should_be_16_byte_aligned);
+COMPILE_ASSERT(!(PROBE_CPU_D0_OFFSET & 0x7), ProbeContext_cpu_d0_offset_should_be_8_byte_aligned);
 
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d0]) == PROBE_CPU_D0_OFFSET, ProbeContext_cpu_d0_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d1]) == PROBE_CPU_D1_OFFSET, ProbeContext_cpu_d1_offset_matches_ctiMasmProbeTrampoline);
@@ -233,7 +238,6 @@ COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d30]) == PROBE_CPU_D30_OFFS
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d31]) == PROBE_CPU_D31_OFFSET, ProbeContext_cpu_d31_offset_matches_ctiMasmProbeTrampoline);
 
 COMPILE_ASSERT(sizeof(ProbeContext) == PROBE_SIZE, ProbeContext_size_matches_ctiMasmProbeTrampoline);
-COMPILE_ASSERT(!(PROBE_ALIGNED_SIZE & 0xf), ProbeContext_aligned_size_offset_should_be_16_byte_aligned);
 #undef PROBE_OFFSETOF
 
 asm (
@@ -254,11 +258,11 @@ asm (
 
     "mov       ip, sp" "\n"
     "mov       r3, sp" "\n"
-    "sub       r3, r3, #" STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) "\n"
+    "sub       r3, r3, #" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) "\n"
 
     // The ARM EABI specifies that the stack needs to be 16 byte aligned.
     "bic       r3, r3, #0xf" "\n"
-    "mov       sp, r3" "\n"
+    "mov       sp, r3" "\n" // Set the sp to protect the ProbeContext from interrupts before we initialize it.
 
     "str       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
     "add       lr, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_R0_OFFSET) "\n"
@@ -284,15 +288,59 @@ asm (
     "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
 
     "add       ip, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_D0_OFFSET) "\n"
-    "vstmia.64 ip, { d0-d15 }" "\n"
-    "vstmia.64 ip, { d16-d31 }" "\n"
+    "vstmia.64 ip!, { d0-d15 }" "\n"
+    "vstmia.64 ip!, { d16-d31 }" "\n"
 
     "mov       fp, sp" "\n" // Save the ProbeContext*.
 
+    // Initialize ProbeContext::initializeStackFunction to zero.
+    "mov       r0, #0" "\n"
+    "str       r0, [fp, #" STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "]" "\n"
+
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_PROBE_FUNCTION_OFFSET) "]" "\n"
     "mov       r0, sp" "\n" // the ProbeContext* arg.
     "blx       ip" "\n"
 
+    // Make sure the ProbeContext is entirely below the result stack pointer so
+    // that register values are still preserved when we call the initializeStack
+    // function.
+    "ldr       r1, [fp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n" // Result sp.
+    "add       r2, fp, #" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) "\n" // End of ProveContext + buffer.
+    "cmp       r1, r2" "\n"
+    "bge     " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) "\n"
+
+    // Allocate a safe place on the stack below the result stack pointer to stash the ProbeContext.
+    "sub       r1, r1, #" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) "\n"
+    "bic       r1, r1, #0xf" "\n" // The ARM EABI specifies that the stack needs to be 16 byte aligned.
+    "mov       sp, r1" "\n" // Set the new sp to protect that memory from interrupts before we copy the ProbeContext.
+
+    // Copy the ProbeContext to the safe place.
+    // Note: we have to copy from low address to higher address because we're moving the
+    // ProbeContext to a lower address.
+    "mov       r5, fp" "\n"
+    "mov       r6, r1" "\n"
+    "add       r7, fp, #" STRINGIZE_VALUE_OF(PROBE_SIZE) "\n"
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) ":" "\n"
+    "ldr       r3, [r5], #4" "\n"
+    "ldr       r4, [r5], #4" "\n"
+    "str       r3, [r6], #4" "\n"
+    "str       r4, [r6], #4" "\n"
+    "cmp       r5, r7" "\n"
+    "blt     " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) "\n"
+
+    "mov       fp, r1" "\n"
+
+    // Call initializeStackFunction if present.
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) ":" "\n"
+    "ldr       r2, [fp, #" STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "]" "\n"
+    "cbz       r2, " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) "\n"
+
+    "mov       r0, fp" "\n" // Set the ProbeContext* arg.
+    "blx       r2" "\n" // Call the initializeStackFunction (loaded into r2 above).
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) ":" "\n"
+
     "mov       sp, fp" "\n"
 
     // To enable probes to modify register state, we copy all registers
@@ -307,82 +355,13 @@ asm (
     "vmsr      FPSCR, ip" "\n"
 
     // There are 5 more registers left to restore: ip, sp, lr, pc, and apsr.
-    // There are 2 issues that complicate the restoration of these last few
-    // registers:
-    //
-    // 1. Normal ARM calling convention relies on moving lr to pc to return to
-    //    the caller. In our case, the address to return to is specified by
-    //    ProbeContext.cpu.pc. And at that moment, we won't have any available
-    //    scratch registers to hold the return address (lr needs to hold
-    //    ProbeContext.cpu.lr, not the return address).
-    //
-    //    The solution is to store the return address on the stack and load the
-    //     pc from there.
-    //
-    // 2. Issue 1 means we will need to write to the stack location at
-    //    ProbeContext.cpu.gprs[sp] - PTR_SIZE. But if the user probe function had
-    //    modified the value of ProbeContext.cpu.gprs[sp] to point in the range between
-    //    &ProbeContext.cpu.gprs[ip] thru &ProbeContext.cpu.sprs[aspr], then the action
-    //    for Issue 1 may trash the values to be restored before we can restore them.
-    //
-    //    The solution is to check if ProbeContext.cpu.gprs[sp] contains a value in
-    //    the undesirable range. If so, we copy the remaining ProbeContext
-    //    register data to a safe area first, and restore the remaining register
-    //    from this new safe area.
-
-    // The restore area for the pc will be located at 1 word below the resultant sp.
-    // All restore values are located at offset <= PROBE_CPU_APSR_OFFSET. Hence,
-    // we need to make sure that resultant sp > offset of apsr + 1.
-    "add       ip, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET + PTR_SIZE) "\n"
-    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "cmp       lr, ip" "\n"
-    "bgt     " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
-
-    // Getting here means that the restore area will overlap the ProbeContext data
-    // that we will need to get the restoration values from. So, let's move that
-    // data to a safe place before we start writing into the restore area.
-    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
-    // restore area. This ensures that:
-    // 1. The safe area does not overlap the restore area.
-    // 2. The safe area does not overlap the ProbeContext.
-    //    This makes it so that we can use memcpy (does not require memmove) semantics
-    //    to copy the restore values to the safe area.
-    
-    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
-    "sub       lr, lr, #(2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ")" "\n"
-
-    "mov       ip, sp" "\n" // Save the original ProbeContext*.
-
-    // Make sure the stack pointer points to the safe area. This ensures that the
-    // safe area is protected from interrupt handlers overwriting it.
-    "mov       sp, lr" "\n" // sp now points to the new ProbeContext in the safe area.
-
-    "mov       lr, ip" "\n" // Use lr as the old ProbeContext*.
-
-    // Copy the restore data to the new ProbeContext*.
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
-
-    // ctiMasmProbeTrampolineEnd expects lr to contain the sp value to be restored.
-    // Since we used it as scratch above, let's restore it.
-    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-
-    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
 
     // Set up the restore area for sp and pc.
-    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
+    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
 
     // Push the pc on to the restore area.
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "sub       lr, lr, #" STRINGIZE_VALUE_OF(PTR_SIZE) "\n"
+    "sub       lr, lr, #" STRINGIZE_VALUE_OF(OUT_SIZE) "\n"
     "str       ip, [lr]" "\n"
     // Point sp to the restore area.
     "str       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
index 60b3dcb..c2cd4b8 100644 (file)
@@ -43,8 +43,10 @@ using namespace ARM64Registers;
 #define PTR_SIZE 8
 #define PROBE_PROBE_FUNCTION_OFFSET (0 * PTR_SIZE)
 #define PROBE_ARG_OFFSET (1 * PTR_SIZE)
+#define PROBE_INIT_STACK_FUNCTION_OFFSET (2 * PTR_SIZE)
+#define PROBE_INIT_STACK_ARG_OFFSET (3 * PTR_SIZE)
 
-#define PROBE_FIRST_GPREG_OFFSET (2 * PTR_SIZE)
+#define PROBE_FIRST_GPREG_OFFSET (4 * PTR_SIZE)
 
 #define GPREG_SIZE 8
 #define PROBE_CPU_X0_OFFSET (PROBE_FIRST_GPREG_OFFSET + (0 * GPREG_SIZE))
@@ -131,6 +133,8 @@ using namespace ARM64Registers;
 #define PROBE_OFFSETOF(x) offsetof(struct ProbeContext, x)
 COMPILE_ASSERT(PROBE_OFFSETOF(probeFunction) == PROBE_PROBE_FUNCTION_OFFSET, ProbeContext_probeFunction_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(arg) == PROBE_ARG_OFFSET, ProbeContext_arg_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackFunction) == PROBE_INIT_STACK_FUNCTION_OFFSET, ProbeContext_initializeStackFunction_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackArg) == PROBE_INIT_STACK_ARG_OFFSET, ProbeContext_initializeStackArg_offset_matches_ctiMasmProbeTrampoline);
 
 COMPILE_ASSERT(!(PROBE_CPU_X0_OFFSET & 0x7), ProbeContext_cpu_r0_offset_should_be_8_byte_aligned);
 
@@ -212,6 +216,7 @@ COMPILE_ASSERT(sizeof(ProbeContext) == PROBE_SIZE, ProbeContext_size_matches_cti
 
 // Conditions for using ldp and stp.
 static_assert(PROBE_CPU_PC_OFFSET == PROBE_CPU_SP_OFFSET + PTR_SIZE, "PROBE_CPU_SP_OFFSET and PROBE_CPU_PC_OFFSET must be adjacent");
+static_assert(!(PROBE_SIZE_PLUS_EXTRAS & 0xf), "PROBE_SIZE_PLUS_EXTRAS should be 16 byte aligned"); // the ProbeContext copying code relies on this.
 
 #undef PROBE_OFFSETOF
 
@@ -269,7 +274,7 @@ static_assert(OUT_X28_OFFSET == offsetof(OutgoingProbeRecord, x28), "OUT_X28_OFF
 static_assert(OUT_FP_OFFSET == offsetof(OutgoingProbeRecord, fp), "OUT_FP_OFFSET is incorrect");
 static_assert(OUT_LR_OFFSET == offsetof(OutgoingProbeRecord, lr), "OUT_LR_OFFSET is incorrect");
 static_assert(OUT_SIZE == sizeof(OutgoingProbeRecord), "OUT_SIZE is incorrect");
-static_assert(!(sizeof(OutgoingProbeRecord) & 0xf), "IncomingProbeStack must be 16-byte aligned");
+static_assert(!(sizeof(OutgoingProbeRecord) & 0xf), "OutgoingProbeStack must be 16-byte aligned");
 
 #define STATE_PC_NOT_CHANGED 0
 #define STATE_PC_CHANGED 1
@@ -290,9 +295,9 @@ asm (
     "mov       x26, sp" "\n"
     "mov       x27, sp" "\n"
 
-    "sub       x27, x27, #" STRINGIZE_VALUE_OF(PROBE_SIZE_PLUS_EXTRAS) "\n"
+    "sub       x27, x27, #" STRINGIZE_VALUE_OF(PROBE_SIZE_PLUS_EXTRAS + OUT_SIZE) "\n"
     "bic       x27, x27, #0xf" "\n" // The ARM EABI specifies that the stack needs to be 16 byte aligned.
-    "mov       sp, x27" "\n" // Make sure interrupts don't over-write our data on the stack.
+    "mov       sp, x27" "\n" // Set the sp to protect the ProbeContext from interrupts before we initialize it.
 
     "stp       x0, x1, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_X0_OFFSET) "]" "\n"
     "mrs       x0, nzcv" "\n" // Preload nzcv.
@@ -350,11 +355,52 @@ asm (
 
     "mov       x27, sp" "\n" // Save the ProbeContext* in a callee saved register.
 
+    // Initialize ProbeContext::initializeStackFunction to zero.
+    "str       xzr, [x27, #" STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "]" "\n"
+
     // Note: we haven't changed the value of fp. Hence, it is still pointing to the frame of
     // the caller of the probe (which is what we want in order to play nice with debuggers e.g. lldb).
     "mov       x0, sp" "\n" // Set the ProbeContext* arg.
     "blr       x2" "\n" // Call the probe handler function (loaded into x2 above).
 
+    // Make sure the ProbeContext is entirely below the result stack pointer so
+    // that register values are still preserved when we call the initializeStack
+    // function.
+    "ldr       x1, [x27, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n" // Result sp.
+    "add       x2, x27, #" STRINGIZE_VALUE_OF(PROBE_SIZE_PLUS_EXTRAS + OUT_SIZE) "\n" // End of ProbeContext + buffer.
+    "cmp       x1, x2" "\n"
+    "bge     " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) "\n"
+
+    // Allocate a safe place on the stack below the result stack pointer to stash the ProbeContext.
+    "sub       x1, x1, #" STRINGIZE_VALUE_OF(PROBE_SIZE_PLUS_EXTRAS + OUT_SIZE) "\n"
+    "bic       x1, x1, #0xf" "\n" // The ARM EABI specifies that the stack needs to be 16 byte aligned.
+    "mov       sp, x1" "\n" // Set the new sp to protect that memory from interrupts before we copy the ProbeContext.
+
+    // Copy the ProbeContext to the safe place.
+    // Note: we have to copy from low address to higher address because we're moving the
+    // ProbeContext to a lower address.
+    "mov       x5, x27" "\n"
+    "mov       x6, x1" "\n"
+    "add       x7, x27, #" STRINGIZE_VALUE_OF(PROBE_SIZE_PLUS_EXTRAS) "\n"
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) ":" "\n"
+    "ldp       x3, x4, [x5], #16" "\n"
+    "stp       x3, x4, [x6], #16" "\n"
+    "cmp       x5, x7" "\n"
+    "blt     " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) "\n"
+
+    "mov       x27, x1" "\n"
+
+    // Call initializeStackFunction if present.
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) ":" "\n"
+    "ldr       x2, [x27, #" STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "]" "\n"
+    "cbz       x2, " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) "\n"
+
+    "mov       x0, x27" "\n" // Set the ProbeContext* arg.
+    "blr       x2" "\n" // Call the initializeStackFunction (loaded into x2 above).
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) ":" "\n"
+
     "mov       sp, x27" "\n"
 
     // To enable probes to modify register state, we copy all registers
@@ -428,48 +474,6 @@ asm (
 
     "ldr       x29, [sp, #" STRINGIZE_VALUE_OF(SAVED_PROBE_RETURN_PC_OFFSET) "]" "\n" // Preload the probe return site pc.
 
-    // The probe handler may have moved the sp. For the return process, we may need
-    // space for 2 OutgoingProbeRecords below the final sp value. We need to make
-    // sure that the space for these 2 OutgoingProbeRecords do not overlap the
-    // restore values of the registers.
-
-    // All restore values are located at offset <= PROBE_CPU_FPSR_OFFSET. Hence,
-    // we need to make sure that resultant sp > offset of fpsr + 2 * sizeof(OutgoingProbeRecord).
-
-    "add       x27, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_FPSR_OFFSET + 2 * OUT_SIZE) "\n"
-    "ldr       x28, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "cmp       x28, x27" "\n"
-    "bgt     " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineFillOutgoingProbeRecords) "\n"
-
-    // There is overlap. We need to copy the ProbeContext to a safe area first.
-    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the OutgoingProbeRecords are.
-    // This ensures that:
-    // 1. The safe area does not overlap the OutgoingProbeRecords.
-    // 2. The safe area does not overlap the ProbeContext.
-
-    // x28 already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
-    "sub       x28, x28, #" STRINGIZE_VALUE_OF(2 * PROBE_SIZE) "\n"
-
-    "mov       x27, sp" "\n" // Save the original ProbeContext*.
-
-    // Make sure the stack pointer points to the safe area. This ensures that the
-    // safe area is protected from interrupt handlers overwriting it.
-    "mov       sp, x28" "\n" // sp now points to the new ProbeContext in the safe area.
-
-    // Copy the relevant restore data to the new ProbeContext*.
-    "str       x30, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_X0_OFFSET) "]" "\n" // Stash the pc changed state away so that we can use lr.
-
-    "ldp       x28, x30, [x27, #" STRINGIZE_VALUE_OF(PROBE_CPU_X27_OFFSET) "]" "\n" // copy x27 and x28.
-    "stp       x28, x30, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_X27_OFFSET) "]" "\n"
-    "ldp       x28, x30, [x27, #" STRINGIZE_VALUE_OF(PROBE_CPU_FP_OFFSET) "]" "\n" // copy fp and lr.
-    "stp       x28, x30, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_FP_OFFSET) "]" "\n"
-    "ldp       x28, x30, [x27, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n" // copy sp and pc.
-    "stp       x28, x30, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "ldp       x28, x30, [x27, #" STRINGIZE_VALUE_OF(PROBE_CPU_NZCV_OFFSET) "]" "\n" // copy nzcv and fpsr.
-    "stp       x28, x30, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_NZCV_OFFSET) "]" "\n"
-
-    "ldr       x30, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_X0_OFFSET) "]" "\n" // Retrieve the stashed the pc changed state.
-
     LOCAL_LABEL_STRING(ctiMasmProbeTrampolineFillOutgoingProbeRecords) ":" "\n"
 
     "cbnz       x30, " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineEnd) "\n" // Skip lr restoration setup if state (in lr) == STATE_PC_CHANGED.
index 1d68704..e98aee2 100644 (file)
@@ -42,8 +42,10 @@ extern "C" void ctiMasmProbeTrampoline();
 #define PTR_SIZE 4
 #define PROBE_PROBE_FUNCTION_OFFSET (0 * PTR_SIZE)
 #define PROBE_ARG_OFFSET (1 * PTR_SIZE)
+#define PROBE_INIT_STACK_FUNCTION_OFFSET (2 * PTR_SIZE)
+#define PROBE_INIT_STACK_ARG_OFFSET (3 * PTR_SIZE)
 
-#define PROBE_FIRST_GPREG_OFFSET (2 * PTR_SIZE)
+#define PROBE_FIRST_GPREG_OFFSET (4 * PTR_SIZE)
 
 #define GPREG_SIZE 4
 #define PROBE_CPU_R0_OFFSET (PROBE_FIRST_GPREG_OFFSET + (0 * GPREG_SIZE))
@@ -101,14 +103,18 @@ extern "C" void ctiMasmProbeTrampoline();
 #define PROBE_CPU_D29_OFFSET (PROBE_FIRST_FPREG_OFFSET + (29 * FPREG_SIZE))
 #define PROBE_CPU_D30_OFFSET (PROBE_FIRST_FPREG_OFFSET + (30 * FPREG_SIZE))
 #define PROBE_CPU_D31_OFFSET (PROBE_FIRST_FPREG_OFFSET + (31 * FPREG_SIZE))
+
 #define PROBE_SIZE (PROBE_FIRST_FPREG_OFFSET + (32 * FPREG_SIZE))
-#define PROBE_ALIGNED_SIZE (PROBE_SIZE)
+
+#define OUT_SIZE GPREG_SIZE
 
 // These ASSERTs remind you that if you change the layout of ProbeContext,
 // you need to change ctiMasmProbeTrampoline offsets above to match.
 #define PROBE_OFFSETOF(x) offsetof(struct ProbeContext, x)
 COMPILE_ASSERT(PROBE_OFFSETOF(probeFunction) == PROBE_PROBE_FUNCTION_OFFSET, ProbeContext_probeFunction_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(arg) == PROBE_ARG_OFFSET, ProbeContext_arg_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackFunction) == PROBE_INIT_STACK_FUNCTION_OFFSET, ProbeContext_initializeStackFunction_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackArg) == PROBE_INIT_STACK_ARG_OFFSET, ProbeContext_initializeStackArg_offset_matches_ctiMasmProbeTrampoline);
 
 COMPILE_ASSERT(!(PROBE_CPU_R0_OFFSET & 0x3), ProbeContext_cpu_r0_offset_should_be_4_byte_aligned);
 
@@ -132,7 +138,7 @@ COMPILE_ASSERT(PROBE_OFFSETOF(cpu.gprs[ARMRegisters::pc]) == PROBE_CPU_PC_OFFSET
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.sprs[ARMRegisters::apsr]) == PROBE_CPU_APSR_OFFSET, ProbeContext_cpu_apsr_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.sprs[ARMRegisters::fpscr]) == PROBE_CPU_FPSCR_OFFSET, ProbeContext_cpu_fpscr_offset_matches_ctiMasmProbeTrampoline);
 
-COMPILE_ASSERT(!(PROBE_CPU_D0_OFFSET & 0xf), ProbeContext_cpu_d0_offset_should_be_16_byte_aligned);
+COMPILE_ASSERT(!(PROBE_CPU_D0_OFFSET & 0x7), ProbeContext_cpu_d0_offset_should_be_8_byte_aligned);
 
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d0]) == PROBE_CPU_D0_OFFSET, ProbeContext_cpu_d0_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d1]) == PROBE_CPU_D1_OFFSET, ProbeContext_cpu_d1_offset_matches_ctiMasmProbeTrampoline);
@@ -169,8 +175,6 @@ COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d30]) == PROBE_CPU_D30_OFFS
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[ARMRegisters::d31]) == PROBE_CPU_D31_OFFSET, ProbeContext_cpu_d31_offset_matches_ctiMasmProbeTrampoline);
 
 COMPILE_ASSERT(sizeof(ProbeContext) == PROBE_SIZE, ProbeContext_size_matches_ctiMasmProbeTrampoline);
-COMPILE_ASSERT(!(PROBE_ALIGNED_SIZE & 0xf), ProbeContext_aligned_size_offset_should_be_16_byte_aligned);
-
 #undef PROBE_OFFSETOF
     
 asm (
@@ -193,11 +197,11 @@ asm (
 
     "mov       ip, sp" "\n"
     "mov       r0, sp" "\n"
-    "sub       r0, r0, #" STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) "\n"
+    "sub       r0, r0, #" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) "\n"
 
     // The ARM EABI specifies that the stack needs to be 16 byte aligned.
     "bic       r0, r0, #0xf" "\n"
-    "mov       sp, r0" "\n"
+    "mov       sp, r0" "\n" // Set the sp to protect the ProbeContext from interrupts before we initialize it.
 
     "str       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
     "add       lr, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_R1_OFFSET) "\n"
@@ -228,10 +232,56 @@ asm (
 
     "mov       fp, sp" "\n" // Save the ProbeContext*.
 
+    // Initialize ProbeContext::initializeStackFunction to zero.
+    "mov       r0, #0" "\n"
+    "str       r0, [fp, #" STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "]" "\n"
+
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_PROBE_FUNCTION_OFFSET) "]" "\n"
     "mov       r0, sp" "\n" // the ProbeContext* arg.
     "blx       ip" "\n"
 
+    // Make sure the ProbeContext is entirely below the result stack pointer so
+    // that register values are still preserved when we call the initializeStack
+    // function.
+    "ldr       r1, [fp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n" // Result sp.
+    "add       r2, fp, #" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) "\n" // End of ProveContext + buffer.
+    "cmp       r1, r2" "\n"
+    "it        ge" "\n"
+    "bge     " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) "\n"
+
+    // Allocate a safe place on the stack below the result stack pointer to stash the ProbeContext.
+    "sub       r1, r1, #" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) "\n"
+    "bic       r1, r1, #0xf" "\n" // The ARM EABI specifies that the stack needs to be 16 byte aligned.
+    "mov       sp, r1" "\n" // Set the new sp to protect that memory from interrupts before we copy the ProbeContext.
+
+    // Copy the ProbeContext to the safe place.
+    // Note: we have to copy from low address to higher address because we're moving the
+    // ProbeContext to a lower address.
+    "mov       r5, fp" "\n"
+    "mov       r6, r1" "\n"
+    "add       r7, fp, #" STRINGIZE_VALUE_OF(PROBE_SIZE) "\n"
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) ":" "\n"
+    "ldr       r3, [r5], #4" "\n"
+    "ldr       r4, [r5], #4" "\n"
+    "str       r3, [r6], #4" "\n"
+    "str       r4, [r6], #4" "\n"
+    "cmp       r5, r7" "\n"
+    "it        lt" "\n"
+    "blt     " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) "\n"
+
+    "mov       fp, r1" "\n"
+
+    // Call initializeStackFunction if present.
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) ":" "\n"
+    "ldr       r2, [fp, #" STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "]" "\n"
+    "cbz       r2, " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) "\n"
+
+    "mov       r0, fp" "\n" // Set the ProbeContext* arg.
+    "blx       r2" "\n" // Call the initializeStackFunction (loaded into r2 above).
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) ":" "\n"
+
     "mov       sp, fp" "\n"
 
     // To enable probes to modify register state, we copy all registers
@@ -247,84 +297,13 @@ asm (
     "vmsr      FPSCR, ip" "\n"
 
     // There are 5 more registers left to restore: ip, sp, lr, pc, and apsr.
-    // There are 2 issues that complicate the restoration of these last few
-    // registers:
-    //
-    // 1. Normal ARM calling convention relies on moving lr to pc to return to
-    //    the caller. In our case, the address to return to is specified by
-    //    ProbeContext.cpu.gprs[pc]. And at that moment, we won't have any available
-    //    scratch registers to hold the return address (lr needs to hold
-    //    ProbeContext.cpu.gprs[lr], not the return address).
-    //
-    //    The solution is to store the return address on the stack and load the
-    //    pc from there.
-    //
-    // 2. Issue 1 means we will need to write to the stack location at
-    //    ProbeContext.cpu.gprs[sp] - PTR_SIZE. But if the user probe function had
-    //    modified the value of ProbeContext.cpu.gprs[sp] to point in the range between
-    //    &ProbeContext.cpu.gprs[ip] thru &ProbeContext.cpu.sprs[aspr], then the action
-    //    for Issue 1 may trash the values to be restored before we can restore them.
-    //
-    //    The solution is to check if ProbeContext.cpu.gprs[sp] contains a value in
-    //    the undesirable range. If so, we copy the remaining ProbeContext
-    //    register data to a safe area first, and restore the remaining register
-    //    from this new safe area.
-
-    // The restore area for the pc will be located at 1 word below the resultant sp.
-    // All restore values are located at offset <= PROBE_CPU_APSR_OFFSET. Hence,
-    // we need to make sure that resultant sp > offset of apsr + 1.
-    "add       ip, sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET + PTR_SIZE) "\n"
-    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "cmp       lr, ip" "\n"
-    "it        gt" "\n"
-    "bgt     " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
-
-    // Getting here means that the restore area will overlap the ProbeContext data
-    // that we will need to get the restoration values from. So, let's move that
-    // data to a safe place before we start writing into the restore area.
-    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
-    // restore area. This ensures that:
-    // 1. The safe area does not overlap the restore area.
-    // 2. The safe area does not overlap the ProbeContext.
-    //    This makes it so that we can use memcpy (does not require memmove) semantics
-    //    to copy the restore values to the safe area.
-
-    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
-    "sub       lr, lr, #(2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ")" "\n"
-    
-    "mov       ip, sp" "\n" // Save the original ProbeContext*.
-    
-    // Make sure the stack pointer points to the safe area. This ensures that the
-    // safe area is protected from interrupt handlers overwriting it.
-    "mov       sp, lr" "\n" // sp now points to the new ProbeContext in the safe area.
-    
-    "mov       lr, ip" "\n" // Use lr as the old ProbeContext*.
-    
-    // Copy the restore data to the new ProbeContext*.
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_IP_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_LR_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "ldr       ip, [lr, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
-    "str       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_APSR_OFFSET) "]" "\n"
-
-    // ctiMasmProbeTrampolineEnd expects lr to contain the sp value to be restored.
-    // Since we used it as scratch above, let's restore it.
-    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
-
-    ".thumb_func " THUMB_FUNC_PARAM(ctiMasmProbeTrampolineEnd) "\n"
-    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
 
     // Set up the restore area for sp and pc.
-    // lr already contains [sp, #STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET)].
+    "ldr       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
 
     // Push the pc on to the restore area.
     "ldr       ip, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_PC_OFFSET) "]" "\n"
-    "sub       lr, lr, #" STRINGIZE_VALUE_OF(PTR_SIZE) "\n"
+    "sub       lr, lr, #" STRINGIZE_VALUE_OF(OUT_SIZE) "\n"
     "str       ip, [lr]" "\n"
     // Point sp to the restore area.
     "str       lr, [sp, #" STRINGIZE_VALUE_OF(PROBE_CPU_SP_OFFSET) "]" "\n"
index a7fe01a..1a56605 100644 (file)
@@ -46,8 +46,10 @@ extern "C" void ctiMasmProbeTrampoline();
 
 #define PROBE_PROBE_FUNCTION_OFFSET (0 * PTR_SIZE)
 #define PROBE_ARG_OFFSET (1 * PTR_SIZE)
+#define PROBE_INIT_STACK_FUNCTION_OFFSET (2 * PTR_SIZE)
+#define PROBE_INIT_STACK_ARG_OFFSET (3 * PTR_SIZE)
 
-#define PROBE_FIRST_GPR_OFFSET (2 * PTR_SIZE)
+#define PROBE_FIRST_GPR_OFFSET (4 * PTR_SIZE)
 #define PROBE_CPU_EAX_OFFSET (PROBE_FIRST_GPR_OFFSET + (0 * PTR_SIZE))
 #define PROBE_CPU_ECX_OFFSET (PROBE_FIRST_GPR_OFFSET + (1 * PTR_SIZE))
 #define PROBE_CPU_EDX_OFFSET (PROBE_FIRST_GPR_OFFSET + (2 * PTR_SIZE))
@@ -87,7 +89,6 @@ extern "C" void ctiMasmProbeTrampoline();
 
 #if CPU(X86)
 #define PROBE_SIZE (PROBE_CPU_XMM7_OFFSET + XMM_SIZE)
-#define PROBE_ALIGNED_SIZE (PROBE_SIZE + (2 * XMM_SIZE))
 #else // CPU(X86_64)
 #define PROBE_CPU_XMM8_OFFSET (PROBE_FIRST_XMM_OFFSET + (8 * XMM_SIZE))
 #define PROBE_CPU_XMM9_OFFSET (PROBE_FIRST_XMM_OFFSET + (9 * XMM_SIZE))
@@ -98,14 +99,19 @@ extern "C" void ctiMasmProbeTrampoline();
 #define PROBE_CPU_XMM14_OFFSET (PROBE_FIRST_XMM_OFFSET + (14 * XMM_SIZE))
 #define PROBE_CPU_XMM15_OFFSET (PROBE_FIRST_XMM_OFFSET + (15 * XMM_SIZE))
 #define PROBE_SIZE (PROBE_CPU_XMM15_OFFSET + XMM_SIZE)
-#define PROBE_ALIGNED_SIZE (PROBE_SIZE + (4 * XMM_SIZE))
 #endif // CPU(X86_64)
 
+// The outgoing record to be popped off the stack at the end consists of:
+// eflags, eax, ecx, ebp, eip.
+#define OUT_SIZE        (5 * PTR_SIZE)
+
 // These ASSERTs remind you that if you change the layout of ProbeContext,
 // you need to change ctiMasmProbeTrampoline offsets above to match.
 #define PROBE_OFFSETOF(x) offsetof(struct ProbeContext, x)
 COMPILE_ASSERT(PROBE_OFFSETOF(probeFunction) == PROBE_PROBE_FUNCTION_OFFSET, ProbeContext_probeFunction_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(arg) == PROBE_ARG_OFFSET, ProbeContext_arg_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackFunction) == PROBE_INIT_STACK_FUNCTION_OFFSET, ProbeContext_initializeStackFunction_offset_matches_ctiMasmProbeTrampoline);
+COMPILE_ASSERT(PROBE_OFFSETOF(initializeStackArg) == PROBE_INIT_STACK_ARG_OFFSET, ProbeContext_initializeStackArg_offset_matches_ctiMasmProbeTrampoline);
 
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.gprs[X86Registers::eax]) == PROBE_CPU_EAX_OFFSET, ProbeContext_cpu_eax_offset_matches_ctiMasmProbeTrampoline);
 COMPILE_ASSERT(PROBE_OFFSETOF(cpu.gprs[X86Registers::ecx]) == PROBE_CPU_ECX_OFFSET, ProbeContext_cpu_ecx_offset_matches_ctiMasmProbeTrampoline);
@@ -152,7 +158,6 @@ COMPILE_ASSERT(PROBE_OFFSETOF(cpu.fprs[X86Registers::xmm15]) == PROBE_CPU_XMM15_
 #endif // CPU(X86_64)
 
 COMPILE_ASSERT(sizeof(ProbeContext) == PROBE_SIZE, ProbeContext_size_matches_ctiMasmProbeTrampoline);
-COMPILE_ASSERT(!(PROBE_ALIGNED_SIZE & 0x1f), ProbeContext_aligned_size_offset_should_be_32_byte_aligned);
 
 #undef PROBE_OFFSETOF
 
@@ -175,10 +180,9 @@ asm (
     //     esp[5 * ptrSize]: saved esp
 
     "movl %esp, %eax" "\n"
-    "subl $" STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ", %esp" "\n"
+    "subl $" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) ", %esp" "\n"
 
-    // The X86_64 ABI specifies that the worse case stack alignment requirement
-    // is 32 bytes.
+    // The X86_64 ABI specifies that the worse case stack alignment requirement is 32 bytes.
     "andl $~0x1f, %esp" "\n"
 
     "movl %ebp, " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%esp)" "\n"
@@ -212,6 +216,9 @@ asm (
     "movq %xmm6, " STRINGIZE_VALUE_OF(PROBE_CPU_XMM6_OFFSET) "(%ebp)" "\n"
     "movq %xmm7, " STRINGIZE_VALUE_OF(PROBE_CPU_XMM7_OFFSET) "(%ebp)" "\n"
 
+    "xorl %eax, %eax" "\n"
+    "movl %eax, " STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "(%ebp)" "\n"
+
     // Reserve stack space for the arg while maintaining the required stack
     // pointer 32 byte alignment:
     "subl $0x20, %esp" "\n"
@@ -219,6 +226,48 @@ asm (
 
     "call *" STRINGIZE_VALUE_OF(PROBE_PROBE_FUNCTION_OFFSET) "(%ebp)" "\n"
 
+    // Make sure the ProbeContext is entirely below the result stack pointer so
+    // that register values are still preserved when we call the initializeStack
+    // function.
+    "movl $" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) ", %ecx" "\n"
+    "movl %ebp, %eax" "\n"
+    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp), %edx" "\n"
+    "addl %ecx, %eax" "\n"
+    "cmpl %eax, %edx" "\n"
+    "jge " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) "\n"
+
+    // Allocate a safe place on the stack below the result stack pointer to stash the ProbeContext.
+    "subl %ecx, %edx" "\n"
+    "andl $~0x1f, %edx" "\n" // Keep the stack pointer 32 bytes aligned.
+    "xorl %eax, %eax" "\n"
+    "movl %edx, %esp" "\n"
+
+    "movl $" STRINGIZE_VALUE_OF(PROBE_SIZE) ", %ecx" "\n"
+
+    // Copy the ProbeContext to the safe place.
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) ":" "\n"
+    "movl (%ebp, %eax), %edx" "\n"
+    "movl %edx, (%esp, %eax)" "\n"
+    "addl $" STRINGIZE_VALUE_OF(PTR_SIZE) ", %eax" "\n"
+    "cmpl %eax, %ecx" "\n"
+    "jg " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) "\n"
+
+    "movl %esp, %ebp" "\n"
+
+    // Call initializeStackFunction if present.
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) ":" "\n"
+    "xorl %ecx, %ecx" "\n"
+    "addl " STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "(%ebp), %ecx" "\n"
+    "je " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) "\n"
+
+    // Reserve stack space for the arg while maintaining the required stack
+    // pointer 32 byte alignment:
+    "subl $0x20, %esp" "\n"
+    "movl %ebp, 0(%esp)" "\n" // the ProbeContext* arg.
+    "call *%ecx" "\n"
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) ":" "\n"
+
     // To enable probes to modify register state, we copy all registers
     // out of the ProbeContext before returning.
 
@@ -249,59 +298,10 @@ asm (
 
     // ecx now points to the restore area.
 
-    // Before we copy values from the ProbeContext to the restore area, we need to
-    // make sure that the restore area does not overlap any of the values that we'll
-    // be copying from in the ProbeContext. All the restore values to be copied from
-    // comes from offset <= PROBE_CPU_EFLAGS_OFFSET in the ProbeContext.
-    "movl %ebp, %eax" "\n"
-    "addl $" STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) ", %eax" "\n"
-    "cmpl %eax, %ecx" "\n"
-    "jg " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
-
-    // Getting here means that the restore area will overlap the ProbeContext data
-    // that we will need to get the restoration values from. So, let's move that
-    // data to a safe place before we start writing into the restore area.
-    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
-    // restore area. This ensures that:
-    // 1. The safe area does not overlap the restore area.
-    // 2. The safe area does not overlap the ProbeContext.
-    //    This makes it so that we can use memcpy (does not require memmove) semantics
-    //    to copy the restore values to the safe area.
-    // Note: the safe area does not have to 32-byte align it because we're not using
-    // it to store any xmm regs.
-    "movl %ecx, %eax" "\n"
-    "subl $2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ", %eax" "\n"
-
-    // eax now points to the safe area.
-
-    // Make sure the stack pointer points to the safe area. This ensures that the
-    // safe area is protected from interrupt handlers overwriting it.
-    "movl %eax, %esp" "\n"
-
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%eax)" "\n"
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%ebp), %ecx" "\n"
-    "movl %ecx, " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%eax)" "\n"
-    "movl %eax, %ebp" "\n"
-
-    // We used ecx above as scratch register. Let's restore it to points to the
-    // restore area.
-    "movl " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%ebp), %ecx" "\n"
-    "subl $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %ecx" "\n"
-
-    // ecx now points to the restore area.
-
-    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
-
     // Copy remaining restore values from the ProbeContext to the restore area.
+    // Note: We already ensured above that the ProbeContext is in a safe location before
+    // calling the initializeStackFunction. The initializeStackFunction is not allowed to
+    // change the stack pointer again.
     "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%ebp), %eax" "\n"
     "movl %eax, 0 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%ecx)" "\n"
     "movl " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%ebp), %eax" "\n"
@@ -342,11 +342,11 @@ asm (
     //     esp[5 * ptrSize]: saved rsp
 
     "movq %rsp, %rax" "\n"
-    "subq $" STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ", %rsp" "\n"
+    "subq $" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) ", %rsp" "\n"
 
-    // The X86_64 ABI specifies that the worse case stack alignment requirement
-    // is 32 bytes.
+    // The X86_64 ABI specifies that the worse case stack alignment requirement is 32 bytes.
     "andq $~0x1f, %rsp" "\n"
+    // Since sp points to the ProbeContext, we've ensured that it's protected from interrupts before we initialize it.
 
     "movq %rbp, " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%rsp)" "\n"
     "movq %rsp, %rbp" "\n" // Save the ProbeContext*.
@@ -396,9 +396,51 @@ asm (
     "movq %xmm14, " STRINGIZE_VALUE_OF(PROBE_CPU_XMM14_OFFSET) "(%rbp)" "\n"
     "movq %xmm15, " STRINGIZE_VALUE_OF(PROBE_CPU_XMM15_OFFSET) "(%rbp)" "\n"
 
+    "xorq %rax, %rax" "\n"
+    "movq %rax, " STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "(%rbp)" "\n"
+
     "movq %rbp, %rdi" "\n" // the ProbeContext* arg.
     "call *" STRINGIZE_VALUE_OF(PROBE_PROBE_FUNCTION_OFFSET) "(%rbp)" "\n"
 
+    // Make sure the ProbeContext is entirely below the result stack pointer so
+    // that register values are still preserved when we call the initializeStack
+    // function.
+    "movq $" STRINGIZE_VALUE_OF(PROBE_SIZE + OUT_SIZE) ", %rcx" "\n"
+    "movq %rbp, %rax" "\n"
+    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp), %rdx" "\n"
+    "addq %rcx, %rax" "\n"
+    "cmpq %rax, %rdx" "\n"
+    "jge " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) "\n"
+    // Allocate a safe place on the stack below the result stack pointer to stash the ProbeContext.
+    "subq %rcx, %rdx" "\n"
+    "andq $~0x1f, %rdx" "\n" // Keep the stack pointer 32 bytes aligned.
+    "xorq %rax, %rax" "\n"
+    "movq %rdx, %rsp" "\n"
+
+    "movq $" STRINGIZE_VALUE_OF(PROBE_SIZE) ", %rcx" "\n"
+
+    // Copy the ProbeContext to the safe place.
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) ":" "\n"
+    "movq (%rbp, %rax), %rdx" "\n"
+    "movq %rdx, (%rsp, %rax)" "\n"
+    "addq $" STRINGIZE_VALUE_OF(PTR_SIZE) ", %rax" "\n"
+    "cmpq %rax, %rcx" "\n"
+    "jg " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineCopyLoop) "\n"
+
+    "movq %rsp, %rbp" "\n"
+
+    // Call initializeStackFunction if present.
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineProbeContextIsSafe) ":" "\n"
+    "xorq %rcx, %rcx" "\n"
+    "addq " STRINGIZE_VALUE_OF(PROBE_INIT_STACK_FUNCTION_OFFSET) "(%rbp), %rcx" "\n"
+    "je " LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) "\n"
+
+    "movq %rbp, %rdi" "\n" // the ProbeContext* arg.
+    "call *%rcx" "\n"
+
+    LOCAL_LABEL_STRING(ctiMasmProbeTrampolineRestoreRegisters) ":" "\n"
+
     // To enable probes to modify register state, we copy all registers
     // out of the ProbeContext before returning.
 
@@ -446,59 +488,10 @@ asm (
 
     // rcx now points to the restore area.
 
-    // Before we copy values from the ProbeContext to the restore area, we need to
-    // make sure that the restore area does not overlap any of the values that we'll
-    // be copying from in the ProbeContext. All the restore values to be copied from
-    // comes from offset <= PROBE_CPU_EFLAGS_OFFSET in the ProbeContext.
-    "movq %rbp, %rax" "\n"
-    "addq $" STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) ", %rax" "\n"
-    "cmpq %rax, %rcx" "\n"
-    "jg " SYMBOL_STRING(ctiMasmProbeTrampolineEnd) "\n"
-
-    // Getting here means that the restore area will overlap the ProbeContext data
-    // that we will need to get the restoration values from. So, let's move that
-    // data to a safe place before we start writing into the restore area.
-    // Let's locate the "safe area" at 2x sizeof(ProbeContext) below where the
-    // restore area. This ensures that:
-    // 1. The safe area does not overlap the restore area.
-    // 2. The safe area does not overlap the ProbeContext.
-    //    This makes it so that we can use memcpy (does not require memmove) semantics
-    //    to copy the restore values to the safe area.
-    // Note: the safe area does not have to 32-byte align it because we're not using
-    // it to store any xmm regs.
-    "movq %rcx, %rax" "\n"
-    "subq $2 * " STRINGIZE_VALUE_OF(PROBE_ALIGNED_SIZE) ", %rax" "\n"
-
-    // rax now points to the safe area.
-
-    // Make sure the stack pointer points to the safe area. This ensures that the
-    // safe area is protected from interrupt handlers overwriting it.
-    "movq %rax, %rsp" "\n"
-
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, " STRINGIZE_VALUE_OF(PROBE_CPU_ECX_OFFSET) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, " STRINGIZE_VALUE_OF(PROBE_CPU_EBP_OFFSET) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, " STRINGIZE_VALUE_OF(PROBE_CPU_EIP_OFFSET) "(%rax)" "\n"
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%rbp), %rcx" "\n"
-    "movq %rcx, " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%rax)" "\n"
-    "movq %rax, %rbp" "\n"
-
-    // We used rcx above as scratch register. Let's restore it to points to the
-    // restore area.
-    "movq " STRINGIZE_VALUE_OF(PROBE_CPU_ESP_OFFSET) "(%rbp), %rcx" "\n"
-    "subq $5 * " STRINGIZE_VALUE_OF(PTR_SIZE) ", %rcx" "\n"
-
-    // rcx now points to the restore area.
-
-    SYMBOL_STRING(ctiMasmProbeTrampolineEnd) ":" "\n"
-
     // Copy remaining restore values from the ProbeContext to the restore area.
+    // Note: We already ensured above that the ProbeContext is in a safe location before
+    // calling the initializeStackFunction. The initializeStackFunction is not allowed to
+    // change the stack pointer again.
     "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EFLAGS_OFFSET) "(%rbp), %rax" "\n"
     "movq %rax, 0 * " STRINGIZE_VALUE_OF(PTR_SIZE) "(%rcx)" "\n"
     "movq " STRINGIZE_VALUE_OF(PROBE_CPU_EAX_OFFSET) "(%rbp), %rax" "\n"
index 825dd83..80c05f7 100644 (file)
@@ -63,6 +63,8 @@ using namespace JSC;
 
 namespace {
 
+using CPUState = MacroAssembler::CPUState;
+
 StaticLock crashLock;
 
 typedef WTF::Function<void(CCallHelpers&)> Generator;
@@ -267,7 +269,7 @@ void testProbePreservesGPRS()
     // having already validated that we can read and write from registers. We'll use these abilities
     // to validate that the probe preserves register values.
     unsigned probeCallCount = 0;
-    MacroAssembler::CPUState originalState;
+    CPUState originalState;
 
     compileAndRun<void>([&] (CCallHelpers& jit) {
         jit.emitFunctionPrologue();
@@ -347,7 +349,7 @@ void testProbePreservesGPRS()
 void testProbeModifiesStackPointer(WTF::Function<void*(ProbeContext*)> computeModifiedStack)
 {
     unsigned probeCallCount = 0;
-    MacroAssembler::CPUState originalState;
+    CPUState originalState;
     void* originalSP { nullptr };
     void* modifiedSP { nullptr };
     uintptr_t modifiedFlags { 0 };
@@ -508,6 +510,158 @@ void testProbeModifiesProgramCounter()
     CHECK_EQ(continuationWasReached, true);
 }
 
+struct FillStackData {
+    CPUState originalState;
+    void* originalSP { nullptr };
+    void* newSP { nullptr };
+    uintptr_t modifiedFlags { 0 };
+    MacroAssembler::SPRegisterID flagsSPR;
+};
+
+static void fillStack(ProbeContext* context)
+{
+    auto& cpu = context->cpu;
+
+    FillStackData& data = *reinterpret_cast<FillStackData*>(context->initializeStackArg);
+    CPUState& originalState = data.originalState;
+    void*& originalSP = data.originalSP;
+    void*& newSP = data.newSP;
+    uintptr_t& modifiedFlags = data.modifiedFlags;
+    MacroAssembler::SPRegisterID& flagsSPR = data.flagsSPR;
+
+    CHECK_EQ(reinterpret_cast<void*>(context->initializeStackFunction), reinterpret_cast<void*>(fillStack));
+    CHECK_EQ(cpu.sp(), newSP);
+
+    // Verify that the probe has put the ProbeContext out of harm's way.
+    CHECK_EQ((reinterpret_cast<void*>(context + 1) <= cpu.sp()), true);
+
+    // Verify the CPU state.
+    for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+        if (isFP(id)) {
+            CHECK_EQ(cpu.gpr(id), originalState.gpr(id));
+            continue;
+        }
+        if (isSpecialGPR(id))
+            continue;
+        CHECK_EQ(cpu.gpr(id), testWord(id));
+    }
+    for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
+        CHECK_EQ(cpu.fpr<uint64_t>(id), testWord64(id));
+    CHECK_EQ(cpu.spr(flagsSPR), modifiedFlags);
+
+    // Fill the stack with values.
+    uintptr_t* p = reinterpret_cast<uintptr_t*>(newSP);
+    int count = 0;
+    while (p < reinterpret_cast<uintptr_t*>(originalSP))
+        *p++ = testWord(count++);
+};
+
+void testProbeModifiesStackWithCallback()
+{
+    unsigned probeCallCount = 0;
+    FillStackData data;
+    CPUState& originalState = data.originalState;
+    void*& originalSP = data.originalSP;
+    void*& newSP = data.newSP;
+    uintptr_t& modifiedFlags = data.modifiedFlags;
+    size_t numberOfExtraEntriesToWrite = 10; // ARM64 requires that this be 2 word aligned.
+    MacroAssembler::SPRegisterID& flagsSPR = data.flagsSPR;
+
+#if CPU(X86) || CPU(X86_64)
+    flagsSPR = X86Registers::eflags;
+    uintptr_t flagsMask = 0xc5;
+#elif CPU(ARM_THUMB2) || CPU(ARM_TRADITIONAL)
+    flagsSPR = ARMRegisters::apsr;
+    uintptr_t flagsMask = 0xf0000000;
+#elif CPU(ARM64)
+    flagsSPR = ARM64Registers::nzcv;
+    uintptr_t flagsMask = 0xf0000000;
+#endif
+
+    compileAndRun<void>([&] (CCallHelpers& jit) {
+        jit.emitFunctionPrologue();
+
+        // Write expected values into the registers.
+        jit.probe([&] (ProbeContext* context) {
+            auto& cpu = context->cpu;
+            probeCallCount++;
+
+            // Preserve the original CPU state.
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                originalState.gpr(id) = cpu.gpr(id);
+                if (isSpecialGPR(id))
+                    continue;
+                cpu.gpr(id) = testWord(static_cast<int>(id));
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id)) {
+                originalState.fpr(id) = cpu.fpr(id);
+                cpu.fpr(id) = bitwise_cast<double>(testWord64(id));
+            }
+            originalState.spr(flagsSPR) = cpu.spr(flagsSPR);
+            modifiedFlags = originalState.spr(flagsSPR) ^ flagsMask;
+            cpu.spr(flagsSPR) = modifiedFlags;
+
+            CHECK_EQ(reinterpret_cast<void*>(context->initializeStackFunction), 0);
+
+            // Prepare for initializeStack callback.
+            context->initializeStackFunction = fillStack;
+            context->initializeStackArg = &data;
+
+            // Ensure that we'll be writing over the regions of the stack where the ProbeContext is.
+            originalSP = cpu.sp();
+            newSP = reinterpret_cast<uintptr_t*>(context) - numberOfExtraEntriesToWrite;
+            cpu.sp() = newSP;
+        });
+
+        // Validate that the registers and stack have the expected values.
+        jit.probe([&] (ProbeContext* context) {
+            auto& cpu = context->cpu;
+            probeCallCount++;
+
+            // Validate the register values.
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                if (isFP(id)) {
+                    CHECK_EQ(cpu.gpr(id), originalState.gpr(id));
+                    continue;
+                }
+                if (isSpecialGPR(id))
+                    continue;
+                CHECK_EQ(cpu.gpr(id), testWord(id));
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
+                CHECK_EQ(cpu.fpr<uint64_t>(id), testWord64(id));
+            CHECK_EQ(cpu.spr(flagsSPR), modifiedFlags);
+            CHECK_EQ(cpu.sp(), newSP);
+
+            // Validate the stack with values.
+            uintptr_t* p = reinterpret_cast<uintptr_t*>(newSP);
+            int count = 0;
+            while (p < reinterpret_cast<uintptr_t*>(originalSP))
+                CHECK_EQ(*p++, testWord(count++));
+        });
+
+        // Restore the original state.
+        jit.probe([&] (ProbeContext* context) {
+            auto& cpu = context->cpu;
+            probeCallCount++;
+            for (auto id = CCallHelpers::firstRegister(); id <= CCallHelpers::lastRegister(); id = nextID(id)) {
+                if (isSpecialGPR(id))
+                    continue;
+                cpu.gpr(id) = originalState.gpr(id);
+            }
+            for (auto id = CCallHelpers::firstFPRegister(); id <= CCallHelpers::lastFPRegister(); id = nextID(id))
+                cpu.fpr(id) = originalState.fpr(id);
+            cpu.spr(flagsSPR) = originalState.spr(flagsSPR);
+            cpu.sp() = originalSP;
+        });
+
+        jit.emitFunctionEpilogue();
+        jit.ret();
+    });
+
+    CHECK_EQ(probeCallCount, 3);
+}
+
 #define RUN(test) do {                          \
         if (!shouldRun(#test))                  \
             break;                              \
@@ -540,6 +694,7 @@ void run(const char* filter)
     RUN(testProbeModifiesStackPointerToInsideProbeContextOnStack());
     RUN(testProbeModifiesStackPointerToNBytesBelowSP());
     RUN(testProbeModifiesProgramCounter());
+    RUN(testProbeModifiesStackWithCallback());
 
     if (tasks.isEmpty())
         usage();