Base: Add count argument to Base.summarysize

Frequently the number of allocations can be more important for
performance than the amount of memory allocated.

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.

Author: topolarity

base/summarysize.jl

@@ -6,17 +6,20 @@ 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 their quantity.
 
 # Keyword Arguments
+- `count`: if true, returns object quantity rather than total object size.
 - `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,29 +36,33 @@ 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)
-                np = datatype_npointers(T)
+                np = Base.datatype_npointers(T)
                 nf = length(x) * np
                 idx = (i-1) ÷ np + 1
                 if @inbounds @inline isassigned(x, idx)
@@ -72,7 +79,7 @@ function summarysize(obj;
                 end
             end
         else
-            nf = datatype_npointers(typeof(x))
+            nf = Base.datatype_npointers(typeof(x))
             if nth_pointer_isdefined(x, i)
                 val = get_nth_pointer(x, i)
             end
@@ -90,15 +97,15 @@ 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)
-    issingletontype(typeof(obj)) && return 0
+@noinline function _summarysize(ss::SummarySize, @nospecialize(obj), count::Bool)
+    Base.issingletontype(typeof(obj)) && return 0
     # NOTE: this attempts to discover multiple copies of the same immutable value,
     # and so is somewhat approximate.
     key = ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), obj)
     haskey(ss.seen, key) ? (return 0) : (ss.seen[key] = true)
-    if datatype_npointers(typeof(obj)) > 0
+    if Base.datatype_npointers(typeof(obj)) > 0
         push!(ss.frontier_x, obj)
         push!(ss.frontier_i, 1)
     end
@@ -112,7 +119,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 +128,14 @@ 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 : (7 * Core.sizeof(Int) + 6 * Core.sizeof(Int32)))
+    size += (ss.count ? 1 : (4 * nfields(obj) + ifelse(Sys.WORD_SIZE == 64, 4, 0)))
     size += ss(obj.parameters)::Int
     if isdefined(obj, :types)
         size += ss(obj.types)::Int
@@ -139,17 +146,21 @@ 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
-    datakey = unsafe_convert(Ptr{Cvoid}, obj)
+    headersize = 2 * sizeof(Int)
+    size::Int = (ss.count ? 1 : headersize)
+    datakey = Base.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) + headersize != datakey
+            size += 1
+        end
         T = eltype(obj)
         if !isempty(obj) && T !== Symbol && (!Base.allocatedinline(T) || (T isa DataType && !Base.datatype_pointerfree(T)))
             push!(ss.frontier_x, obj)
@@ -159,10 +170,10 @@ function (ss::SummarySize)(obj::GenericMemory)
     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) + obj.alloc * (ss.count ? 1 : sizeof(Base.GMP.Limb))