From 748150c56df8e29ce125d8aa6ec2a1459674f257 Mon Sep 17 00:00:00 2001
From: Bernhard Manfred Gruber <bernhardmgruber@gmail.com>
Date: Fri, 11 Nov 2022 00:56:27 +0100
Subject: [PATCH] Support stateful accessors

Fixes: #604
---
 include/llama/View.hpp | 20 ++++++++++++--------
 tests/view.cpp         | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/include/llama/View.hpp b/include/llama/View.hpp
index 69a07c11d3..78794a4e9d 100644
--- a/include/llama/View.hpp
+++ b/include/llama/View.hpp
@@ -54,10 +54,11 @@ namespace llama
     LLAMA_FN_HOST_ACC_INLINE auto allocViewUninitialized(
         Mapping mapping = {},
         const Allocator& alloc = {},
-        Accessor = {}) -> View<Mapping, internal::AllocatorBlobType<Allocator, typename Mapping::RecordDim>, Accessor>
+        Accessor accessor = {})
+        -> View<Mapping, internal::AllocatorBlobType<Allocator, typename Mapping::RecordDim>, Accessor>
     {
         auto blobs = internal::makeBlobArray(alloc, mapping, std::make_index_sequence<Mapping::blobCount>{});
-        return {std::move(mapping), std::move(blobs)};
+        return {std::move(mapping), std::move(blobs), std::move(accessor)};
     }
 
     namespace internal
@@ -407,7 +408,6 @@ namespace llama
 #endif
     {
         static_assert(!std::is_const_v<TMapping>);
-        static_assert(std::is_empty_v<TAccessor>, "Stateful accessors are not implemented");
         using Mapping = TMapping;
         using BlobType = TBlobType;
         using ArrayExtents = typename Mapping::ArrayExtents;
@@ -432,8 +432,9 @@ namespace llama
         View() = default;
 
         LLAMA_FN_HOST_ACC_INLINE
-        View(Mapping mapping, Array<BlobType, Mapping::blobCount> storageBlobs)
+        View(Mapping mapping, Array<BlobType, Mapping::blobCount> storageBlobs, Accessor accessor = {})
             : Mapping(std::move(mapping))
+            , Accessor(std::move(accessor))
             , storageBlobs(std::move(storageBlobs))
         {
         }
@@ -629,8 +630,8 @@ namespace llama
             std::make_index_sequence<blobCount>{});
         return llama::View<typename View::Mapping, typename decltype(blobs)::value_type, typename View::Accessor>{
             view.mapping(),
-            std::move(blobs)/*,
-            view.accessor()*/};
+            std::move(blobs),
+            view.accessor()};
     }
 
     /// Creates a shallow copy of a view. This copy must not outlive the view, since it references its blob array.
@@ -659,9 +660,12 @@ namespace llama
     // \param view A view which's mapping and blobs are copied into a new view with the different accessor. If you no
     // longer need the old view, consider moving it into the argument of this function.
     template<typename NewAccessor, typename Mapping, typename BlobType, typename OldAccessor>
-    LLAMA_FN_HOST_ACC_INLINE auto withAccessor(View<Mapping, BlobType, OldAccessor> view)
+    LLAMA_FN_HOST_ACC_INLINE auto withAccessor(View<Mapping, BlobType, OldAccessor> view, NewAccessor newAccessor = {})
     {
-        return View<Mapping, BlobType, NewAccessor>{std::move(view.mapping()), std::move(view.storageBlobs)};
+        return View<Mapping, BlobType, NewAccessor>{
+            std::move(view.mapping()),
+            std::move(view.storageBlobs),
+            std::move(newAccessor)};
     }
 
     /// Like a \ref View, but array indices are shifted.
diff --git a/tests/view.cpp b/tests/view.cpp
index 62f1bc6229..418f84d739 100644
--- a/tests/view.cpp
+++ b/tests/view.cpp
@@ -429,4 +429,39 @@ TEMPLATE_TEST_CASE("view.withAccessor.shallowCopy.Restrict", "", llama::bloballo
     auto view2 = llama::withAccessor<llama::accessor::Restrict>(llama::shallowCopy(view));
     iotaFillView(view2);
     iotaCheckView(view);
-}
\ No newline at end of file
+}
+
+namespace
+{
+    struct OffsetFloatAccessor
+    {
+        float offset;
+
+        template<typename T>
+        auto operator()(T& ref) -> decltype(auto)
+        {
+            if constexpr(std::is_same_v<T, float>)
+                return ref + offset;
+            else
+                return ref;
+        }
+    };
+} // namespace
+
+TEST_CASE("view.withAccessor.OffsetFloatAccessor")
+{
+    auto view
+        = llama::allocView(llama::mapping::AoS{llama::ArrayExtents{4}, Particle{}}, llama::bloballoc::SharedPtr{});
+    view(0)(tag::Pos{})(tag::X{}) = 2.0;
+    view(0)(tag::Mass{}) = 2.0f;
+
+    auto view2 = llama::withAccessor(view, OffsetFloatAccessor{3});
+
+    CHECK(view2(0)(tag::Pos{})(tag::X{}) == 2.0);
+    CHECK(view2(0)(tag::Mass{}) == 5.0f);
+
+    view2.accessor().offset = 10;
+
+    CHECK(view2(0)(tag::Pos{})(tag::X{}) == 2.0);
+    CHECK(view2(0)(tag::Mass{}) == 12.0f);
+}