DistributedDataShared.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once
 
#include "shambase/aliases_int.hpp"
#include <functional>
#include <map>
#include <utility>
#include <vector>
 
namespace shambase {
 
    template<class T>
    class DistributedDataShared {
 
        std::multimap<std::pair<u64, u64>, T> data;
 
        using iterator = typename std::multimap<std::pair<u64, u64>, T>::iterator;
 
        public:
        inline std::multimap<std::pair<u64, u64>, T> &get_native() { return data; }
 
        inline const std::multimap<std::pair<u64, u64>, T> &get_native() const { return data; }
 
        inline auto begin() { return data.begin(); }
        inline auto end() { return data.end(); }
        inline auto begin() const { return data.begin(); }
        inline auto end() const { return data.end(); }
        inline auto cbegin() const { return data.cbegin(); }
        inline auto cend() const { return data.cend(); }
 
        inline iterator add_obj(u64 left_id, u64 right_id, T &&obj) {
            return data.emplace(std::pair<u64, u64>{left_id, right_id}, std::forward<T>(obj));
        }
 
        inline void for_each(std::function<void(u64, u64, T &)> &&f) {
            for (auto &[id, obj] : data) {
                f(id.first, id.second, obj);
            }
        }
 
        inline void for_each(std::function<void(u64, u64, const T &)> &&f) const {
            for (auto &[id, obj] : data) {
                f(id.first, id.second, obj);
            }
        }
 
        inline void tranfer_all(std::function<bool(u64, u64)> cd, DistributedDataShared &other) {
 
            std::vector<std::pair<u64, u64>> occurences;
 
            // whoa i forgot the & here and triggered the copy constructor of every patch
            // like do not forget it or it will be a disaster waiting to come
            // i did throw up a 64 GPUs run because of that
            for (auto &[k, v] : data) {
                if (cd(k.first, k.second)) {
                    occurences.push_back(k);
                }
            }
 
            for (auto p : occurences) {
                auto ext = data.extract(p);
                other.data.insert(std::move(ext));
            }
        }
 
        inline bool has_key(u64 left_id, u64 right_id) const {
            return (data.find({left_id, right_id}) != data.end());
        }
 
        inline u64 get_element_count() const { return data.size(); }
 
        template<class Tmap>
        inline DistributedDataShared<Tmap> map(std::function<Tmap(u64, u64, T &)> map_func) {
            DistributedDataShared<Tmap> ret;
            for_each([&](u64 left, u64 right, T &ref) {
                ret.add_obj(left, right, map_func(left, right, ref));
            });
            return ret;
        }
 
        template<class Tmap>
        inline DistributedDataShared<Tmap> map(
            std::function<Tmap(u64, u64, const T &)> map_func) const {
            DistributedDataShared<Tmap> ret;
            for_each([&](u64 left, u64 right, const T &ref) {
                ret.add_obj(left, right, map_func(left, right, ref));
            });
            return ret;
        }
 
        inline void reset() { data.clear(); }
 
        inline bool is_empty() const { return data.empty(); }
    };
 
} // namespace shambase