1 files changed, 81 insertions, 27 deletions

diff --git a/src/pal/src/thread/context.cpp b/src/pal/src/thread/context.cpp
index f832015710..0449df568b 100644
--- a/src/pal/src/thread/context.cpp
+++ b/src/pal/src/thread/context.cpp
@@ -33,9 +33,7 @@ SET_DEFAULT_DEBUG_CHANNEL(THREAD); // some headers have code with asserts, so do
 
 extern PGET_GCMARKER_EXCEPTION_CODE g_getGcMarkerExceptionCode;
 
-// in context2.S
-extern void CONTEXT_CaptureContext(LPCONTEXT lpContext);
-
+#define CONTEXT_AREA_MASK 0xffff
 #ifdef _X86_
 #define CONTEXT_ALL_FLOATING (CONTEXT_FLOATING_POINT | CONTEXT_EXTENDED_REGISTERS)
 #elif defined(_AMD64_)
@@ -288,7 +286,7 @@ CONTEXT_GetThreadContext(
     }
 
     if (lpContext->ContextFlags & 
-        (CONTEXT_CONTROL | CONTEXT_INTEGER))
+        (CONTEXT_CONTROL | CONTEXT_INTEGER) & CONTEXT_AREA_MASK)
     {        
         if (CONTEXT_GetRegisters(dwProcessId, lpContext) == FALSE)
         {
@@ -348,7 +346,7 @@ CONTEXT_SetThreadContext(
     }
     
     if (lpContext->ContextFlags  & 
-        (CONTEXT_CONTROL | CONTEXT_INTEGER))
+        (CONTEXT_CONTROL | CONTEXT_INTEGER) & CONTEXT_AREA_MASK)
     {   
 #if HAVE_PT_REGS
         if (ptrace((__ptrace_request)PT_GETREGS, dwProcessId, (caddr_t)&ptrace_registers, 0) == -1)
@@ -371,11 +369,11 @@ CONTEXT_SetThreadContext(
 	ASSERT("Don't know how to set the context of another process on this platform!");
 	return FALSE;
 #endif
-        if (lpContext->ContextFlags & CONTEXT_CONTROL)
+        if (lpContext->ContextFlags & CONTEXT_CONTROL & CONTEXT_AREA_MASK)
         {
             ASSIGN_CONTROL_REGS
         }
-        if (lpContext->ContextFlags & CONTEXT_INTEGER)
+        if (lpContext->ContextFlags & CONTEXT_INTEGER & CONTEXT_AREA_MASK)
         {
             ASSIGN_INTEGER_REGS
         }
@@ -467,13 +465,13 @@ void CONTEXTToNativeContext(CONST CONTEXT *lpContext, native_context_t *native)
     }
 
     // TODO: Enable for all Unix systems
-#if defined(_AMD64_) && defined(__linux__)
+#if defined(_AMD64_) && defined(XSTATE_SUPPORTED)
     if ((lpContext->ContextFlags & CONTEXT_XSTATE) == CONTEXT_XSTATE)
     {
         _ASSERTE(FPREG_HasExtendedState(native));
         memcpy_s(FPREG_Xstate_Ymmh(native), sizeof(M128A) * 16, lpContext->VectorRegister, sizeof(M128A) * 16);
     }
-#endif // _AMD64_
+#endif //_AMD64_ && XSTATE_SUPPORTED
 }
 
 /*++
@@ -564,22 +562,24 @@ void CONTEXTFromNativeContext(const native_context_t *native, LPCONTEXT lpContex
 #endif
     }
 
-    // TODO: Enable for all Unix systems
-#if defined(_AMD64_) && defined(__linux__)
+#ifdef _AMD64_
     if ((contextFlags & CONTEXT_XSTATE) == CONTEXT_XSTATE)
     {
+    // TODO: Enable for all Unix systems
+#if XSTATE_SUPPORTED
         if (FPREG_HasExtendedState(native))
         {
             memcpy_s(lpContext->VectorRegister, sizeof(M128A) * 16, FPREG_Xstate_Ymmh(native), sizeof(M128A) * 16);
         }
         else
+#endif // XSTATE_SUPPORTED
         {
             // Reset the CONTEXT_XSTATE bit(s) so it's clear that the extended state data in
             // the CONTEXT is not valid.
             const ULONG xstateFlags = CONTEXT_XSTATE & ~(CONTEXT_CONTROL & CONTEXT_INTEGER);
             lpContext->ContextFlags &= ~xstateFlags;
         }
-    } 
+    }
 #endif // _AMD64_
 }
 
@@ -855,7 +855,7 @@ CONTEXT_GetThreadContextFromPort(
     mach_msg_type_number_t StateCount;
     thread_state_flavor_t StateFlavor;
   
-    if (lpContext->ContextFlags & (CONTEXT_CONTROL | CONTEXT_INTEGER | CONTEXT_SEGMENTS))
+    if (lpContext->ContextFlags & (CONTEXT_CONTROL | CONTEXT_INTEGER | CONTEXT_SEGMENTS) & CONTEXT_AREA_MASK)
     {
 #ifdef _X86_  
         x86_thread_state32_t State;
@@ -877,7 +877,7 @@ CONTEXT_GetThreadContextFromPort(
         CONTEXT_GetThreadContextFromThreadState(StateFlavor, (thread_state_t)&State, lpContext);
     }
     
-    if (lpContext->ContextFlags & CONTEXT_ALL_FLOATING) {
+    if (lpContext->ContextFlags & CONTEXT_ALL_FLOATING & CONTEXT_AREA_MASK) {
 #ifdef _X86_
         x86_float_state32_t State;
         StateFlavor = x86_FLOAT_STATE32;
@@ -898,6 +898,22 @@ CONTEXT_GetThreadContextFromPort(
         CONTEXT_GetThreadContextFromThreadState(StateFlavor, (thread_state_t)&State, lpContext);
     }
 
+#if defined(_AMD64_) && defined(XSTATE_SUPPORTED)
+    if (lpContext->ContextFlags & CONTEXT_XSTATE & CONTEXT_AREA_MASK) {
+        x86_avx_state64_t State;
+        StateFlavor = x86_AVX_STATE64;
+        StateCount = sizeof(State) / sizeof(natural_t);
+        MachRet = thread_get_state(Port, StateFlavor, (thread_state_t)&State, &StateCount);
+        if (MachRet != KERN_SUCCESS)
+        {
+            ASSERT("thread_get_state(XSTATE) failed: %d\n", MachRet);
+            goto exit;
+        }
+
+        CONTEXT_GetThreadContextFromThreadState(StateFlavor, (thread_state_t)&State, lpContext);
+    }
+#endif
+
 exit:
     return MachRet;
 }
@@ -917,7 +933,7 @@ CONTEXT_GetThreadContextFromThreadState(
     {
 #ifdef _X86_
         case x86_THREAD_STATE32:
-            if (lpContext->ContextFlags & (CONTEXT_CONTROL | CONTEXT_INTEGER | CONTEXT_SEGMENTS))
+            if (lpContext->ContextFlags & (CONTEXT_CONTROL | CONTEXT_INTEGER | CONTEXT_SEGMENTS) & CONTEXT_AREA_MASK)
             {
                 x86_thread_state32_t *pState = (x86_thread_state32_t *)threadState;
 
@@ -944,7 +960,7 @@ CONTEXT_GetThreadContextFromThreadState(
         {
             x86_float_state32_t *pState = (x86_float_state32_t *)threadState;
 
-            if (lpContext->ContextFlags & CONTEXT_FLOATING_POINT)
+            if (lpContext->ContextFlags & CONTEXT_FLOATING_POINT & CONTEXT_AREA_MASK)
             {
                 lpContext->FloatSave.ControlWord = *(DWORD*)&pState->fpu_fcw;
                 lpContext->FloatSave.StatusWord = *(DWORD*)&pState->fpu_fsw;
@@ -963,7 +979,7 @@ CONTEXT_GetThreadContextFromThreadState(
                     memcpy(&lpContext->FloatSave.RegisterArea[i * 10], (&pState->fpu_stmm0)[i].mmst_reg, 10);
             }
 
-            if (lpContext->ContextFlags & CONTEXT_EXTENDED_REGISTERS)
+            if (lpContext->ContextFlags & CONTEXT_EXTENDED_REGISTERS & CONTEXT_AREA_MASK)
             {
                 // The only extended register information that Mach will tell us about are the xmm register values.
                 // Both Windows and Mach store the registers in a packed layout (each of the 8 registers is 16 bytes)
@@ -975,7 +991,7 @@ CONTEXT_GetThreadContextFromThreadState(
 
 #elif defined(_AMD64_)
         case x86_THREAD_STATE64:
-            if (lpContext->ContextFlags & (CONTEXT_CONTROL | CONTEXT_INTEGER | CONTEXT_SEGMENTS))
+            if (lpContext->ContextFlags & (CONTEXT_CONTROL | CONTEXT_INTEGER | CONTEXT_SEGMENTS) & CONTEXT_AREA_MASK)
             {
                 x86_thread_state64_t *pState = (x86_thread_state64_t *)threadState;
 
@@ -1009,7 +1025,7 @@ CONTEXT_GetThreadContextFromThreadState(
             break;
 
         case x86_FLOAT_STATE64:
-            if (lpContext->ContextFlags & CONTEXT_FLOATING_POINT)
+            if (lpContext->ContextFlags & CONTEXT_FLOATING_POINT & CONTEXT_AREA_MASK)
             {
                 x86_float_state64_t *pState = (x86_float_state64_t *)threadState;
 
@@ -1031,9 +1047,19 @@ CONTEXT_GetThreadContextFromThreadState(
                     memcpy(&lpContext->FltSave.FloatRegisters[i], (&pState->__fpu_stmm0)[i].__mmst_reg, 10);
 
                 // AMD64's FLOATING_POINT includes the xmm registers.
-                memcpy(&lpContext->Xmm0, &pState->__fpu_xmm0, 8 * 16);
+                memcpy(&lpContext->Xmm0, &pState->__fpu_xmm0, 16 * 16);
+            }
+            break;
+
+#ifdef XSTATE_SUPPORTED
+        case x86_AVX_STATE64:
+            if (lpContext->ContextFlags & CONTEXT_XSTATE & CONTEXT_AREA_MASK)
+            {
+                x86_avx_state64_t *pState = (x86_avx_state64_t *)threadState;
+                memcpy(&lpContext->VectorRegister, &pState->__fpu_ymmh0, 16 * 16);
             }
             break;
+#endif
 #else
 #error Unexpected architecture.
 #endif
@@ -1120,7 +1146,7 @@ CONTEXT_SetThreadContextOnPort(
     mach_msg_type_number_t StateCount;
     thread_state_flavor_t StateFlavor;
 
-    if (lpContext->ContextFlags & (CONTEXT_CONTROL|CONTEXT_INTEGER)) 
+    if (lpContext->ContextFlags & (CONTEXT_CONTROL|CONTEXT_INTEGER) & CONTEXT_AREA_MASK)
     {
 #ifdef _X86_
         x86_thread_state32_t State;
@@ -1187,21 +1213,42 @@ CONTEXT_SetThreadContextOnPort(
         }
     }
 
-    if (lpContext->ContextFlags & CONTEXT_ALL_FLOATING)
+    if (lpContext->ContextFlags & CONTEXT_ALL_FLOATING & CONTEXT_AREA_MASK)
     {
         
 #ifdef _X86_
         x86_float_state32_t State;
         StateFlavor = x86_FLOAT_STATE32;
+        StateCount = sizeof(State) / sizeof(natural_t);
 #elif defined(_AMD64_)
+#ifdef XSTATE_SUPPORTED
+        // We're relying on the fact that the initial portion of
+        // x86_avx_state64_t is identical to x86_float_state64_t.
+        // Check a few fields to make sure the assumption is correct.
+        static_assert_no_msg(sizeof(x86_avx_state64_t) > sizeof(x86_float_state64_t));
+        static_assert_no_msg(offsetof(x86_avx_state64_t, __fpu_fcw) == offsetof(x86_float_state64_t, __fpu_fcw));
+        static_assert_no_msg(offsetof(x86_avx_state64_t, __fpu_xmm0) == offsetof(x86_float_state64_t, __fpu_xmm0));
+
+        x86_avx_state64_t State;
+        if (lpContext->ContextFlags & CONTEXT_XSTATE & CONTEXT_AREA_MASK)
+        {
+            StateFlavor = x86_AVX_STATE64;
+            StateCount = sizeof(State) / sizeof(natural_t);
+        }
+        else
+        {
+            StateFlavor = x86_FLOAT_STATE64;
+            StateCount = sizeof(x86_float_state64_t) / sizeof(natural_t);
+        }
+#else
         x86_float_state64_t State;
         StateFlavor = x86_FLOAT_STATE64;
+        StateCount = sizeof(State) / sizeof(natural_t);
+#endif
 #else
 #error Unexpected architecture.
 #endif
 
-        StateCount = sizeof(State) / sizeof(natural_t);
-
         // If we're setting only one of the floating point or extended registers (of which Mach supports only
         // the xmm values) then we don't have values for the other set. This is a problem since Mach only
         // supports setting both groups as a single unit. So in this case we'll need to fetch the current
@@ -1222,7 +1269,7 @@ CONTEXT_SetThreadContextOnPort(
             _ASSERTE(StateCountGet == StateCount);
         }
 
-        if (lpContext->ContextFlags & CONTEXT_FLOATING_POINT)
+        if (lpContext->ContextFlags & CONTEXT_FLOATING_POINT & CONTEXT_AREA_MASK)
         {
 #ifdef _X86_
             *(DWORD*)&State.fpu_fcw = lpContext->FloatSave.ControlWord;
@@ -1258,14 +1305,14 @@ CONTEXT_SetThreadContextOnPort(
             for (int i = 0; i < 8; i++)
                 memcpy((&State.__fpu_stmm0)[i].__mmst_reg, &lpContext->FltSave.FloatRegisters[i], 10);
 
-            memcpy(&State.__fpu_xmm0, &lpContext->Xmm0, 8 * 16);
+            memcpy(&State.__fpu_xmm0, &lpContext->Xmm0, 16 * 16);
 #else
 #error Unexpected architecture.
 #endif
         }
 
 #ifdef _X86_
-        if (lpContext->ContextFlags & CONTEXT_EXTENDED_REGISTERS)
+        if (lpContext->ContextFlags & CONTEXT_EXTENDED_REGISTERS & CONTEXT_AREA_MASK)
         {
             // The only extended register information that Mach will tell us about are the xmm register
             // values. Both Windows and Mach store the registers in a packed layout (each of the 8 registers
@@ -1274,6 +1321,13 @@ CONTEXT_SetThreadContextOnPort(
         }
 #endif // _X86_
 
+#if defined(_AMD64_) && defined(XSTATE_SUPPORTED)
+        if (lpContext->ContextFlags & CONTEXT_XSTATE & CONTEXT_AREA_MASK)
+        {
+            memcpy(&State.__fpu_ymmh0, lpContext->VectorRegister, 16 * 16);
+        }
+#endif
+
         MachRet = thread_set_state(Port,
                                    StateFlavor,
                                    (thread_state_t)&State,