diff --git a/include/ygm/container/array.hpp b/include/ygm/container/array.hpp index 43628367..e3f440a1 100644 --- a/include/ygm/container/array.hpp +++ b/include/ygm/container/array.hpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -36,11 +37,13 @@ class array public detail::base_misc, std::tuple>, public detail::base_async_visit, std::tuple>, + public detail::base_iterators>, public detail::base_iteration_key_value, std::tuple>, public detail::base_async_reduce, std::tuple> { - friend struct detail::base_misc, std::tuple>; + friend struct detail::base_misc, + std::tuple>; public: using self_type = array; @@ -57,6 +60,86 @@ class array using detail::base_async_insert_key_value, for_all_args>::async_insert; + /* + * @brief Proxy class for returning pair-like objects that contain references + */ + template + struct iterator_proxy { + public: + using value_ref_type = std::conditional_t; + + iterator_proxy(const key_type i, value_ref_type v) : index(i), value(v) {}; + + /* + * @brief `get()` method for tuple-like access + * + * Allows use of structured bindings + */ + template + decltype(auto) get() const { + if constexpr (N == 0) + return index; + else if constexpr (N == 1) + return value; + } + + const key_type index; + value_ref_type value; + }; + + /* + * @brief Iterator for array that gives access to items and their indices + */ + template + class array_iterator { + public: + using value_type = std::pair; + using iterator_proxy_type = iterator_proxy; + + array_iterator(self_type* arr, const key_type offset, const key_type index) + : p_arr(arr), m_offset(offset), m_index(index) {}; + + iterator_proxy_type operator*() const { + return iterator_proxy_type(m_index + m_offset, + p_arr->m_local_vec[m_index]); + } + + struct arrow_proxy { + iterator_proxy_type m_proxy; + iterator_proxy_type* operator->() { return &m_proxy; } + }; + + arrow_proxy operator->() const { return arrow_proxy{**this}; } + + array_iterator& operator++() { + m_index++; + + return *this; + } + + array_iterator operator++(int) { + iterator tmp(*this); + ++(*this); + return tmp; + } + + bool operator==(const array_iterator& other) { + return m_index == other.m_index; + } + + bool operator!=(const array_iterator& other) { + return m_index != other.m_index; + } + + private: + self_type* p_arr; + key_type m_offset; + key_type m_index; + }; + + using iterator = array_iterator; + using const_iterator = array_iterator; + array() = delete; /** @@ -404,6 +487,69 @@ class array return *this; } + /** + * @brief Access to begin iterator of locally-held items + * + * @return Local iterator to beginning of items held by process. + * @details Does not call `barrier()`. + */ + iterator local_begin() { + return iterator(this, partitioner.local_start(), 0); + } + + /** + * @brief Access to begin const_iterator of locally-held items for const array + * + * @return Local const iterator to beginning of items held by process. + * @details Does not call `barrier()`. + */ + const_iterator local_begin() const { + return const_iterator(this, partitioner.local_start(), 0); + } + + /** + * @brief Access to begin const_iterator of locally-held items for const array + * + * @return Local const iterator to beginning of items held by process. + * @details Does not call `barrier()`. + */ + const_iterator local_cbegin() const { + return const_iterator(const_cast(this), + partitioner.local_start(), 0); + } + + /** + * @brief Access to end iterator of locally-held items + * + * @return Local iterator to ending of items held by process. + * @details Does not call `barrier()`. + */ + iterator local_end() { + return iterator(this, partitioner.local_start(), partitioner.local_size()); + } + + /** + * @brief Access to end const_iterator of locally-held items for const array + * + * @return Local const iterator to ending of items held by process. + * @details Does not call `barrier()`. + */ + const_iterator local_end() const { + return const_iterator(this, partitioner.local_start(), + partitioner.local_size()); + } + + /** + * @brief Access to end const_iterator of locally-held items for const array + * + * @return Local const iterator to ending of items held by process. + * @details Does not call `barrier()`. + */ + const_iterator local_cend() const { + return const_iterator(const_cast(this), + partitioner.local_start(), partitioner.local_size()); + } + /** * @brief Insert a key and value into local storage. * diff --git a/include/ygm/container/detail/block_partitioner.hpp b/include/ygm/container/detail/block_partitioner.hpp index 282d9273..77cb502e 100644 --- a/include/ygm/container/detail/block_partitioner.hpp +++ b/include/ygm/container/detail/block_partitioner.hpp @@ -91,7 +91,7 @@ struct block_partitioner { * @param global_index Index into global container * @return Index into local storage */ - index_type local_index(const index_type &global_index) { + index_type local_index(const index_type &global_index) const { index_type to_return = global_index - m_local_start_index; YGM_ASSERT_RELEASE((to_return >= 0) && (to_return < m_local_size)); return to_return; @@ -104,7 +104,7 @@ struct block_partitioner { * @param local_index Index into locally-held items * @return Global index associated to local item index */ - index_type global_index(const index_type &local_index) { + index_type global_index(const index_type &local_index) const { index_type to_return = m_local_start_index + local_index; YGM_ASSERT_RELEASE(to_return < m_partitioned_size); return to_return; @@ -115,14 +115,14 @@ struct block_partitioner { * * @return Number of items stored on this rank */ - index_type local_size() { return m_local_size; } + index_type local_size() const { return m_local_size; } /** * @brief Global index of first local item * * @return Beginning of global index space assigned to current rank */ - index_type local_start() { return m_local_start_index; } + index_type local_start() const { return m_local_start_index; } private: int m_comm_size; diff --git a/test/test_array.cpp b/test/test_array.cpp index 03acd031..3551f4df 100644 --- a/test/test_array.cpp +++ b/test/test_array.cpp @@ -46,6 +46,22 @@ int main(int argc, char **argv) { arr.for_all([](const auto index, const auto value) { YGM_ASSERT_RELEASE(index == size_t(value)); }); + + // test range-based for + for (const auto &index_item : arr) { + YGM_ASSERT_RELEASE(index_item.index == size_t(index_item.value)); + } + + for (auto iter = arr.cbegin(); iter != arr.cend(); ++iter) { + YGM_ASSERT_RELEASE(iter->index == size_t(iter->value)); + } + + auto const_for_loop = [](const ygm::container::array &c_arr) { + for (const auto &index_item : c_arr) { + YGM_ASSERT_RELEASE(index_item.index == size_t(index_item.value)); + } + }; + const_for_loop(arr); } // Test async_binary_op_update_value @@ -306,9 +322,8 @@ int main(int argc, char **argv) { }); // Double all values in copy - arr_copy.for_all([]([[maybe_unused]] const auto &index, auto &value) { - value *= 2; - }); + arr_copy.for_all( + []([[maybe_unused]] const auto &index, auto &value) { value *= 2; }); world.barrier();