Particle Data Structure

Creating a general AoSoA data structure with general types.

packages/Core/src/Cabana_AoSoA.hpp

 #ifndef CABANA_AOSOA_HPP
 #define CABANA_AOSOA_HPP
 
-#include <Cabana_SoA.hpp>
+#include "Cabana_SoA.hpp"
 
 #include <type_traits>
 #include <cmath>
+#include <cstdlib>
 
 namespace Cabana
 {
 //---------------------------------------------------------------------------//
-template<typename Device, size_t ArraySize, typename... Types>
+template<typename Device, std::size_t ArraySize, typename... Types>
 class AoSoA
 {
   public:
@@ -18,27 +19,31 @@ class AoSoA
     using device_type = Device;
 
     // Inner array size (size of the arrays held by the structs).
-    static constexpr size_t array_size = ArraySize;
+    static constexpr std::size_t array_size = ArraySize;
 
     // SoA type.
     using struct_type = SoA<array_size,Types...>;
 
     // Struct member array return type at a given index I.
-    template<size_t I>
+    template<std::size_t I>
     using struct_member_array_type =
         typename ArrayTypeAtIndex<I,array_size,Types...>::return_type;
 
+    // Struct member array data type at a given index I.
+    template<std::size_t I>
+    using struct_member_data_type =
+        typename std::remove_pointer<struct_member_array_type<I> >::type;
+
     // Struct member array element value type at a given index I.
-    template<size_t I>
+    template<std::size_t I>
     using struct_member_value_type =
-        typename std::remove_all_extents<struct_member_array_type<I>>::type;
+        typename std::remove_all_extents<struct_member_data_type<I> >::type;
 
     // Struct member array element reference type at a given index I.
-    template<size_t I>
+    template<std::size_t I>
     using struct_member_reference_type =
         typename std::add_lvalue_reference<struct_member_value_type<I>>::type;
 
-
   public:
 
     //! AoSoA index.
@@ -47,21 +52,29 @@ class AoSoA
       private:
 
         // The struct index.
-        size_t _s;
+        std::size_t _s;
 
         // Array offset index in the struct.
-        size_t _i;
+        std::size_t _i;
 
       public:
 
+        // Constructor.
+        Index( const std::size_t struct_id, const std::size_t offset )
+            : _s( struct_id )
+            , _i( offset )
+        {};
+
         // Get the struct index.
-        inline size_t structIndex() const { return _s; }
+        inline std::size_t s() const
+        { return _s; }
 
-        // Get the array offset index in the struct.
-        inline size_t arrayIndex() const { return _i; }
+        // Get the struct offset.
+        inline std::size_t i() const
+        { return _i; }
 
         // Prefix increment operator.
-        ParticleIndex& operator++()
+        Index& operator++()
         {
             _i = ( array_size - 1 == _i ) ? 0 : _i + 1;
             _s = ( 0 == _i ) ? _s + 1: _s;
@@ -69,21 +82,21 @@ class AoSoA
         };
 
         // Postfix increment operator.
-        ParticleIndex operator++(int)
+        Index operator++(int)
         {
-            ParticleIndex temp = *this;
+            Index temp = *this;
             ++*this;
             return temp;
         };
 
         // Equality comparator.
-        bool operator==( const ParticleIndex& rhs ) const
+        bool operator==( const Index& rhs ) const
         {
-            return (_s == rhs._s) && (_i == rhs._i)
+            return (_s == rhs._s) && (_i == rhs._i);
         }
 
         // Inequality comparator.
-        bool operator!=( const ParticleIndex& rhs ) const
+        bool operator!=( const Index& rhs ) const
         {
             return (_s != rhs._s) || (_i != rhs._i);
         }
@@ -92,81 +105,112 @@ class AoSoA
   public:
 
     // Constructor.
-    AoSoA( const size_t size );
+    AoSoA( const std::size_t size )
+        : _size( size )
+    {
+        _num_struct = _size / ArraySize;
+        if ( 0 < _size % ArraySize ) ++_num_struct;
+
+        Device::allocate( _structs, _num_struct );
+    }
 
     // Destructor.
-    ~AoSoA();
+    ~AoSoA()
+    {
+        Device::deallocate( _structs );
+    }
 
     // Get the requested size of the AoSoA (global number of requested
     // elements). This may or not be equal to the total number of allocated
     // elements in the AoSoA if the size of the struct arrays is not evenly
     // divisible by the requested size.
-    size_t size() const { return _size; }
+    std::size_t size() const { return _size; }
 
     // Get the allocated size of the AoSoA (global number of allocated
     // elements). This may or not be equal to the total number of requested
     // elements in the AoSoA if the size of the struct arrays is not evenly
     // divisible by the requested size.
-    size_t allocatedSize() const { return _num_struct * array_size; }
+    std::size_t allocatedSize() const { return _num_struct * array_size; }
 
     // Get the number of structs in the array.
-    size_t numStruct() const { return _num_struct; }
+    std::size_t numStruct() const { return _num_struct; }
 
     // Get the beginning of the array.
-    Index begin() const;
+    Index begin() const
+    {
+        return Index( 0, 0 );
+    }
 
     // Get the end of the array.
-    Index end() const;
+    Index end() const
+    {
+        return Index( _num_struct - 1, _size % ArraySize );
+    }
 
     // Access the data array at a given struct member index.
-    template<size_t I>
-    inline struct_member_array_type<I> array( const ParticleIndex& idx );
+    template<std::size_t I>
+    inline struct_member_array_type<I> array( const Index& idx )
+    {
+        return getStructMember<I>( _structs[idx.s()] );
+    }
 
     // Access the data value at a given struct member index and offset index
     // for Rank 0 data.
-    template<size_t I>
+    template<std::size_t I>
     inline
-    std::enable_if<(1==std::rank<struct_member_array_type<I>>),
-                   struct_member_reference_type<I> >::type
-    operator()( const ParticleIndex& idx );
+    typename std::enable_if<(0==std::rank<struct_member_data_type<I> >::value),
+                            struct_member_reference_type<I> >::type
+    get( const Index& idx )
+    {
+        return array<I>(idx)[idx.i()];
+    }
 
     // Access the data value at a given struct member index and offset index
     // for Rank 1 data.
-    template<size_t I>
+    template<std::size_t I>
     inline
-    std::enable_if<(2==std::rank<struct_member_array_type<I>>),
-                   struct_member_reference_type<I> >::type
-    operator()( const ParticleIndex& idx,
-                const int d0 );
+    typename std::enable_if<(1==std::rank<struct_member_data_type<I> >::value),
+                            struct_member_reference_type<I> >::type
+    get( const Index& idx,
+         const int d0 )
+    {
+        return array<I>(idx)[idx.i()][d0];
+    }
 
     // Access the data value at a given struct member index and offset index
     // for Rank 2 data.
-    template<size_t I>
+    template<std::size_t I>
     inline
-    std::enable_if<(3==std::rank<struct_member_array_type<I>>),
-                   struct_member_reference_type<I> >::type
-    operator()( const ParticleIndex& idx,
-                const int d0,
-                const int d1 );
+    typename std::enable_if<(2==std::rank<struct_member_data_type<I> >::value),
+                            struct_member_reference_type<I> >::type
+    get( const Index& idx,
+         const int d0,
+         const int d1 )
+    {
+        return array<I>(idx)[idx.i()][d0][d1];
+    }
 
     // Access the data value at a given struct member index and offset index
     // for Rank 3 data.
-    template<size_t I>
+    template<std::size_t I>
     inline
-    std::enable_if<(4==std::rank<struct_member_array_type<I>>),
-                   struct_member_reference_type<I> >::type
-    operator()( const ParticleIndex& idx,
-                const int d0,
-                const int d1,
-                const int d2 );
+    typename std::enable_if<(3==std::rank<struct_member_data_type<I> >::value),
+                            struct_member_reference_type<I> >::type
+    get( const Index& idx,
+         const int d0,
+         const int d1,
+         const int d2 )
+    {
+        return array<I>(idx)[idx.i()][d0][d1][d2];
+    }
 
   private:
 
     // Total size of the array.
-    size_t _size;
+    std::size_t _size;
 
     // Number of structs in the array.
-    size_t _num_struct;
+    std::size_t _num_struct;
 
     // Structs-of-Arrays.
     struct_type* _structs;
@@ -176,116 +220,4 @@ class AoSoA
 
 } // end namespace Cabana
 
-//---------------------------------------------------------------------------//
-// Implementation.
-//---------------------------------------------------------------------------//
-
-namespace Cabana
-{
-//---------------------------------------------------------------------------//
-// Constructor.
-template<typename Device, size_t ArraySize, typename... Types>
-AoSoA<Device,ArraySize,Types...>::AoSoA( const size_t size )
-    : _size( size )
-    , _num_struct( std::ceil(size/ArraySize) )
-{
-    Device::allocate( _structs, _num_struct );
-}
-
-//---------------------------------------------------------------------------//
-// Get the beginning of the array.
-template<typename Device, size_t ArraySize, typename... Types>
-Index AoSoA<Device,ArraySize,Types...>::begin() const
-{
-    Index idx;
-    idx._s = 0;
-    idx._i = 0;
-    return idx;
-}
-
-//---------------------------------------------------------------------------//
-// Get the end of the array.
-template<typename Device, size_t ArraySize, typename... Types>
-Index AoSoA<Device,ArraySize,Types...>::end() const
-{
-    Index idx;
-    idx._s = _num_struct - 1;
-    idx._i = _size % ArraySize;
-    return idx;
-}
-
-//---------------------------------------------------------------------------//
-// Access the data array at a given struct member index.
-template<size_t I>
-template<typename Device, size_t ArraySize, typename... Types>
-inline struct_member_array_type<I>
-AoSoA<Device,ArraySize,Types...>::array( const ParticleIndex& idx )
-{
-    return getStructMember<I>( _structs[idx._s] );
-}
-
-//---------------------------------------------------------------------------//
-// Access the data value at a given struct member index and offset index
-// for Rank 0 data.
-template<size_t I>
-template<typename Device, size_t ArraySize, typename... Types>
-inline
-std::enable_if<(1==std::rank<struct_member_array_type<I>>),
-               struct_member_reference_type<I> >::type
-AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx )
-{
-    return array(idx)[idx._i];
-}
-
-//---------------------------------------------------------------------------//
-// Access the data value at a given struct member index and offset index
-// for Rank 1 data.
-template<size_t I>
-template<typename Device, size_t ArraySize, typename... Types>
-inline
-std::enable_if<(2==std::rank<struct_member_array_type<I>>),
-               struct_member_reference_type<I> >::type
-AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx,
-                                              const int d0 )
-{
-    return array(idx)[idx._i][d0];
-}
-
-//---------------------------------------------------------------------------//
-// Access the data value at a given struct member index and offset index
-// for Rank 2 data.
-template<size_t I>
-template<typename Device, size_t ArraySize, typename... Types>
-inline
-std::enable_if<(3==std::rank<struct_member_array_type<I>>),
-               struct_member_reference_type<I> >::type
-AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx,
-                                              const int d0,
-                                              const int d1 )
-{
-    return array(idx)[idx._i][d0][d1];
-}
-
-//---------------------------------------------------------------------------//
-// Access the data value at a given struct member index and offset index
-// for Rank 3 data.
-template<size_t I>
-template<typename Device, size_t ArraySize, typename... Types>
-inline
-std::enable_if<(4==std::rank<struct_member_array_type<I>>),
-               struct_member_reference_type<I> >::type
-AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx,
-                                              const int d0,
-                                              const int d1,
-                                              const int d2 )
-{
-    return array(idx)[idx._i][d0][d1][d2];
-}
-
-//---------------------------------------------------------------------------//
-
-} // end namespace Cabana
-
-//---------------------------------------------------------------------------//
-
 #endif // CABANA_AOSOA_HPP