Repository: dotnet/runtime
File: src/coreclr/nativeaot/Runtime/arm64/GcProbe.S (Assembly)

Code before:

    // Restore callee saved registers
    EPILOG_RESTORE_REG_PAIR x19, x20, 0x20
    EPILOG_RESTORE_REG_PAIR x21, x22, 0x30
    EPILOG_RESTORE_REG_PAIR x23, x24, 0x40
    EPILOG_RESTORE_REG_PAIR x25, x26, 0x50
    EPILOG_RESTORE_REG_PAIR x27, x28, 0x60

    EPILOG_RESTORE_REG_PAIR_INDEXED   fp, lr, PROBE_FRAME_SIZE
.endm

//
// The prolog for all GC suspension hijacks (normal and stress). Fixes up the hijacked return address, and
// clears the hijack state.
//
// Register state on entry:
//  All registers correct for return to the original return address.
//
// Register state on exit:
//  x4: thread pointer
//  x2: preserved
//
.macro FixupHijackedCallstack

    // Preserve x2 since it may contain an async continuation
    str         x2, [sp, #-16]!

    // x4 <- GetThread()
#ifdef FEATURE_EMULATED_TLS
    GETTHREAD_ETLS_4
#else
    INLINE_GETTHREAD x4
#endif

    //
    // Fix the stack by restoring the original return address
    //
    ldr         lr, [x4, #OFFSETOF__Thread__m_pvHijackedReturnAddress]

    //
    // Clear hijack state
    //
    // Clear m_ppvHijackedReturnAddressLocation and m_pvHijackedReturnAddress
    stp         xzr, xzr, [x4, #OFFSETOF__Thread__m_ppvHijackedReturnAddressLocation]

    // Restore x2
    ldr         x2, [sp], #16
.endm

//
// GC Probe Hijack target

Review comment (suggestion):
It looks like `INLINE_GETTHREAD` does not trash anything except what it uses for the return value. It trashes x0 internally, but takes care to push/pop it. 

The `FEATURE_EMULATED_TLS` ends up making a real function call (`RhpGetThread`) which formally trashes all scratch registers. Thus GETTHREAD_ETLS_4 pushes/pops x0, x1 around the call.  

I think saving/restoring x2 is only relevant on FEATURE_EMULATED_TLS path where registers are trashed, so it makes more sense to do save/restore inside GETTHREAD_ETLS_4. 
I.E - the `INLINE_GETTHREAD` does not save/restore x1, so I assume the functionality that we use does not trash x2 either. Thus no need to save/restore x2 on that code path.



Code after:

    // Restore callee saved registers
    EPILOG_RESTORE_REG_PAIR x19, x20, 0x20
    EPILOG_RESTORE_REG_PAIR x21, x22, 0x30
    EPILOG_RESTORE_REG_PAIR x23, x24, 0x40
    EPILOG_RESTORE_REG_PAIR x25, x26, 0x50
    EPILOG_RESTORE_REG_PAIR x27, x28, 0x60

    EPILOG_RESTORE_REG_PAIR_INDEXED   fp, lr, PROBE_FRAME_SIZE
.endm

//
// The prolog for all GC suspension hijacks (normal and stress). Fixes up the hijacked return address, and
// clears the hijack state.
//
// Register state on entry:
//  All registers correct for return to the original return address.
//
// Register state on exit:
//  x4: thread pointer
//  x0, x1, x2: preserved
//
.macro FixupHijackedCallstack

    // Store thread pointer temporarily in x3, then move it to x4
    // x4 <- GetThread()

#ifdef FEATURE_EMULATED_TLS
    GETTHREAD_ETLS_3
#else
    INLINE_GETTHREAD x3
#endif

    mov         x4, x3

    //
    // Fix the stack by restoring the original return address
    //
    ldr         lr, [x4, #OFFSETOF__Thread__m_pvHijackedReturnAddress]

    //
    // Clear hijack state
    //
    // Clear m_ppvHijackedReturnAddressLocation and m_pvHijackedReturnAddress
    stp         xzr, xzr, [x4, #OFFSETOF__Thread__m_ppvHijackedReturnAddressLocation]
.endm

//
// GC Probe Hijack target
//
NESTED_ENTRY RhpGcProbeHijack, _TEXT, NoHandler