data_layout_policy.h

#pragma once
#include "gaia/config/config.h"
 
#include <tuple>
#include <type_traits>
#include <utility>
 
#include "gaia/core/span.h"
#include "gaia/core/utility.h"
#include "gaia/mem/mem_alloc.h"
#include "gaia/mem/mem_sani.h"
#include "gaia/meta/reflection.h"
 
namespace gaia {
  namespace mem {
    enum class DataLayout : uint8_t {
      AoS = 0, 
      SoA = 1, 
      SoA8 = 2, 
      SoA16 = 3, 
 
      Count = 4 
    };
 
#define GAIA_LAYOUT(layout_name) static constexpr auto gaia_Data_Layout = ::gaia::mem::DataLayout::layout_name
 
    // Helper templates
    namespace detail {
      template <typename T>
      constexpr uint32_t get_alignment() {
        if constexpr (std::is_empty_v<T>)
          // Always consider 0 for empty types
          return 0U;
        else {
          // Use at least 4 (32-bit systems) or 8 (64-bit systems) bytes for alignment
          return (uint32_t)core::get_min(sizeof(uintptr_t), alignof(T));
        }
      }
 
      constexpr size_t get_aligned_byte_offset(uintptr_t address, size_t alig, size_t itemSize, size_t cnt) {
        const auto padding = mem::padding(address, alig);
        address += padding + itemSize * cnt;
        return address;
      }
 
      template <typename T, size_t Alignment>
      constexpr size_t get_aligned_byte_offset(uintptr_t address, size_t cnt) {
        const auto padding = mem::padding<Alignment>(address);
        address += padding + sizeof(T) * cnt;
        return address;
      }
    } // namespace detail
 
    template <DataLayout TDataLayout, typename TItem>
    struct data_layout_properties;
    template <typename TItem>
    struct data_layout_properties<DataLayout::AoS, TItem> {
      constexpr static DataLayout Layout = DataLayout::AoS;
      constexpr static size_t PackSize = 1;
      constexpr static size_t Alignment = detail::get_alignment<TItem>();
    };
    template <typename TItem>
    struct data_layout_properties<DataLayout::SoA, TItem> {
      constexpr static DataLayout Layout = DataLayout::SoA;
      constexpr static size_t PackSize = 4;
      constexpr static size_t Alignment = PackSize * 4;
    };
    template <typename TItem>
    struct data_layout_properties<DataLayout::SoA8, TItem> {
      constexpr static DataLayout Layout = DataLayout::SoA8;
      constexpr static size_t PackSize = 8;
      constexpr static size_t Alignment = PackSize * 4;
    };
    template <typename TItem>
    struct data_layout_properties<DataLayout::SoA16, TItem> {
      constexpr static DataLayout Layout = DataLayout::SoA16;
      constexpr static size_t PackSize = 16;
      constexpr static size_t Alignment = PackSize * 4;
    };
 
    template <DataLayout TDataLayout, typename TItem>
    struct data_view_policy;
 
    template <DataLayout TDataLayout, typename TItem>
    struct data_view_policy_get;
    template <DataLayout TDataLayout, typename TItem>
    struct data_view_policy_set;
 
    template <DataLayout TDataLayout, typename TItem, size_t Ids>
    struct data_view_policy_get_idx;
    template <DataLayout TDataLayout, typename TItem, size_t Ids>
    struct data_view_policy_set_idx;
 
    template <typename ValueType>
    struct data_view_policy_aos {
      using TargetCastType = std::add_pointer_t<ValueType>;
 
      constexpr static DataLayout Layout = data_layout_properties<DataLayout::AoS, ValueType>::Layout;
      constexpr static size_t Alignment = data_layout_properties<DataLayout::AoS, ValueType>::Alignment;
 
      GAIA_NODISCARD static constexpr uint32_t get_min_byte_size(uintptr_t addr, size_t cnt) noexcept {
        const auto offset = detail::get_aligned_byte_offset<ValueType, Alignment>(addr, cnt);
        return (uint32_t)(offset - addr);
      }
 
