soa_aos_aosao_cpp.cpp

#include <boost/pfr.hpp>
#include <vector>
#include <tuple>
#include <type_traits>
#include <cassert>
#include <array>

namespace storage_policy
{
    struct aos;
    struct soa;
    template <std::size_t ChunkSize> struct aosoa;
}

template <typename, typename> struct storage;

// ---------------------------------------------------------------------------------------------

template <typename T>
struct storage<storage_policy::aos, T>
{
    std::vector<T> _data;
 
    void add(auto&&... xs)  {  _data.push_back({xs...}); }
    auto get(std::size_t i) { return boost::pfr::structure_tie(_data[i]); }

    void call(auto&& f)
    {
        for (auto& c : _data) std::apply(f, boost::pfr::structure_tie(c));
    }
};

// ---------------------------------------------------------------------------------------------

template <typename> struct aos_to_soa;

template <typename... Ts> 
struct aos_to_soa<std::tuple<Ts...>>
    : std::type_identity<std::tuple<std::vector<Ts>...>> { };

template <typename T>
struct storage<storage_policy::soa, T>
{
    typename aos_to_soa<
        decltype(boost::pfr::structure_to_tuple(std::declval<const T&>()))
    >::type _data;

    void add(auto&&... xs) 
    { 
        std::apply([&](auto&&... vs) { (vs.emplace_back(xs), ...); }, _data);
    }

    auto get(std::size_t i) 
    { 
        using tuple_type = decltype(boost::pfr::structure_tie(std::declval<const T&>()));
        return std::apply([&](auto&&... vs) { return tuple_type{vs[i]...}; }, _data);
    }

    void call(auto&& f)
    {
        std::apply([&](auto&& v, auto&&... vs)
        {
            for (std::size_t i = 0; i < v.size(); ++i) f(v[i], vs[i]...);
        }, _data);
    }
};

// ---------------------------------------------------------------------------------------------

template <std::size_t, typename> struct aos_to_aosoa;

template <std::size_t ChunkSize, typename... Ts> 
struct aos_to_aosoa<ChunkSize, std::tuple<Ts...>>
    : std::type_identity<std::tuple<std::array<Ts, ChunkSize>...>> { };

template <std::size_t ChunkSize, typename T>
struct storage<storage_policy::aosoa<ChunkSize>, T>
{
    using member_tuple = decltype(boost::pfr::structure_to_tuple(std::declval<const T&>()));
    using chunkified = typename aos_to_aosoa<ChunkSize, member_tuple>::type;
    
    std::array<chunkified, 256 /* TODO: unhardcode */> _data{};
    std::size_t _next = 0;

    auto do_with_elem(std::size_t elem_idx, auto&& f)
    {
        auto& chunk = _data[elem_idx / ChunkSize];

        return [&]<auto... Is>(std::index_sequence<Is...>)
        {
            return f(std::get<Is>(chunk)[elem_idx % ChunkSize]...);
        }(std::make_index_sequence<std::tuple_size_v<member_tuple>>{});
    }

    void add(auto&&... xs) 
    { 
        do_with_elem(_next++, [&](auto&&... members){ ((members = xs), ...); });
    }

    auto get(std::size_t i) 
    { 
        return do_with_elem(i, [&](auto&&... members){ return std::tie(members...); });
    }

    void call(auto&& f)
    {
        for (std::size_t i = 0; i < _next; ++i)
            do_with_elem(i, [&](auto&&... members){ f(members...); });
    }
};

// ---------------------------------------------------------------------------------------------

struct c_physics { float pos, vel, acc; };

void integrate(float& pos, float& vel, float& acc)
{
    vel += acc; 
    pos += vel;
}

// ---------------------------------------------------------------------------------------------

int main()
{
    storage<storage_policy::aos, c_physics> aos_physics;
    // [{pos0, vel0, acc0}, {pos1, vel1, acc1}, ...]
    aos_physics.add(3.f, 2.f, 1.f);
    aos_physics.call(&integrate);
    auto [p0, v0, a0] = aos_physics.get(0);
    assert(p0 == 6.f);

    storage<storage_policy::soa, c_physics> soa_physics;
    // {[pos0, pos1, ...], [vel0, vel1, ...], [acc0, acc1, ...]}
    soa_physics.add(3.f, 2.f, 1.f);
    soa_physics.call(&integrate);
    auto [p1, v1, a1] = soa_physics.get(0);
    assert(p1 == 6.f);

    storage<storage_policy::aosoa<4>, c_physics> aosoa_physics;
    // [{[pos0, pos1, pos2, pos3], [vel0, vel1, vel2, vel3], [acc0, acc1, acc2, acc3]}, {[pos4, pos5, pos6, pos7], ...}}
    aosoa_physics.add(3.f, 2.f, 1.f);
    aosoa_physics.call(&integrate);
    auto [p2, v2, a2] = aosoa_physics.get(0);
    assert(p2 == 6.f);
}