From e87d936583c5b376e8bc74ceeca09adb6ad8f186 Mon Sep 17 00:00:00 2001
From: Johan Mabille <johan.mabille@gmail.com>
Date: Fri, 27 Aug 2021 18:39:23 +0200
Subject: [PATCH] Fixed simd_traits implementation

---
 include/xsimd/types/xsimd_register.hpp | 17 ++++++-
 include/xsimd/types/xsimd_traits.hpp   | 64 +++++++++++++++++---------
 test/test_traits.cpp                   |  5 ++
 3 files changed, 63 insertions(+), 23 deletions(-)

diff --git a/include/xsimd/types/xsimd_register.hpp b/include/xsimd/types/xsimd_register.hpp
index 78278e70d..d28ccefc4 100644
--- a/include/xsimd/types/xsimd_register.hpp
+++ b/include/xsimd/types/xsimd_register.hpp
@@ -1,6 +1,8 @@
 #ifndef XSIMD_REGISTER_HPP
 #define XSIMD_REGISTER_HPP
 
+#include <type_traits>
+
 namespace xsimd
 {
 
@@ -16,6 +18,11 @@ namespace xsimd
         template<class T, class Arch>
         struct simd_register;
 
+        template <class T, class A>
+        struct has_simd_register : std::false_type
+        {
+        };
+
 #define XSIMD_DECLARE_SIMD_REGISTER(SCALAR_TYPE, ISA, VECTOR_TYPE) \
     template<> \
     struct simd_register<SCALAR_TYPE, ISA>\
@@ -23,7 +30,10 @@ namespace xsimd
         using register_type = VECTOR_TYPE;\
         register_type data;\
         operator register_type() const { return data; }\
-    }
+    };\
+    template <>\
+    struct has_simd_register<SCALAR_TYPE, ISA> : std::true_type\
+    {}
 
 #define XSIMD_DECLARE_SIMD_REGISTER_ALIAS(ISA, ISA_BASE)\
     template<class T> \
@@ -32,7 +42,10 @@ namespace xsimd
         using register_type = typename simd_register<T, ISA_BASE>::register_type;\
         simd_register(register_type reg) : simd_register<T, ISA_BASE>{reg} {}\
         simd_register() = default;\
-    }
+    };\
+    template<class T>\
+    struct has_simd_register<T, ISA> : has_simd_register<T, ISA_BASE>\
+    {}
 
         template <class T, class Arch>
         struct get_bool_simd_register
diff --git a/include/xsimd/types/xsimd_traits.hpp b/include/xsimd/types/xsimd_traits.hpp
index 1ce0ef3ff..50f965dc9 100644
--- a/include/xsimd/types/xsimd_traits.hpp
+++ b/include/xsimd/types/xsimd_traits.hpp
@@ -17,50 +17,72 @@
 
 namespace xsimd
 {
+
+    /**************************************
+     * simd_traits and revert_simd_traits *
+     **************************************/
+
+    template <class T, class A = default_arch>
+    struct has_simd_register : types::has_simd_register<T, A>
+    {
+    };
+
     namespace detail
     {
-        template <class T, class = void>
-        struct has_batch : std::false_type
+        template <class T, bool>
+        struct simd_traits_impl;
+
+        template <class T>
+        struct simd_traits_impl<T, false>
         {
+            using type = T;
+            using bool_type = bool;
+            static constexpr size_t size = 1;
         };
 
         template <class T>
-        struct has_batch<T, check_size_t<sizeof(batch<T>)>> : std::true_type
+        constexpr size_t simd_traits_impl<T, false>::size;
+
+        template <class T>
+        struct simd_traits_impl<T, true>
         {
+            using type = batch<T>;
+            using bool_type = typename type::batch_bool_type;
+            static constexpr size_t size = type::size;
         };
+
+        template <class T>
+        constexpr size_t simd_traits_impl<T, true>::size;
     }
 
-    template <class T, bool = detail::has_batch<T>::value>
-    struct simd_traits
+    template <class T>
+    struct simd_traits : detail::simd_traits_impl<T, has_simd_register<T>::value>
     {
-        using type = T;
-        using bool_type = bool;
-        static constexpr size_t size = 1;
     };
 
-    template <class T, bool B>
-    constexpr size_t simd_traits<T, B>::size;
-
     template <class T>
-    struct revert_simd_traits
+    struct simd_traits<std::complex<T>>
+        : detail::simd_traits_impl<std::complex<T>, has_simd_register<T>::value>
     {
-        using type = T;
-        static constexpr size_t size = simd_traits<type>::size;
     };
 
-    template <class T>
-    constexpr size_t revert_simd_traits<T>::size;
+#ifdef XSIMD_ENABLE_XTL_COMPLEX
+    template <class T, bool i3ec>
+    struct simd_traits<xtl::xcomplex<T, T, i3ec>>
+        : detail::simd_traits_impl<std::complex<T>, has_simd_register<T>::value>
+    {
+    };
+#endif
 
     template <class T>
-    struct simd_traits<T, true>
+    struct revert_simd_traits
     {
-        using type = batch<T>;
-        using bool_type = typename type::batch_bool_type;
-        static constexpr size_t size = type::size;
+        using type = T;
+        static constexpr size_t size = simd_traits<type>::size;
     };
 
     template <class T>
-    constexpr size_t simd_traits<T, true>::size;
+    constexpr size_t revert_simd_traits<T>::size;
 
     template <class T>
     struct revert_simd_traits<batch<T>>
diff --git a/test/test_traits.cpp b/test/test_traits.cpp
index 134802952..ea31bff04 100644
--- a/test/test_traits.cpp
+++ b/test/test_traits.cpp
@@ -27,6 +27,11 @@ class traits_test : public testing::Test
         using batch_bool_type = xsimd::batch_bool<value_type>;
         constexpr bool same_bool_type = std::is_same<batch_bool_type, typename traits_type::bool_type>::value;
         EXPECT_TRUE(same_bool_type);
+
+        using vector_traits_type = xsimd::simd_traits<std::vector<value_type>>;
+        EXPECT_EQ(vector_traits_type::size, 1);
+        constexpr bool vec_same_type = std::is_same<typename vector_traits_type::type, std::vector<value_type>>::value;
+        EXPECT_TRUE(vec_same_type);
     }
 
     void test_revert_simd_traits()