      template <typename Allocator>
      GAIA_NODISCARD static uint8_t* alloc(size_t cnt) noexcept {
        const auto bytes = get_min_byte_size(0, cnt);
        auto* pData = (ValueType*)mem::AllocHelper::alloc<uint8_t, Allocator>(bytes);
        core::call_ctor_raw_n(pData, cnt);
        return (uint8_t*)pData;
      }
 
      template <typename Allocator>
      static void free(void* pData, size_t cap, size_t cnt) noexcept {
        if (pData == nullptr)
          return;
        core::call_dtor_n((ValueType*)pData, cnt);
        GAIA_MEM_SANI_DEL_BLOCK(sizeof(ValueType), pData, cap, cnt);
        return mem::AllocHelper::free<Allocator>(pData);
      }
 
      template <typename Allocator>
      static void free(void* pData, size_t cnt) noexcept {
        if (pData == nullptr)
          return;
        core::call_dtor_n((ValueType*)pData, cnt);
        return mem::AllocHelper::free<Allocator>(pData);
      }
 
      GAIA_NODISCARD constexpr static ValueType get_value(std::span<const ValueType> s, size_t idx) noexcept {
        return s[idx];
      }
 
      GAIA_NODISCARD constexpr static const ValueType& get(std::span<const ValueType> s, size_t idx) noexcept {
        return s[idx];
      }
 
      GAIA_NODISCARD constexpr static ValueType& set(std::span<ValueType> s, size_t idx) noexcept {
        return s[idx];
      }
    };
 
    template <typename ValueType>
    struct data_view_policy<DataLayout::AoS, ValueType>: data_view_policy_aos<ValueType> {};
 
    template <typename ValueType>
    struct data_view_policy_aos_get {
      using view_policy = data_view_policy_aos<ValueType>;
 
      std::span<const uint8_t> m_data;
 
      data_view_policy_aos_get(std::span<uint8_t> data): m_data({(const uint8_t*)data.data(), data.size()}) {}
      data_view_policy_aos_get(std::span<const uint8_t> data): m_data({(const uint8_t*)data.data(), data.size()}) {}
      template <typename C>
      data_view_policy_aos_get(const C& c): m_data({(const uint8_t*)c.data(), c.size()}) {
        static_assert(!std::is_same_v<C, data_view_policy_aos_get>);
      }
 
      GAIA_NODISCARD const ValueType& operator[](size_t idx) const noexcept {
        GAIA_ASSERT(idx < m_data.size());
        return ((const ValueType*)m_data.data())[idx];
      }
 
      GAIA_NODISCARD decltype(auto) data() const noexcept {
        return (const uint8_t*)m_data.data();
      }
 
      GAIA_NODISCARD auto size() const noexcept {
        return m_data.size();
      }
    };
 
    template <typename ValueType>
    struct data_view_policy_get<DataLayout::AoS, ValueType>: data_view_policy_aos_get<ValueType> {};
 
    template <typename ValueType>
    struct data_view_policy_aos_set {
      using view_policy = data_view_policy_aos<ValueType>;
 
      std::span<uint8_t> m_data;
 
      data_view_policy_aos_set(std::span<uint8_t> data): m_data({(uint8_t*)data.data(), data.size()}) {}
      data_view_policy_aos_set(std::span<const uint8_t> data): m_data({(uint8_t*)data.data(), data.size()}) {}
      template <typename C>
      data_view_policy_aos_set(const C& c): m_data({(uint8_t*)c.data(), c.size()}) {
        static_assert(!std::is_same_v<C, data_view_policy_aos_set>);
      }
 
      GAIA_NODISCARD ValueType& operator[](size_t idx) noexcept {
        GAIA_ASSERT(idx < m_data.size());
        return ((ValueType*)m_data.data())[idx];
      }
 
      GAIA_NODISCARD const ValueType& operator[](size_t idx) const noexcept {
        GAIA_ASSERT(idx < m_data.size());
        return ((const ValueType*)m_data.data())[idx];
      }
 
      GAIA_NODISCARD auto data() const noexcept {
        return m_data.data();
      }
 
      GAIA_NODISCARD auto size() const noexcept {
        return m_data.size();
      }
    };
 
