Introduce a kTotalMorePrecise accounting mode for Cord::EstimatedMemoryUsage(). This mode avoids double-counting blocks that a Cord references more than once; otherwise it is similar to the existing kTotal mode.

There's no change to the existing kTotal or kFairShare accounting modes.

PiperOrigin-RevId: 544378591
Change-Id: I7b4ae55cd93d631194e59a9cd0ff07f47611219e

1 files changed, 23 insertions, 1 deletions

diff --git a/absl/strings/cord_analysis.cc b/absl/strings/cord_analysis.cc
index 73d3c4e6..e859b0db 100644
--- a/absl/strings/cord_analysis.cc
+++ b/absl/strings/cord_analysis.cc
@@ -16,6 +16,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <unordered_set>
 
 #include "absl/base/attributes.h"
 #include "absl/base/config.h"
@@ -37,7 +38,7 @@ namespace cord_internal {
 namespace {
 
 // Accounting mode for analyzing memory usage.
-enum class Mode { kTotal, kFairShare };
+enum class Mode { kFairShare, kTotal, kTotalMorePrecise };
 
 // CordRepRef holds a `const CordRep*` reference in rep, and depending on mode,
 // holds a 'fraction' representing a cumulative inverse refcount weight.
@@ -62,6 +63,23 @@ struct RawUsage {
   void Add(size_t size, CordRepRef<mode>) { total += size; }
 };
 
+// Overloaded representation of RawUsage that tracks the set of objects
+// counted, and avoids double-counting objects referenced more than once
+// by the same Cord.
+template <>
+struct RawUsage<Mode::kTotalMorePrecise> {
+  size_t total = 0;
+  // TODO(b/289250880): Replace this with a flat_hash_set.
+  std::unordered_set<const CordRep*> counted;
+
+  void Add(size_t size, CordRepRef<Mode::kTotalMorePrecise> repref) {
+    if (counted.find(repref.rep) == counted.end()) {
+      counted.insert(repref.rep);
+      total += size;
+    }
+  }
+};
+
 // Returns n / refcount avoiding a div for the common refcount == 1.
 template <typename refcount_t>
 double MaybeDiv(double d, refcount_t refcount) {
@@ -183,6 +201,10 @@ size_t GetEstimatedFairShareMemoryUsage(const CordRep* rep) {
   return GetEstimatedUsage<Mode::kFairShare>(rep);
 }
 
+size_t GetMorePreciseMemoryUsage(const CordRep* rep) {
+  return GetEstimatedUsage<Mode::kTotalMorePrecise>(rep);
+}
+
 }  // namespace cord_internal
 ABSL_NAMESPACE_END
 }  // namespace absl