Base: Add count argument to Base.summarysize (#58874)

In many cases, number of allocations is more important than memory
allocated (w.r.t. performance).

This adds a small utility to `summarysize` to allow counting the number
of unique reachable objects, which can be useful to understand, e.g., if
you are performing many small allocations or how many of the allocations
in a function ended up reachable from the result.

---------

Co-authored-by: Jameson Nash <[email protected]>

Author: topolarity

base/summarysize.jl

@@ -6,17 +6,21 @@ struct SummarySize
     frontier_i::Vector{Int}
     exclude::Any
     chargeall::Any
+    count::Bool
 end
 
 nth_pointer_isdefined(obj, i::Int) = ccall(:jl_nth_pointer_isdefined, Cint, (Any, Csize_t), obj, i-1) != 0
 get_nth_pointer(obj, i::Int) = ccall(:jl_get_nth_pointer, Any, (Any, Csize_t), obj, i-1)
 
 """
-    Base.summarysize(obj; exclude=Union{...}, chargeall=Union{...})::Int
+    Base.summarysize(obj; count = false, exclude=Union{...}, chargeall=Union{...})::Int
 
-Compute the amount of memory, in bytes, used by all unique objects reachable from the argument.
+Compute all unique objects reachable from the argument and return either their size in
+memory (in bytes) or the number of allocations they span.
 
 # Keyword Arguments
+- `count`: if false, return the total size of the objects in memory. if true, return the
+  number of allocations spanned by the object.
 - `exclude`: specifies the types of objects to exclude from the traversal.
 - `chargeall`: specifies the types of objects to always charge the size of all of their
   fields, even if those fields would normally be excluded.
@@ -33,28 +37,32 @@ julia> Base.summarysize(Ref(rand(100)))
 
 julia> sizeof(Ref(rand(100)))
 8
+
+julia> Base.summarysize(Core.svec(1.0, "testing", true); count=true)
+4
 ```
 """
 function summarysize(obj;
+                     count::Bool = false,
                      exclude = Union{DataType, Core.TypeName, Core.MethodInstance},
                      chargeall = Union{Core.TypeMapEntry, Method})
     @nospecialize obj exclude chargeall
-    ss = SummarySize(IdDict(), Any[], Int[], exclude, chargeall)
+    ss = SummarySize(IdDict(), Any[], Int[], exclude, chargeall, count)
     size::Int = ss(obj)
     while !isempty(ss.frontier_x)
         # DFS heap traversal of everything without a specialization
         # BFS heap traversal of anything with a specialization
         x = ss.frontier_x[end]
         i = ss.frontier_i[end]
         val = nothing
-        if isa(x, SimpleVector)
+        if isa(x, Core.SimpleVector)
             nf = length(x)
             if isassigned(x, i)
                 val = x[i]
             end
         elseif isa(x, GenericMemory)
             T = eltype(x)
-            if Base.allocatedinline(T)
+            if allocatedinline(T)
                 np = datatype_npointers(T)
                 nf = length(x) * np
                 idx = (i-1) ÷ np + 1
@@ -90,9 +98,9 @@ function summarysize(obj;
     return size
 end
 
-(ss::SummarySize)(@nospecialize obj) = _summarysize(ss, obj)
+(ss::SummarySize)(@nospecialize obj) = _summarysize(ss, obj, ss.count)
 # define the general case separately to make sure it is not specialized for every type
-@noinline function _summarysize(ss::SummarySize, @nospecialize obj)
+@noinline function _summarysize(ss::SummarySize, @nospecialize(obj), count::Bool)
     issingletontype(typeof(obj)) && return 0
     # NOTE: this attempts to discover multiple copies of the same immutable value,
     # and so is somewhat approximate.
@@ -112,7 +120,7 @@ end
         # 0-field mutable structs are not unique
         return gc_alignment(0)
     end
-    return sz
+    return count ? 1 : sz
 end
 
 (::SummarySize)(obj::Symbol) = 0
@@ -121,14 +129,13 @@ end
 function (ss::SummarySize)(obj::String)
     key = ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), obj)
     haskey(ss.seen, key) ? (return 0) : (ss.seen[key] = true)