    template <typename ValueType>
    struct data_view_policy_set<DataLayout::AoS, ValueType>: data_view_policy_aos_set<ValueType> {};
 
    struct data_view_policy_soa_erased {
      GAIA_NODISCARD static const uint8_t*
      get(const void* pData, uint32_t alignment, std::span<const uint8_t> fieldSizes, uint32_t fieldIdx, uint32_t row,
          uint32_t capacity) noexcept {
        GAIA_ASSERT(pData != nullptr);
        GAIA_ASSERT(alignment != 0 && fieldIdx < fieldSizes.size() && row < capacity);
 
        auto address = (uintptr_t)pData;
        GAIA_FOR(fieldIdx) {
          address = detail::get_aligned_byte_offset(address, alignment, fieldSizes[i], capacity);
        }
        address += mem::padding(address, alignment);
        return (const uint8_t*)address + ((uintptr_t)fieldSizes[fieldIdx] * row);
      }
 
      GAIA_NODISCARD static uint8_t*
      set(void* pData, uint32_t alignment, std::span<const uint8_t> fieldSizes, uint32_t fieldIdx, uint32_t row,
          uint32_t capacity) noexcept {
        return const_cast<uint8_t*>(get(pData, alignment, fieldSizes, fieldIdx, row, capacity));
      }
    };
 
    template <DataLayout TDataLayout, typename ValueType>
    struct data_view_policy_soa {
      static_assert(std::is_copy_assignable_v<ValueType>);
 
      using TTuple = decltype(meta::struct_to_tuple(std::declval<ValueType>()));
      using TargetCastType = uint8_t*;
 
      constexpr static DataLayout Layout = data_layout_properties<TDataLayout, ValueType>::Layout;
      constexpr static size_t Alignment = data_layout_properties<TDataLayout, ValueType>::Alignment;
      constexpr static size_t TTupleItems = std::tuple_size<TTuple>::value;
      static_assert(Alignment > 0U, "SoA data can't be zero-aligned");
 
      template <size_t Item>
      using value_type = typename std::tuple_element<Item, TTuple>::type;
      template <size_t Item>
      using const_value_type = typename std::add_const<value_type<Item>>::type;
 
      GAIA_NODISCARD constexpr static uint32_t get_min_byte_size(uintptr_t addr, size_t cnt) noexcept {
        const auto offset = get_aligned_byte_offset<TTupleItems>(addr, cnt);
        return (uint32_t)(offset - addr);
      }
 
      template <typename Allocator>
      GAIA_NODISCARD static uint8_t* alloc(size_t cnt) noexcept {
        const auto bytes = get_min_byte_size(0, cnt);
        return mem::AllocHelper::alloc_alig<uint8_t, Allocator>(Alignment, bytes);
      }
 
      template <typename Allocator>
      static void free(void* pData, size_t cap, size_t cnt) noexcept {
        if (pData == nullptr)
          return;
 
        mem_del_block(pData, cap, cnt);
        return mem::AllocHelper::free_alig<Allocator>(pData);
      }
 
      static void mem_add_block(void* pData, size_t cap, size_t count) {
        meta::each_member(ValueType{}, [&](auto&&... item) {
          auto address = mem::align<Alignment>((uintptr_t)pData);
          ((
               //
               GAIA_MEM_SANI_ADD_BLOCK(sizeof(item), (void*)address, cap, count),
               // Skip towards the next element and make sure the address is aligned properly
               address = mem::align<Alignment>(address + (sizeof(item) * cap))),
           ...);
        });
      }
 
      static void mem_del_block(void* pData, size_t cap, size_t count) {
        meta::each_member(ValueType{}, [&](auto&&... item) {
          auto address = mem::align<Alignment>((uintptr_t)pData);
          ((
               //
               GAIA_MEM_SANI_DEL_BLOCK(sizeof(item), (void*)address, cap, count),
               // Skip towards the next element and make sure the address is aligned properly
               address = mem::align<Alignment>(address + (sizeof(item) * cap))),
           ...);
        });
      }
 
      static void mem_push_block(void* pData, size_t cap, size_t count, size_t n) {
        meta::each_member(ValueType{}, [&](auto&&... item) {
          auto address = mem::align<Alignment>((uintptr_t)pData);
          ((
               //
               GAIA_MEM_SANI_PUSH_N(sizeof(item), (void*)address, cap, count, n),
               // Skip towards the next element and make sure the address is aligned properly
               address = mem::align<Alignment>(address + (sizeof(item) * cap))),
           ...);
        });
      }
 
      static void mem_pop_block(void* pData, size_t cap, size_t count, size_t n) {
        meta::each_member(ValueType{}, [&](auto&&... item) {
          auto address = mem::align<Alignment>((uintptr_t)pData);
          ((
               //
               GAIA_MEM_SANI_POP_N(sizeof(item), (void*)address, cap, count, n),
               // Skip towards the next element and make sure the address is aligned properly
               address = mem::align<Alignment>(address + (sizeof(item) * cap))),
           ...);
        });
      }
 
      GAIA_NODISCARD constexpr static ValueType get(std::span<const uint8_t> s, size_t idx) noexcept {
        return get_inter(meta::struct_to_tuple(ValueType{}), s, idx, std::make_index_sequence<TTupleItems>());
      }
 
      template <size_t Item>
      GAIA_NODISCARD constexpr static auto get(std::span<const uint8_t> s, size_t idx = 0) noexcept {
        const auto offset = get_aligned_byte_offset<Item>((uintptr_t)s.data(), s.size());
        const auto& ref = get_ref<const value_type<Item>>(reinterpret_cast<const uint8_t*>(offset), idx);
        return std::span{&ref, s.size() - idx};
      }
 
      class accessor {
        std::span<uint8_t> m_data;
        size_t m_idx;
 
      public:
        constexpr accessor(std::span<uint8_t> data, size_t idx): m_data(data), m_idx(idx) {}
 
        constexpr void operator=(const ValueType& val) noexcept {
          set_inter(meta::struct_to_tuple(val), m_data, m_idx, std::make_index_sequence<TTupleItems>());
        }
 
        constexpr void operator=(ValueType&& val) noexcept {
          set_inter(meta::struct_to_tuple(GAIA_MOV(val)), m_data, m_idx, std::make_index_sequence<TTupleItems>());
        }
 
        GAIA_NODISCARD constexpr operator ValueType() const noexcept {
          return get_inter(
              meta::struct_to_tuple(ValueType{}), {(const uint8_t*)m_data.data(), m_data.size()}, m_idx,
              std::make_index_sequence<TTupleItems>());
        }
      };
 
      GAIA_NODISCARD constexpr static auto set(std::span<uint8_t> s, size_t idx) noexcept {
        return accessor(s, idx);
      }
 
      template <size_t Item>
      GAIA_NODISCARD constexpr static auto set(std::span<uint8_t> s, size_t idx = 0) noexcept {
        const auto offset = get_aligned_byte_offset<Item>((uintptr_t)s.data(), s.size());
        auto& ref = get_ref<value_type<Item>>((const uint8_t*)offset, idx);
        return std::span{&ref, s.size() - idx};
      }
 
    private:
      template <size_t... Ids>
      GAIA_NODISCARD constexpr static size_t
      get_aligned_byte_offset_seq(uintptr_t address, size_t cnt, std::index_sequence<Ids...> /*no_name*/) {
        ((address = detail::get_aligned_byte_offset(address, Alignment, sizeof(value_type<Ids>), cnt)), ...);
        address += mem::padding(address, Alignment);
        return address;
      }
 
      template <uint32_t N>
      GAIA_NODISCARD constexpr static size_t get_aligned_byte_offset(uintptr_t address, size_t cnt) {
        return get_aligned_byte_offset_seq(address, cnt, std::make_index_sequence<N>());
      }
 
      template <typename TMemberType>
      GAIA_NODISCARD constexpr static TMemberType& get_ref(const uint8_t* data, size_t idx) noexcept {
        // Write the value directly to the memory address.
        // Usage of unaligned_ref is not necessary because the memory is aligned.
        auto* pCastData = (TMemberType*)data;
        return pCastData[idx];
      }
 
