Always call MemUsage::Collect to collect metrics from a field (#4480)

Previously Collect() was used for types that implemented
CollectMemUsage() but otherwise Add() was used. This required the caller
to think about the type of the field and know/decide which method to
use.

Now, the caller always uses Collect() unless they are adding specific
byte values, in which case Add is used. Typically then, Add will only be
used to implement the CollectMemUsage() function.

To do this we require all Collect() methods to be templates so that they
all be a single overload set. The Collect on BumpPtrAllocator is
converted to a template that checks
`std::same_as<llvm::BumpPtrAllocator, T>`.

toolchain/base/mem_usage.h

@@ -47,27 +47,28 @@ class MemUsage {
                           .reserved_bytes = reserved_bytes});
   }
 
-  // Adds usage tracking for an allocator.
-  auto Add(std::string label, const llvm::BumpPtrAllocator& allocator) -> void {
+  // Adds memory usage for a `llvm::BumpPtrAllocator`.
+  auto Collect(std::string label, const llvm::BumpPtrAllocator& allocator)
+      -> void {
     Add(std::move(label), allocator.getBytesAllocated(),
         allocator.getTotalMemory());
   }
 
-  // Adds usage tracking for a map.
+  // Adds memory usage for a `Map`.
   template <typename KeyT, typename ValueT, ssize_t SmallSize,
             typename KeyContextT>
-  auto Add(std::string label, Map<KeyT, ValueT, SmallSize, KeyContextT> map,
-           KeyContextT key_context = KeyContextT()) -> void {
+  auto Collect(std::string label, Map<KeyT, ValueT, SmallSize, KeyContextT> map,
+               KeyContextT key_context = KeyContextT()) -> void {
     // These don't track used bytes, so we set the same value for used and
     // reserved bytes.
     auto bytes = map.ComputeMetrics(key_context).storage_bytes;
     Add(std::move(label), bytes, bytes);
   }
 
-  // Adds usage tracking for a set.
+  // Adds memory usage for a `Set`.
   template <typename KeyT, ssize_t SmallSize, typename KeyContextT>
-  auto Add(std::string label, Set<KeyT, SmallSize, KeyContextT> set,
-           KeyContextT key_context = KeyContextT()) -> void {
+  auto Collect(std::string label, Set<KeyT, SmallSize, KeyContextT> set,
+               KeyContextT key_context = KeyContextT()) -> void {
     // These don't track used bytes, so we set the same value for used and
     // reserved bytes.
     auto bytes = set.ComputeMetrics(key_context).storage_bytes;
@@ -81,7 +82,8 @@ class MemUsage {
   // This uses SmallVector in order to get proper inference for T, which
   // ArrayRef misses.
   template <typename T, unsigned N>
-  auto Add(std::string label, const llvm::SmallVector<T, N>& array) -> void {
+  auto Collect(std::string label, const llvm::SmallVector<T, N>& array)
+      -> void {
     Add(std::move(label), array.size_in_bytes(), array.capacity_in_bytes());
   }
 
@@ -90,7 +92,11 @@ class MemUsage {
   // The expected signature of `CollectMemUsage` is above, in MemUsage class
   // comments.
   template <typename T>
-  auto Collect(llvm::StringRef label, const T& arg) -> void {
+    requires requires(MemUsage& mem_usage, llvm::StringRef label,
+                      const T& arg) {
+      { arg.CollectMemUsage(mem_usage, label) };
+    }
+  auto Collect(std::string label, const T& arg) -> void {
     arg.CollectMemUsage(*this, label);
   }
 

toolchain/base/value_store.h

@@ -96,7 +96,7 @@ class ValueStore
   // Collects memory usage of the values.
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
-    mem_usage.Add(label.str(), values_);
+    mem_usage.Collect(label.str(), values_);
   }
 
   auto array_ref() const -> llvm::ArrayRef<ValueType> { return values_; }

toolchain/lex/tokenized_buffer.cpp

@@ -373,10 +373,10 @@ auto TokenizedBuffer::AddComment(int32_t indent, int32_t start, int32_t end)
 
 auto TokenizedBuffer::CollectMemUsage(MemUsage& mem_usage,
                                       llvm::StringRef label) const -> void {
-  mem_usage.Add(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
-  mem_usage.Add(MemUsage::ConcatLabel(label, "token_infos_"), token_infos_);
-  mem_usage.Add(MemUsage::ConcatLabel(label, "line_infos_"), line_infos_);
-  mem_usage.Add(MemUsage::ConcatLabel(label, "comments_"), comments_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "token_infos_"), token_infos_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "line_infos_"), line_infos_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "comments_"), comments_);
 }
 
 auto TokenizedBuffer::SourceBufferDiagnosticConverter::ConvertLoc(

toolchain/parse/tree.cpp

@@ -71,8 +71,8 @@ auto Tree::Verify() const -> ErrorOr<Success> {
 
 auto Tree::CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
     -> void {
-  mem_usage.Add(MemUsage::ConcatLabel(label, "node_impls_"), node_impls_);
-  mem_usage.Add(MemUsage::ConcatLabel(label, "imports_"), imports_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "node_impls_"), node_impls_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "imports_"), imports_);
 }
 
 auto Tree::PostorderIterator::MakeRange(NodeId begin, NodeId end)

toolchain/parse/tree_and_subtrees.cpp

@@ -235,7 +235,8 @@ auto TreeAndSubtrees::PrintPreorder(llvm::raw_ostream& output) const -> void {
 
 auto TreeAndSubtrees::CollectMemUsage(MemUsage& mem_usage,
                                       llvm::StringRef label) const -> void {
-  mem_usage.Add(MemUsage::ConcatLabel(label, "subtree_sizes_"), subtree_sizes_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "subtree_sizes_"),
+                    subtree_sizes_);
 }
 
 auto TreeAndSubtrees::SiblingIterator::Print(llvm::raw_ostream& output) const

toolchain/sem_ir/block_value_store.h

@@ -85,8 +85,8 @@ class BlockValueStore : public Yaml::Printable<BlockValueStore<IdT>> {
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
     mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
-    mem_usage.Add(MemUsage::ConcatLabel(label, "canonical_blocks_"),
-                  canonical_blocks_, KeyContext(this));
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "canonical_blocks_"),
+                      canonical_blocks_, KeyContext(this));
   }
 
   auto size() const -> int { return values_.size(); }

toolchain/sem_ir/constant.h

@@ -105,9 +105,9 @@ class ConstantValueStore {
   // Collects memory usage of members.
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
-    mem_usage.Add(MemUsage::ConcatLabel(label, "values_"), values_);
-    mem_usage.Add(MemUsage::ConcatLabel(label, "symbolic_constants_"),
-                  symbolic_constants_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "symbolic_constants_"),
+                      symbolic_constants_);
   }
 
   // Returns the constant values mapping as an ArrayRef whose keys are
@@ -155,8 +155,8 @@ class ConstantStore {
   // Collects memory usage of members.
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
-    mem_usage.Add(MemUsage::ConcatLabel(label, "map_"), map_);
-    mem_usage.Add(MemUsage::ConcatLabel(label, "constants_"), constants_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "map_"), map_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "constants_"), constants_);
   }
 
   // Returns a copy of the constant IDs as a vector, in an arbitrary but

