+#include <xbt/sysdep.h>
+
+#include "ampi.hpp"
+#include <smpi/sampi.h>
+
+class AllocTracker {
+ std::vector<size_t> memory_size_;
+ std::map</*address*/ void*, size_t> alloc_table_; // Keep track of all allocations
+
+public:
+ void record(void* ptr, size_t size)
+ {
+ alloc_table_.emplace(ptr, size);
+ if (not simgrid::s4u::this_actor::is_maestro()) {
+ auto actor = static_cast<size_t>(simgrid::s4u::this_actor::get_pid());
+ if (actor >= memory_size_.size())
+ memory_size_.resize(actor + 1, 0);
+ memory_size_[actor] += size;
+ }
+ }
+
+ void forget(void* ptr)
+ {
+ auto elm = alloc_table_.find(ptr);
+ xbt_assert(elm != alloc_table_.end());
+ if (not simgrid::s4u::this_actor::is_maestro()) {
+ auto actor = static_cast<size_t>(simgrid::s4u::this_actor::get_pid());
+ memory_size_.at(actor) -= elm->second; // size
+ }
+ alloc_table_.erase(elm);
+ }
+
+ size_t summary() const
+ {
+ if (simgrid::s4u::this_actor::is_maestro())
+ return 0;
+ auto actor = static_cast<size_t>(simgrid::s4u::this_actor::get_pid());
+ return actor < memory_size_.size() ? memory_size_[actor] : 0;
+ }
+};