Add extended state to RtlRestoreContext

2 files changed, 35 insertions, 3 deletions

diff --git a/src/pal/src/arch/i386/context2.S b/src/pal/src/arch/i386/context2.S
index 6320446a51..0e93e81a55 100644
--- a/src/pal/src/arch/i386/context2.S
+++ b/src/pal/src/arch/i386/context2.S
@@ -126,6 +126,28 @@ LOCAL_LABEL(Done_Restore_CONTEXT_DEBUG_REGISTERS):
     fxrstor [rdi + CONTEXT_FltSave]
 LOCAL_LABEL(Done_Restore_CONTEXT_FLOATING_POINT):
 
+    test    BYTE PTR [rdi + CONTEXT_ContextFlags], CONTEXT_XSTATE
+    je      LOCAL_LABEL(Done_Restore_CONTEXT_XSTATE)
+
+    // Restore the extended state (for now, this is just the upper halves of YMM registers)
+    vinsertf128 ymm0, ymm0, xmmword ptr [rdi + (CONTEXT_VectorRegister + 0 * 16)], 1
+    vinsertf128 ymm1, ymm1, xmmword ptr [rdi + (CONTEXT_VectorRegister + 1 * 16)], 1
+    vinsertf128 ymm2, ymm2, xmmword ptr [rdi + (CONTEXT_VectorRegister + 2 * 16)], 1
+    vinsertf128 ymm3, ymm3, xmmword ptr [rdi + (CONTEXT_VectorRegister + 3 * 16)], 1
+    vinsertf128 ymm4, ymm4, xmmword ptr [rdi + (CONTEXT_VectorRegister + 4 * 16)], 1
+    vinsertf128 ymm5, ymm5, xmmword ptr [rdi + (CONTEXT_VectorRegister + 5 * 16)], 1
+    vinsertf128 ymm6, ymm6, xmmword ptr [rdi + (CONTEXT_VectorRegister + 6 * 16)], 1
+    vinsertf128 ymm7, ymm7, xmmword ptr [rdi + (CONTEXT_VectorRegister + 7 * 16)], 1
+    vinsertf128 ymm8, ymm8, xmmword ptr [rdi + (CONTEXT_VectorRegister + 8 * 16)], 1
+    vinsertf128 ymm9, ymm9, xmmword ptr [rdi + (CONTEXT_VectorRegister + 9 * 16)], 1
+    vinsertf128 ymm10, ymm10, xmmword ptr [rdi + (CONTEXT_VectorRegister + 10 * 16)], 1
+    vinsertf128 ymm11, ymm11, xmmword ptr [rdi + (CONTEXT_VectorRegister + 11 * 16)], 1
+    vinsertf128 ymm12, ymm12, xmmword ptr [rdi + (CONTEXT_VectorRegister + 12 * 16)], 1
+    vinsertf128 ymm13, ymm13, xmmword ptr [rdi + (CONTEXT_VectorRegister + 13 * 16)], 1
+    vinsertf128 ymm14, ymm14, xmmword ptr [rdi + (CONTEXT_VectorRegister + 14 * 16)], 1
+    vinsertf128 ymm15, ymm15, xmmword ptr [rdi + (CONTEXT_VectorRegister + 15 * 16)], 1
+LOCAL_LABEL(Done_Restore_CONTEXT_XSTATE):
+
     test    BYTE PTR [rdi + CONTEXT_ContextFlags], CONTEXT_CONTROL
     je      LOCAL_LABEL(Done_Restore_CONTEXT_CONTROL)
 
diff --git a/src/pal/src/thread/context.cpp b/src/pal/src/thread/context.cpp
index c3412bad7c..49d39d2a3d 100644
--- a/src/pal/src/thread/context.cpp
+++ b/src/pal/src/thread/context.cpp
@@ -563,10 +563,20 @@ void CONTEXTFromNativeContext(const native_context_t *native, LPCONTEXT lpContex
 
     // TODO: Enable for all Unix systems
 #if defined(_AMD64_) && defined(__linux__)
-    if ((contextFlags & CONTEXT_XSTATE) != 0 && FPREG_HasExtendedState(native))
+    if ((contextFlags & CONTEXT_XSTATE) != 0)
     {
-        memcpy_s(lpContext->VectorRegister, sizeof(M128A) * 16, FPREG_Xstate_Ymmh(native), sizeof(M128A) * 16);
-    }
+        if (FPREG_HasExtendedState(native))
+        {
+            memcpy_s(lpContext->VectorRegister, sizeof(M128A) * 16, FPREG_Xstate_Ymmh(native), sizeof(M128A) * 16);
+        }
+        else
+        {
+            // 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_
 }