toolchain/sem_ir/file.cpp

@@ -151,7 +151,7 @@ auto File::OutputYaml(bool include_builtins) const -> Yaml::OutputMapping {
 
 auto File::CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
     -> void {
-  mem_usage.Add(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
   mem_usage.Collect(MemUsage::ConcatLabel(label, "entity_names_"),
                     entity_names_);
   mem_usage.Collect(MemUsage::ConcatLabel(label, "functions_"), functions_);

toolchain/sem_ir/generic.cpp

@@ -47,8 +47,8 @@ auto SpecificStore::GetOrAdd(GenericId generic_id, InstBlockId args_id)
 auto SpecificStore::CollectMemUsage(MemUsage& mem_usage,
                                     llvm::StringRef label) const -> void {
   mem_usage.Collect(MemUsage::ConcatLabel(label, "specifics_"), specifics_);
-  mem_usage.Add(MemUsage::ConcatLabel(label, "lookup_table_"), lookup_table_,
-                KeyContext(specifics_.array_ref()));
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "lookup_table_"),
+                    lookup_table_, KeyContext(specifics_.array_ref()));
 }
 
 auto GetConstantInSpecific(const File& sem_ir, SpecificId specific_id,

toolchain/sem_ir/impl.h

@@ -176,7 +176,7 @@ class ImplStore {
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
     mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
-    mem_usage.Add(MemUsage::ConcatLabel(label, "lookup_"), lookup_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "lookup_"), lookup_);
   }
 
   auto array_ref() const -> llvm::ArrayRef<Impl> { return values_.array_ref(); }

toolchain/sem_ir/inst.h

@@ -423,7 +423,7 @@ class InstStore {
   // Collects memory usage of members.
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
-    mem_usage.Add(MemUsage::ConcatLabel(label, "loc_ids_"), loc_ids_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "loc_ids_"), loc_ids_);
     mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
   }
 

toolchain/sem_ir/name_scope.h

@@ -169,7 +169,7 @@ class NameScopeStore {
   // Collects memory usage of members.
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
-    mem_usage.Collect(label, values_);
+    mem_usage.Collect(std::string(label), values_);
   }
 
  private:

toolchain/sem_ir/type.h

@@ -126,10 +126,10 @@ class TypeStore : public Yaml::Printable<TypeStore> {
 
   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
       -> void {
-    mem_usage.Add(MemUsage::ConcatLabel(label, "complete_type_info_"),
-                  complete_type_info_);
-    mem_usage.Add(MemUsage::ConcatLabel(label, "complete_types_"),
-                  complete_types_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "complete_type_info_"),
+                      complete_type_info_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "complete_types_"),
+                      complete_types_);
   }
 
  private: