From d572dd2bfc4c5dab14ced9a8c0cd4c157f6bce10 Mon Sep 17 00:00:00 2001
From: "Mr. eXoDia" <mr.exodia.tpodt@gmail.com>
Date: Mon, 8 Dec 2014 16:33:04 +0100
Subject: [PATCH] Fixed AVX YMM registers

---
 SDK/C/TitanEngine.h                          | 10 ++-
 TitanEngine/TitanEngine.Debugger.Context.cpp | 64 ++++++--------------
 TitanEngine/stdafx.h                         | 10 ++-
 3 files changed, 36 insertions(+), 48 deletions(-)

diff --git a/SDK/C/TitanEngine.h b/SDK/C/TitanEngine.h
index 8023c74..0576db7 100644
--- a/SDK/C/TitanEngine.h
+++ b/SDK/C/TitanEngine.h
@@ -594,6 +594,12 @@ typedef struct
     DWORD OriginalCOMTableSize;
 } FILE_FIX_INFO, *PFILE_FIX_INFO;
 
+typedef struct
+{
+    M128A Low; //XMM/SSE part
+    M128A High; //AVX part
+} YmmRegister_t;
+
 typedef struct
 {
     BYTE    data[10];
@@ -652,10 +658,10 @@ typedef struct
     DWORD MxCsr;
 #ifdef _WIN64
     M128A XmmRegisters[16];
-    BYTE YmmRegisters[32 * 16];
+    YmmRegister_t YmmRegisters[16];
 #else // x86
     M128A XmmRegisters[8];
-    BYTE YmmRegisters[32 * 8];
+    YmmRegister_t YmmRegisters[8];
 #endif
 } TITAN_ENGINE_CONTEXT_t;
 
diff --git a/TitanEngine/TitanEngine.Debugger.Context.cpp b/TitanEngine/TitanEngine.Debugger.Context.cpp
index db9e61c..ccde362 100644
--- a/TitanEngine/TitanEngine.Debugger.Context.cpp
+++ b/TitanEngine/TitanEngine.Debugger.Context.cpp
@@ -962,32 +962,20 @@ __declspec(dllexport) bool TITCALL SetAVXContext(HANDLE hActiveThread, TITAN_ENG
         return false;
 
     DWORD FeatureLength;
-    PM128A Xmm = (PM128A)_LocateXStateFeature(Context, XSTATE_LEGACY_SSE, &FeatureLength);
-    if(Xmm != NULL) //If the feature is unsupported by the processor it will return NULL
+    PM128A Sse = (PM128A)_LocateXStateFeature(Context, XSTATE_LEGACY_SSE, &FeatureLength);
+    PM128A Avx = (PM128A)_LocateXStateFeature(Context, XSTATE_AVX, NULL);
+    int NumberOfRegisters = FeatureLength / sizeof(Sse[0]);
+
+    if(Sse != NULL) //If the feature is unsupported by the processor it will return NULL
     {
-        for(DWORD Index = 0; Index < FeatureLength / sizeof(* Xmm); Index += 1)
-        {
-            memcpy
-            (
-                Xmm++,
-                ((char*) & (titcontext->YmmRegisters[32 * Index])),
-                sizeof(Xmm[Index])
-            );
-        }
+        for(int i = 0; i < NumberOfRegisters; i++)
+            Sse[i] = titcontext->YmmRegisters[i].Low;
     }
 
-    PM128A Ymm = (PM128A)_LocateXStateFeature(Context, XSTATE_AVX, NULL);
-    if(Ymm != NULL) //If the feature is unsupported by the processor it will return NULL
+    if(Avx != NULL) //If the feature is unsupported by the processor it will return NULL
     {
-        for(DWORD Index = 0; Index < FeatureLength / sizeof(* Ymm); Index += 1)
-        {
-            memcpy
-            (
-                Ymm++,
-                ((char*) & (titcontext->YmmRegisters[32 * Index])) + sizeof(titcontext->XmmRegisters[Index]),
-                sizeof(Ymm[Index])
-            );
-        }
+        for(int i = 0; i < NumberOfRegisters; i++)
+            Avx[i] = titcontext->YmmRegisters[i].High;
     }
 
     return (SetThreadContext(hActiveThread, Context) == TRUE);
@@ -1034,32 +1022,20 @@ __declspec(dllexport) bool TITCALL GetAVXContext(HANDLE hActiveThread, TITAN_ENG
         return false;
 
     DWORD FeatureLength;
-    PM128A Xmm = (PM128A)_LocateXStateFeature(Context, XSTATE_LEGACY_SSE, &FeatureLength);
-    if(Xmm != NULL) //If the feature is unsupported by the processor it will return NULL
+    PM128A Sse = (PM128A)_LocateXStateFeature(Context, XSTATE_LEGACY_SSE, &FeatureLength);
+    PM128A Avx = (PM128A)_LocateXStateFeature(Context, XSTATE_AVX, NULL);
+    int NumberOfRegisters = FeatureLength / sizeof(Sse[0]);
+
+    if(Sse != NULL) //If the feature is unsupported by the processor it will return NULL
     {
-        for(DWORD Index = 0; Index < FeatureLength / sizeof(*Xmm); Index++)
-        {
-            memcpy
-            (
-                (char*) & (titcontext->YmmRegisters[32 * Index]),
-                Xmm++,
-                sizeof(Xmm[Index])
-            );
-        }
+        for(int i = 0; i < NumberOfRegisters; i++)
+            titcontext->YmmRegisters[i].Low = Sse[i];
     }
 
-    PM128A Ymm = (PM128A)_LocateXStateFeature(Context, XSTATE_AVX, &FeatureLength);
-    if(Ymm != NULL) //If the feature is unsupported by the processor it will return NULL
+    if(Avx != NULL) //If the feature is unsupported by the processor it will return NULL
     {
-        for(DWORD Index = 0; Index < FeatureLength / sizeof(*Ymm); Index++)
-        {
-            memcpy
-            (
-                ((char*) & (titcontext->YmmRegisters[32 * Index])) + sizeof(titcontext->XmmRegisters[Index]),
-                Ymm++,
-                sizeof(Ymm[Index])
-            );
-        }
+        for(int i = 0; i < NumberOfRegisters; i++)
+            titcontext->YmmRegisters[i].High = Avx[i];
     }
 
     return true;
diff --git a/TitanEngine/stdafx.h b/TitanEngine/stdafx.h
index 6423b24..4649439 100644
--- a/TitanEngine/stdafx.h
+++ b/TitanEngine/stdafx.h
@@ -93,6 +93,12 @@
 #define CONTEXT_EXTENDED_REGISTERS 0
 #endif
 
+typedef struct
+{
+    M128A Low; //XMM/SSE part
+    M128A High; //AVX part
+} YmmRegister_t;
+
 typedef struct
 {
     BYTE    data[10];
@@ -151,10 +157,10 @@ typedef struct
     DWORD MxCsr;
 #ifdef _WIN64
     M128A XmmRegisters[16];
-    BYTE YmmRegisters[32 * 16];
+    YmmRegister_t YmmRegisters[16];
 #else // x86
     M128A XmmRegisters[8];
-    BYTE YmmRegisters[32 * 8];
+    YmmRegister_t YmmRegisters[8];
 #endif
 } TITAN_ENGINE_CONTEXT_t;