From c4975be2f540c2057e4ed83b30bc267a66c73a62 Mon Sep 17 00:00:00 2001
From: Michael Schellenberger Costa <miscco@nvidia.com>
Date: Sun, 4 Dec 2022 12:42:39 +0100
Subject: [PATCH] Enable construction of vectors from `std::initializer_list`

---
 testing/vector.cu                   | 29 ++++++++++++++++++++++++++++
 thrust/detail/vector_base.h         | 20 +++++++++++++++++++
 thrust/detail/vector_base.inl       | 30 +++++++++++++++++++++++++++++
 thrust/device_vector.h              | 24 +++++++++++++++++++++++
 thrust/host_vector.h                | 25 ++++++++++++++++++++++++
 thrust/system/cpp/detail/vector.inl | 21 ++++++++++++++++++++
 6 files changed, 149 insertions(+)

diff --git a/testing/vector.cu b/testing/vector.cu
index de211af93..4f8d80ed3 100644
--- a/testing/vector.cu
+++ b/testing/vector.cu
@@ -4,6 +4,9 @@
 #include <thrust/sequence.h>
 #include <thrust/device_malloc_allocator.h>
 
+#if THRUST_CPP_DIALECT >= 2011
+#include <initializer_list>
+#endif
 #include <vector>
 #include <list>
 #include <limits>
@@ -37,6 +40,32 @@ void TestVectorBool(void)
 }
 DECLARE_UNITTEST(TestVectorBool);
 
+#if THRUST_CPP_DIALECT >= 2011
+template <class Vector>
+void TestVectorInitializerList(void)
+{
+    Vector v{1, 2, 3};
+    ASSERT_EQUAL(v.size(), 3lu);
+    ASSERT_EQUAL(v[0], 0);
+    ASSERT_EQUAL(v[1], 1);
+    ASSERT_EQUAL(v[2], 2);
+
+    v = {1, 2, 3, 4};
+    ASSERT_EQUAL(v.size(), 4lu);
+    ASSERT_EQUAL(v[0], 1);
+    ASSERT_EQUAL(v[1], 2);
+    ASSERT_EQUAL(v[2], 3);
+    ASSERT_EQUAL(v[3], 4);
+    
+    const auto alloc = v.get_allocator();
+    Vector v2{{1, 2, 3}, alloc};
+    ASSERT_EQUAL(v2.size(), 3lu);
+    ASSERT_EQUAL(v2[0], 0);
+    ASSERT_EQUAL(v2[1], 1);
+    ASSERT_EQUAL(v2[2], 2);
+}
+DECLARE_VECTOR_UNITTEST(TestVectorInitializerList);
+#endif
 
 template <class Vector>
 void TestVectorFrontBack(void)
diff --git a/thrust/detail/vector_base.h b/thrust/detail/vector_base.h
index b05f35194..025082e64 100644
--- a/thrust/detail/vector_base.h
+++ b/thrust/detail/vector_base.h
@@ -28,6 +28,10 @@
 #include <thrust/detail/type_traits.h>
 #include <thrust/detail/config.h>
 #include <thrust/detail/contiguous_storage.h>
+
+#if THRUST_CPP_DIALECT >= 2011
+#include <initializer_list>
+#endif
 #include <vector>
 
 THRUST_NAMESPACE_BEGIN
@@ -128,6 +132,22 @@ template<typename T, typename Alloc>
      *  \param v The vector_base to move.
      */
     vector_base &operator=(vector_base &&v);