-    return Core.sizeof(Int) + Core.sizeof(obj)
+    return (ss.count ? 1 : (Core.sizeof(Int) + Core.sizeof(obj)))
 end
 
 function (ss::SummarySize)(obj::DataType)
     key = pointer_from_objref(obj)
     haskey(ss.seen, key) ? (return 0) : (ss.seen[key] = true)
-    size::Int = 7 * Core.sizeof(Int) + 6 * Core.sizeof(Int32)
-    size += 4 * nfields(obj) + ifelse(Sys.WORD_SIZE == 64, 4, 0)
+    size::Int = ss.count ? 1 : sizeof(DataType)
     size += ss(obj.parameters)::Int
     if isdefined(obj, :types)
         size += ss(obj.types)::Int
@@ -139,30 +146,34 @@ end
 function (ss::SummarySize)(obj::Core.TypeName)
     key = pointer_from_objref(obj)
     haskey(ss.seen, key) ? (return 0) : (ss.seen[key] = true)
-    return Core.sizeof(obj)
+    return (ss.count ? 1 : Core.sizeof(obj))
 end
 
 function (ss::SummarySize)(obj::GenericMemory)
     haskey(ss.seen, obj) ? (return 0) : (ss.seen[obj] = true)
-    headersize = 2*sizeof(Int)
-    size::Int = headersize
+    headersize = 2 * sizeof(Int)
+    size::Int = (ss.count ? 1 : headersize)
     datakey = unsafe_convert(Ptr{Cvoid}, obj)
     if !haskey(ss.seen, datakey)
         ss.seen[datakey] = true
-        size += sizeof(obj)
+        if !ss.count
+            size += sizeof(obj)
+        elseif pointer_from_objref(obj) + 16 != datakey
+            size += 1
+        end
         T = eltype(obj)
-        if !isempty(obj) && T !== Symbol && (!Base.allocatedinline(T) || (T isa DataType && !Base.datatype_pointerfree(T)))
+        if !isempty(obj) && T !== Symbol && (!allocatedinline(T) || (T isa DataType && !datatype_pointerfree(T)))
             push!(ss.frontier_x, obj)
             push!(ss.frontier_i, 1)
         end
     end
     return size
 end
 
-function (ss::SummarySize)(obj::SimpleVector)
+function (ss::SummarySize)(obj::Core.SimpleVector)
     key = pointer_from_objref(obj)
     haskey(ss.seen, key) ? (return 0) : (ss.seen[key] = true)
-    size::Int = Core.sizeof(obj)
+    size::Int = (ss.count ? 1 : Core.sizeof(obj))
     if !isempty(obj)
         push!(ss.frontier_x, obj)
         push!(ss.frontier_i, 1)
@@ -172,7 +183,7 @@ end
 
 function (ss::SummarySize)(obj::Module)
     haskey(ss.seen, obj) ? (return 0) : (ss.seen[obj] = true)
-    size::Int = Core.sizeof(obj)
+    size::Int = (ss.count ? 1 : Core.sizeof(obj))
     for binding in names(obj, all = true)
         if isdefined(obj, binding) && !isdeprecated(obj, binding)
             value = getfield(obj, binding)
@@ -193,7 +204,7 @@ end
 
 function (ss::SummarySize)(obj::Task)
     haskey(ss.seen, obj) ? (return 0) : (ss.seen[obj] = true)
-    size::Int = Core.sizeof(obj)
+    size::Int = (ss.count ? 1 : Core.sizeof(obj))
     if isdefined(obj, :code)
         size += ss(obj.code)::Int
     end
@@ -204,4 +215,4 @@ function (ss::SummarySize)(obj::Task)
     return size
 end
 
-(ss::SummarySize)(obj::BigInt) = _summarysize(ss, obj) + obj.alloc*sizeof(Base.GMP.Limb)
+(ss::SummarySize)(obj::BigInt) = _summarysize(ss, obj, ss.count) + (ss.count ? 1 : obj.alloc * sizeof(GMP.Limb))