From eb59787849400f93167ee7a1945591232e487e78 Mon Sep 17 00:00:00 2001
From: Guillaume Chatelet <gchatelet@google.com>
Date: Fri, 17 Jan 2020 16:53:37 +0100
Subject: [PATCH] Support disabling of extensions Fixes #101

---
 README.md                     | 23 +++++++++++++++++++++++
 include/cpu_features_macros.h | 18 ++++++++++++++++++
 include/cpuinfo_x86.h         |  2 ++
 src/cpuinfo_x86.c             |  7 +++++++
 test/cpuinfo_x86_test.cc      | 20 ++++++++++++++++++++
 5 files changed, 70 insertions(+)

diff --git a/README.md b/README.md
index 29d7946..5b091f4 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,7 @@ instructions) at runtime.
 - [Design Rationale](#rationale)
 - [Code samples](#codesample)
 - [Running sample code](#usagesample)
+- [Mocking the library](#mock)
 - [What's supported](#support)
 - [Android NDK's drop in replacement](#ndk)
 - [License](#license)
@@ -130,6 +131,28 @@ flags           : aes,avx,cx16,smx,sse4_1,sse4_2,ssse3
 {"arch":"x86","brand":"       Intel(R) Xeon(R) CPU E5-1650 0 @ 3.20GHz","family":6,"model":45,"stepping":7,"uarch":"INTEL_SNB","flags":["aes","avx","cx16","smx","sse4_1","sse4_2","ssse3"]}
 ```
 
+<a name="mock"></a>
+### Mocking the library
+
+When testing code depending on `cpu_features` it may be interesting to disable
+support for a particular extension (e.g. test `sse4` support on a `Skylake`
+machine).
+
+It is easily done by setting up an interceptor callback.
+
+```C++
+TEST(X86Test, TestOnlySSE4) {
+  RegisterX86InfoInterceptor([](X86Info* info) {
+    info->features = X86Features{};
+    info->features.sse4 = true;
+  });
+  // 1. Call code that uses GetX86Info(),
+  // 2. Check the expected values,
+  // 3. Don't forget to cancel.
+  RegisterX86InfoInterceptor(NULL);
+}
+```
+
 <a name="support"></a>
 ## What's supported
 
diff --git a/include/cpu_features_macros.h b/include/cpu_features_macros.h
index 2227160..52cb712 100644
--- a/include/cpu_features_macros.h
+++ b/include/cpu_features_macros.h
@@ -140,4 +140,22 @@
 #define CPU_FEATURES_COMPILED_MIPS_MSA defined(__mips_msa)
 #endif
 
+////////////////////////////////////////////////////////////////////////////////
+// Miscellaneous
+////////////////////////////////////////////////////////////////////////////////
+
+#ifndef thread_local
+#if __STDC_VERSION__ >= 201112 && !defined __STDC_NO_THREADS__
+#define thread_local _Thread_local
+#elif defined _WIN32 && (defined _MSC_VER || defined __ICL || \
+                         defined __DMC__ || defined __BORLANDC__)
+#define thread_local __declspec(thread)
+/* note that ICC (linux) and Clang are covered by __GNUC__ */
+#elif defined __GNUC__ || defined __SUNPRO_C || defined __xlC__
+#define thread_local __thread
+#else
+#error "Cannot define thread_local"
+#endif
+#endif
+
 #endif  // CPU_FEATURES_INCLUDE_CPU_FEATURES_MACROS_H_
diff --git a/include/cpuinfo_x86.h b/include/cpuinfo_x86.h
index 24a4f0c..94db414 100644
--- a/include/cpuinfo_x86.h
+++ b/include/cpuinfo_x86.h
@@ -201,6 +201,8 @@ const char* GetX86FeaturesEnumName(X86FeaturesEnum);
 
 const char* GetX86MicroarchitectureName(X86Microarchitecture);
 
+void RegisterX86InfoInterceptor(void (*)(X86Info*));
+
 CPU_FEATURES_END_CPP_NAMESPACE
 
 #if !defined(CPU_FEATURES_ARCH_X86)
diff --git a/src/cpuinfo_x86.c b/src/cpuinfo_x86.c
index ab14053..dca285e 100644
--- a/src/cpuinfo_x86.c
+++ b/src/cpuinfo_x86.c
@@ -23,6 +23,12 @@
 #error "Cannot compile cpuinfo_x86 on a non x86 platform."
 #endif
 
+thread_local void (*X86InfoInterceptor)(X86Info*) = NULL;
+
+void RegisterX86InfoInterceptor(void (*ptr)(X86Info*)) {
+  X86InfoInterceptor = ptr;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Definitions for CpuId and GetXCR0Eax.
 ////////////////////////////////////////////////////////////////////////////////
@@ -615,6 +621,7 @@ X86Info GetX86Info(void) {
   if (IsVendor(leaf_0, "GenuineIntel") || IsVendor(leaf_0, "AuthenticAMD")) {
     ParseCpuId(max_cpuid_leaf, &info);
   }
+  if (X86InfoInterceptor) X86InfoInterceptor(&info);
   return info;
 }
 
diff --git a/test/cpuinfo_x86_test.cc b/test/cpuinfo_x86_test.cc
index 10b9624..8737c81 100644
--- a/test/cpuinfo_x86_test.cc
+++ b/test/cpuinfo_x86_test.cc
@@ -101,6 +101,26 @@ TEST(CpuidX86Test, SandyBridge) {
   EXPECT_FALSE(features.rdrnd);
 }
 
+TEST(CpuidX86Test, ForgingCpuWithInterceptor) {
+  RegisterX86InfoInterceptor([](X86Info* info) {
+    memcpy(info->vendor, "TEST", sizeof("TEST"));
+    info->features = X86Features{};
+    info->features.erms = true;
+  });
+  const auto info = GetX86Info();
+  EXPECT_STREQ(info.vendor, "GenuineIntel");
+
+  const auto& features = info.features;
+  for (size_t i = 0; i < X86_LAST_; ++i)
+    if (i == X86_ERMS)
+      EXPECT_TRUE(GetX86FeaturesEnumValue(&features, (X86FeaturesEnum)i));
+    else
+      EXPECT_FALSE(GetX86FeaturesEnumValue(&features, (X86FeaturesEnum)i));
+
+  RegisterX86InfoInterceptor(NULL);
+  EXPECT_STREQ(GetX86Info().vendor, "GenuineIntel");
+}
+
 const int KiB = 1024;
 const int MiB = 1024 * KiB;