+
+    /*! This constructor builds a \p vector_base from an intializer_list.
+     *  \param il The intializer_list.
+     */
+    vector_base(std::initializer_list<T> il);
+      
+    /*! This constructor builds a \p vector_base from an intializer_list.
+     *  \param il The intializer_list.
+     *  \param alloc The allocator to use by this device_vector.
+     */
+    vector_base(std::initializer_list<T> il, const Alloc &alloc);
+    
+    /*! Assign operator copies from an initializer_list
+     *  \param il The initializer_list.
+     */
+    vector_base &operator=(std::initializer_list<T> il);
   #endif
 
     /*! Copy constructor copies from an exemplar vector_base with different
diff --git a/thrust/detail/vector_base.inl b/thrust/detail/vector_base.inl
index ab94429a8..cdb4bbaab 100644
--- a/thrust/detail/vector_base.inl
+++ b/thrust/detail/vector_base.inl
@@ -195,6 +195,36 @@ template<typename T, typename Alloc>
   return *this;
 } // end vector_base::operator=()
 
+#if THRUST_CPP_DIALECT >= 2011
+  template<typename T, typename Alloc>
+    vector_base<T,Alloc>
+      ::vector_base(std::initializer_list<T> il)
+        :m_storage(),
+         m_size(0)
+  {
+    range_init(il.begin(), il.end());
+  } // end vector_base::vector_base()
+
+  template<typename T, typename Alloc>
+  vector_base<T,Alloc>
+    ::vector_base(std::initializer_list<T> il, const Alloc &alloc)
+    :m_storage(alloc),
+      m_size(0)
+  {
+    range_init(il.begin(), il.end());
+  } // end vector_base::vector_base()
+
+  template<typename T, typename Alloc>
+    vector_base<T,Alloc> &
+      vector_base<T,Alloc>
+      ::operator=(std::initializer_list<T> il)
+  {
+    assign(il.begin(), il.end());
+
+    return *this;
+  } // end vector_base::operator=()
+#endif
+
 template<typename T, typename Alloc>
   template<typename IteratorOrIntegralType>
     void vector_base<T,Alloc>
diff --git a/thrust/device_vector.h b/thrust/device_vector.h
index b00251a0d..3d67e70ca 100644
--- a/thrust/device_vector.h
+++ b/thrust/device_vector.h
@@ -26,6 +26,9 @@
 #include <thrust/detail/vector_base.h>
 #include <thrust/device_allocator.h>
 
+#if THRUST_CPP_DIALECT >= 2011
+#include <initializer_list>
+#endif
 #include <vector>
 #include <utility>
 
@@ -197,6 +200,27 @@ template<typename T, typename Alloc = thrust::device_allocator<T> >
     device_vector &operator=(const detail::vector_base<OtherT,OtherAlloc> &v)
     { Parent::operator=(v); return *this; }
 
+#if THRUST_CPP_DIALECT >= 2011
+    /*! This constructor builds a \p device_vector from an intializer_list.
+     *  \param il The intializer_list.
+     */
+    device_vector(std::initializer_list<T> il)
+      :Parent(il.begin(), il.end()) {}
+      
+    /*! This constructor builds a \p device_vector from an intializer_list.
+     *  \param il The intializer_list.
+     *  \param alloc The allocator to use by this device_vector.
+     */
+    device_vector(std::initializer_list<T> il, const Alloc &alloc)
+      :Parent(il.begin(), il.end(), alloc) {}
+      
+    /*! Assign an \p intializer_list with a matching element type
+     *  \param il The intializer_list.
+     */
+    device_vector &operator=(std::initializer_list<T> il)
+    { Parent::operator=(il); return *this; }
+#endif
+
     /*! This constructor builds a \p device_vector from a range.
      *  \param first The beginning of the range.
      *  \param last The end of the range.
diff --git a/thrust/host_vector.h b/thrust/host_vector.h
index 01bbceb3b..44baea7f8 100644
--- a/thrust/host_vector.h
+++ b/thrust/host_vector.h
@@ -25,6 +25,10 @@
 #include <thrust/detail/config.h>
 #include <thrust/detail/memory_wrapper.h>
 #include <thrust/detail/vector_base.h>
+
+#if THRUST_CPP_DIALECT >= 2011
+#include <initializer_list>
+#endif
 #include <vector>
 #include <utility>
 
@@ -216,6 +220,27 @@ template<typename T, typename Alloc = std::allocator<T> >
     __host__
     host_vector &operator=(const detail::vector_base<OtherT,OtherAlloc> &v)
     { Parent::operator=(v); return *this; }
+    
+#if THRUST_CPP_DIALECT >= 2011
+    /*! This constructor builds a \p host_vector from an intializer_list.
+     *  \param il The intializer_list.
+     */
+    host_vector(std::initializer_list<T> il)
+      :Parent(il.begin(), il.end()) {}
+      
+    /*! This constructor builds a \p host_vector from an intializer_list.
+     *  \param il The intializer_list.
+     *  \param alloc The allocator to use by this host_vector.
+     */
+    host_vector(std::initializer_list<T> il, const Alloc &alloc)
+      :Parent(il.begin(), il.end(), alloc) {}
+      
+    /*! Assign an \p intializer_list with a matching element type
+     *  \param il The intializer_list.
+     */
+    host_vector &operator=(std::initializer_list<T> il)
+    { Parent::operator=(il); return *this; }
+#endif
 
     /*! This constructor builds a \p host_vector from a range.
      *  \param first The beginning of the range.
diff --git a/thrust/system/cpp/detail/vector.inl b/thrust/system/cpp/detail/vector.inl
index d27cdad64..bb987a213 100644
--- a/thrust/system/cpp/detail/vector.inl
+++ b/thrust/system/cpp/detail/vector.inl
@@ -97,6 +97,27 @@ template<typename T, typename Allocator>
     super_t::operator=(std::move(x));
     return *this;
   }
+  
+  template<typename T, typename Allocator>
+    vector<T,Allocator>
+      ::vector(std::initializer_list<T> il)
+        : super_t(il)
+  {}
+  
+  template<typename T, typename Allocator>
+    vector<T,Allocator>
+      ::vector(std::initializer_list<T> il, const Allocator& alloc)
+        : super_t(il, alloc)
+  {}
+
+  template<typename T, typename Allocator>
+    vector<T,Allocator> &
+      vector<T,Allocator>
+        ::operator=(std::initializer_list<T> il)
+  {
+    super_t::operator=(il);
+    return *this;
+  }
 #endif
 
 template<typename T, typename Allocator>