      template <typename Tup, size_t... Ids>
      GAIA_NODISCARD constexpr static ValueType
      get_inter(Tup&& t, std::span<const uint8_t> s, size_t idx, std::index_sequence<Ids...> /*no_name*/) noexcept {
        auto address = mem::align<Alignment>((uintptr_t)s.data());
        ((
             // Put the value at the address into our tuple. Data is aligned so we can read directly.
             std::get<Ids>(t) = get_ref<value_type<Ids>>((const uint8_t*)address, idx),
             // Skip towards the next element and make sure the address is aligned properly
             address = mem::align<Alignment>(address + (sizeof(value_type<Ids>) * s.size()))),
         ...);
        return meta::tuple_to_struct<ValueType, TTuple>(GAIA_FWD(t));
      }
 
      template <typename Tup, size_t... Ids>
      constexpr static void
      set_inter(Tup&& t, std::span<uint8_t> s, size_t idx, std::index_sequence<Ids...> /*no_name*/) noexcept {
        auto address = mem::align<Alignment>((uintptr_t)s.data());
        ((
             // Set the tuple value. Data is aligned so we can write directly.
             get_ref<value_type<Ids>>((uint8_t*)address, idx) = std::get<Ids>(t),
             // Skip towards the next element and make sure the address is aligned properly
             address = mem::align<Alignment>(address + (sizeof(value_type<Ids>) * s.size()))),
         ...);
      }
    };
 
    template <typename ValueType>
    struct data_view_policy<DataLayout::SoA, ValueType>: //
      data_view_policy_soa<DataLayout::SoA, ValueType> {};
    template <typename ValueType>
    struct data_view_policy<DataLayout::SoA8, ValueType>: //
      data_view_policy_soa<DataLayout::SoA8, ValueType> {};
    template <typename ValueType>
    struct data_view_policy<DataLayout::SoA16, ValueType>: //
      data_view_policy_soa<DataLayout::SoA16, ValueType> {};
 
    template <DataLayout TDataLayout, typename ValueType>
    struct data_view_policy_soa_get {
      static_assert(std::is_copy_assignable_v<ValueType>);
 
      using view_policy = data_view_policy_soa<TDataLayout, ValueType>;
 
      template <size_t Item>
      struct data_view_policy_idx_info {
        using const_value_type = typename view_policy::template const_value_type<Item>;
      };
 
      std::span<const uint8_t> m_data;
 
      data_view_policy_soa_get(std::span<uint8_t> data): m_data({(const uint8_t*)data.data(), data.size()}) {}
      data_view_policy_soa_get(std::span<const uint8_t> data): m_data({(const uint8_t*)data.data(), data.size()}) {}
      template <typename C>
      data_view_policy_soa_get(const C& c): m_data({(const uint8_t*)c.data(), c.size()}) {
        static_assert(!std::is_same_v<C, data_view_policy_soa_get>);
      }
 
      GAIA_NODISCARD constexpr decltype(auto) operator[](size_t idx) const noexcept {
        return view_policy::get(m_data, idx);
      }
 
      template <size_t Item>
      GAIA_NODISCARD constexpr auto get() const noexcept {
        auto s = view_policy::template get<Item>(m_data);
        return std::span(s.data(), s.size());
      }
 
      GAIA_NODISCARD decltype(auto) data() const noexcept {
        return (const uint8_t*)m_data.data();
      }
 
      GAIA_NODISCARD auto size() const noexcept {
        return m_data.size();
      }
    };
 
    template <typename ValueType>
    struct data_view_policy_get<DataLayout::SoA, ValueType>: //
      data_view_policy_soa_get<DataLayout::SoA, ValueType> {};
    template <typename ValueType>
    struct data_view_policy_get<DataLayout::SoA8, ValueType>: //
      data_view_policy_soa_get<DataLayout::SoA8, ValueType> {};
    template <typename ValueType>
    struct data_view_policy_get<DataLayout::SoA16, ValueType>: //
        data_view_policy_soa_get<DataLayout::SoA16, ValueType> {};
 
    template <DataLayout TDataLayout, typename ValueType>
    struct data_view_policy_soa_set {
      static_assert(std::is_copy_assignable_v<ValueType>);
 
