BLI: improve api for getting memory count

This introduces `MemoryCount` which can be used across multiple
`MemoryCounter`. Generally, `MemoryCount` is expected to live
longer (e.g. over the entire life-time of a cache), while `MemoryCounter`
is expected to only exists when actually counting the memory.

source/blender/blenlib/BLI_memory_counter.hh

@@ -9,11 +9,29 @@
  */
 
 #include "BLI_function_ref.hh"
-#include "BLI_implicit_sharing.hh"
+#include "BLI_implicit_sharing_ptr.hh"
+#include "BLI_map.hh"
 #include "BLI_memory_counter_fwd.hh"
 #include "BLI_set.hh"
+#include "BLI_vector_set.hh"
 
-namespace blender {
+namespace blender::memory_counter {
+
+class MemoryCount {
+ public:
+  /**
+   * Get the total number of counted bytes.
+   *
+   * \note This is only a rough estimate of the actual used memory. Often, not every little bit of
+   * memory is counted, so this is generally a lower bound. The actual memory usage should not be
+   * significantly higher though.
+   */
+  int64_t total_bytes = 0;
+
+  Set<WeakImplicitSharingPtr> handled_shared_data;
+
+  void reset();
+};
 
 /**
  * #MemoryCounter helps counting the amount of memory used in cases where data is shared and should
@@ -22,20 +40,15 @@ namespace blender {
  */
 class MemoryCounter : NonCopyable, NonMovable {
  private:
-  int64_t owned_bytes_ = 0;
-  Set<const ImplicitSharingInfo *> counted_shared_data_;
+  MemoryCount &count_;
 
  public:
-  MemoryCounter() = default;
-  ~MemoryCounter();
+  MemoryCounter(MemoryCount &count);
 
   /**
    * Add bytes that are uniquely owned, i.e. not shared.
    */
-  void add(const int64_t bytes)
-  {
-    owned_bytes_ += bytes;
-  }
+  void add(const int64_t bytes);
 
   /**
    * Add (potentially) shared data which should not be counted twice.
@@ -54,18 +67,6 @@ class MemoryCounter : NonCopyable, NonMovable {
    * to use in cases where computing the number of bytes is very cheap.
    */
   void add_shared(const ImplicitSharingInfo *sharing_info, const int64_t bytes);
-
-  /**
-   * Get the total number of counted bytes.
-   *
-   * \note This is only a rough estimate of the actual used memory. Often, not every little bit of
-   * memory is counted, so this is generally a lower bound. The actual memory usage should not be
-   * significantly higher though.
-   */
-  int64_t counted_bytes() const
-  {
-    return owned_bytes_;
-  }
 };
 
-}  // namespace blender
+}  // namespace blender::memory_counter

source/blender/blenlib/BLI_memory_counter_fwd.hh

@@ -8,8 +8,14 @@
  * \ingroup bli
  */
 
-namespace blender {
+namespace blender::memory_counter {
 
+class MemoryCount;
 class MemoryCounter;
 
-}
+}  // namespace blender::memory_counter
+
+namespace blender {
+using memory_counter::MemoryCount;
+using memory_counter::MemoryCounter;
+}  // namespace blender

source/blender/blenlib/intern/memory_counter.cc

@@ -4,13 +4,13 @@
 
 #include "BLI_memory_counter.hh"
 
-namespace blender {
+namespace blender::memory_counter {
 
-MemoryCounter::~MemoryCounter()
+MemoryCounter::MemoryCounter(MemoryCount &count) : count_(count) {}
+
+void MemoryCounter::add(const int64_t bytes)
 {
-  for (const ImplicitSharingInfo *sharing_info : counted_shared_data_) {
-    sharing_info->remove_weak_user_and_delete_if_last();
-  }
+  count_.total_bytes += bytes;
 }
 
 void MemoryCounter::add_shared(const ImplicitSharingInfo *sharing_info,
@@ -21,7 +21,7 @@ void MemoryCounter::add_shared(const ImplicitSharingInfo *sharing_info,
     count_fn(*this);
     return;
   }
-  if (!counted_shared_data_.add(sharing_info)) {
+  if (!count_.handled_shared_data.add_as(sharing_info)) {
     /* Data was counted before, avoid counting it again. */
     return;
   }
@@ -36,4 +36,10 @@ void MemoryCounter::add_shared(const ImplicitSharingInfo *sharing_info, const in
   this->add_shared(sharing_info, [&](MemoryCounter &shared_memory) { shared_memory.add(bytes); });
 }
 
-}  // namespace blender
+void MemoryCount::reset()
+{
+  std::destroy_at(this);
+  new (this) MemoryCount();
+}
+
+}  // namespace blender::memory_counter

source/blender/blenlib/tests/BLI_memory_counter_test.cc

@@ -13,12 +13,13 @@ namespace blender::tests {
 
 TEST(memory_counter, Simple)
 {
-  MemoryCounter memory;
-  EXPECT_EQ(memory.counted_bytes(), 0);
+  MemoryCount memory_count;
+  MemoryCounter memory{memory_count};
+  EXPECT_EQ(memory_count.total_bytes, 0);
   memory.add(10);
-  EXPECT_EQ(memory.counted_bytes(), 10);
+  EXPECT_EQ(memory_count.total_bytes, 10);
   memory.add(10);
-  EXPECT_EQ(memory.counted_bytes(), 20);
+  EXPECT_EQ(memory_count.total_bytes, 20);
 
   const int alloc_size = 100;
   void *data1 = MEM_mallocN(alloc_size, __func__);
@@ -27,20 +28,20 @@ TEST(memory_counter, Simple)
   const ImplicitSharingPtr sharing_info2{implicit_sharing::info_for_mem_free(data2)};
 
   memory.add_shared(sharing_info1.get(), alloc_size);
-  EXPECT_EQ(memory.counted_bytes(), 120);
+  EXPECT_EQ(memory_count.total_bytes, 120);
 
   memory.add_shared(sharing_info1.get(), [&](MemoryCounter & /*shared_memory*/) { FAIL(); });
-  EXPECT_EQ(memory.counted_bytes(), 120);
+  EXPECT_EQ(memory_count.total_bytes, 120);
 
   memory.add_shared(sharing_info2.get(),
                     [&](MemoryCounter &shared_memory) { shared_memory.add(alloc_size); });
-  EXPECT_EQ(memory.counted_bytes(), 220);
+  EXPECT_EQ(memory_count.total_bytes, 220);
 
   memory.add_shared(nullptr, 1000);
-  EXPECT_EQ(memory.counted_bytes(), 1220);
+  EXPECT_EQ(memory_count.total_bytes, 1220);
 
   memory.add_shared(nullptr, 1000);
-  EXPECT_EQ(memory.counted_bytes(), 2220);
+  EXPECT_EQ(memory_count.total_bytes, 2220);
 }
 
 }  // namespace blender::tests