      using view_policy = data_view_policy_soa<TDataLayout, ValueType>;
 
      template <size_t Item>
      struct data_view_policy_idx_info {
        using value_type = typename view_policy::template value_type<Item>;
        using const_value_type = typename view_policy::template const_value_type<Item>;
      };
 
      std::span<uint8_t> m_data;
 
      data_view_policy_soa_set(std::span<uint8_t> data): m_data({(uint8_t*)data.data(), data.size()}) {}
      data_view_policy_soa_set(std::span<const uint8_t> data): m_data({(uint8_t*)data.data(), data.size()}) {}
      template <typename C>
      data_view_policy_soa_set(const C& c): m_data({(uint8_t*)c.data(), c.size()}) {
        static_assert(!std::is_same_v<C, data_view_policy_soa_set>);
      }
 
      struct accessor {
        std::span<uint8_t> m_data;
        size_t m_idx;
 
        constexpr void operator=(const ValueType& val) noexcept {
          view_policy::set(m_data, m_idx) = val;
        }
        constexpr void operator=(ValueType&& val) noexcept {
          view_policy::set(m_data, m_idx) = GAIA_FWD(val);
        }
      };
 
      GAIA_NODISCARD constexpr decltype(auto) operator[](size_t idx) const noexcept {
        return view_policy::get({(const uint8_t*)m_data.data(), m_data.size()}, idx);
      }
      GAIA_NODISCARD constexpr auto operator[](size_t idx) noexcept {
        return accessor{m_data, idx};
      }
 
      template <size_t Item>
      GAIA_NODISCARD constexpr auto get() const noexcept {
        auto s = view_policy::template get<Item>(m_data);
        return std::span(s.data(), s.size());
      }
 
      template <size_t Item>
      GAIA_NODISCARD constexpr auto set() noexcept {
        auto s = view_policy::template set<Item>(m_data);
        return std::span(s.data(), s.size());
      }
 
      GAIA_NODISCARD auto data() const noexcept {
        return m_data.data();
      }
 
      GAIA_NODISCARD auto size() const noexcept {
        return m_data.size();
      }
    };
 
    template <typename ValueType>
    struct data_view_policy_set<DataLayout::SoA, ValueType>: //
      data_view_policy_soa_set<DataLayout::SoA, ValueType> {};
    template <typename ValueType>
    struct data_view_policy_set<DataLayout::SoA8, ValueType>: //
      data_view_policy_soa_set<DataLayout::SoA8, ValueType> {};
    template <typename ValueType>
    struct data_view_policy_set<DataLayout::SoA16, ValueType>: //
        data_view_policy_soa_set<DataLayout::SoA16, ValueType> {};
 
    //----------------------------------------------------------------------
    // Helpers
    //----------------------------------------------------------------------
 
    namespace detail {
      template <typename, typename = void>
      struct auto_view_policy_inter {
        static constexpr DataLayout data_layout_type = DataLayout::AoS;
      };
      template <typename T>
      struct auto_view_policy_inter<T, std::void_t<decltype(T::gaia_Data_Layout)>> {
        static constexpr DataLayout data_layout_type = T::gaia_Data_Layout;
      };
 
      template <typename, typename = void>
      struct is_soa_layout: std::false_type {};
      template <typename T>
      struct is_soa_layout<T, std::void_t<decltype(T::gaia_Data_Layout)>>:
          std::bool_constant<!std::is_empty_v<T> && (T::gaia_Data_Layout != DataLayout::AoS)> {};
    } // namespace detail
 
    template <typename T>
    using auto_view_policy = data_view_policy<detail::auto_view_policy_inter<T>::data_layout_type, T>;
    template <typename T>
    using auto_view_policy_get = data_view_policy_get<detail::auto_view_policy_inter<T>::data_layout_type, T>;
    template <typename T>
    using auto_view_policy_set = data_view_policy_set<detail::auto_view_policy_inter<T>::data_layout_type, T>;
 
    template <typename T>
    inline constexpr bool is_soa_layout_v = detail::is_soa_layout<T>::value;
 
  } // namespace mem
} // namespace gaia