Mix-in file for adding Profile.print-like functionality to Profile.Allocs

Profile-Allocs-print.jl

# This file is a part of Julia. License is MIT: https://julialang.org/license

# Mix-in functionality to Profile.Allocs to give it a similar print output as Profile
# Activate with:
#   Base.include(Profile.Allocs, "/Users/jameson/julia/zipyard/profile_alloc_tools.jl")
# or
#   using Revise; Revise.includet(Profile.Allocs, "/Users/jameson/julia/zipyard/profile_alloc_tools.jl")
# Refs: https://github.com/JuliaLang/julia/pull/42768

global print

using Profile: Profile, ProfileFormat, StackFrameTree, print_flat, print_tree
#using Profile.Allocs: fetch, AllocResults, Alloc
using Base.StackTraces: StackFrame

#filter_c!(results) = foreach(x->filter!(y->!y.from_c, x.stacktrace), results)
#
#filter_type!(results, type) = filter!(y->y.type isa type, results)
#
#function sum_jl(results)
#    zero = (count = Int64(0), size = Int64(0))
#    map = Dict{StackFrame, typeof(zero)}()
#    for r in results
#        line = first(r.stacktrace)
#        old = get(map, line, zero)
#        new = typeof(zero)((old.count + 1, old.size + r.size))
#        map[line] = new
#    end
#    map
#end

print(; kwargs...) =
    Profile.print(stdout, fetch(); kwargs...)
print(io::IO; kwargs...) =
    Profile.print(io, fetch(); kwargs...)
print(io::IO, data::AllocResults; kwargs...) =
    Profile.print(io, data; kwargs...)

"""
    Profile.Allocs.print([io::IO = stdout,] [data::AllocResults = fetch()]; kwargs...)

Prints profiling results to `io` (by default, `stdout`). If you do not
supply a `data` vector, the internal buffer of accumulated backtraces
will be used.
"""
function Profile.print(io::IO,
        data::AllocResults,
        ;
        format = :tree,
        C = false,
        #combine = true,
        maxdepth::Int = typemax(Int),
        mincount::Int = 0,
        noisefloor = 0,
        sortedby::Symbol = :filefuncline,
        groupby::Union{Symbol,AbstractVector{Symbol}} = :none,
        recur::Symbol = :off,
        )

    pf = ProfileFormat(;C, maxdepth, mincount, noisefloor, sortedby, recur)
    Profile.print(io, data, pf, format)
    return
end

"""
    Profile.Allocs.print([io::IO = stdout,] data::AllocResults=fetch(); kwargs...)

Prints profiling results to `io`.

See `Profile.Allocs.print([io], data)` for an explanation of the valid keyword arguments.
"""
Profile.print(data::AllocResults; kwargs...) =
    Profile.print(stdout, data; kwargs...)

function Profile.print(io::IO, data::AllocResults, fmt::ProfileFormat, format::Symbol)
    cols::Int = Base.displaysize(io)[2]
    fmt.recur ∈ (:off, :flat, :flatc) || throw(ArgumentError("recur value not recognized"))
    data = data.allocs
    if format === :tree
        tree(io, data, cols, fmt)
    elseif format === :flat
        fmt.recur === :off || throw(ArgumentError("format flat only implements recur=:off"))
        flat(io, data, cols, fmt)
    else
        throw(ArgumentError("output format $(repr(format)) not recognized"))
    end
end


function parse_flat(::Type{T}, data::Vector{Alloc}, C::Bool) where T
    lilist = StackFrame[]
    n = Int[]
    m = Int[]
    lilist_idx = Dict{T, Int}()
    recursive = Set{T}()
    totalbytes = 0
    for r in data
        first = true
        empty!(recursive)
        nb = r.size # or 1 for counting
        totalbytes += nb
        for frame in r.stacktrace
            !C && frame.from_c && continue
            key = (T === UInt64 ? ip : frame)
            idx = get!(lilist_idx, key, length(lilist) + 1)
            if idx > length(lilist)
                push!(recursive, key)
                push!(lilist, frame)
                push!(n, nb)
                push!(m, 0)
            elseif !(key in recursive)
                push!(recursive, key)
                n[idx] += nb
            end
            if first
                m[idx] += nb
                first = false
            end
        end
    end
    @assert length(lilist) == length(n) == length(m) == length(lilist_idx)
    return (lilist, n, m, totalbytes)
end

function flat(io::IO, data::Vector{Alloc}, cols::Int, fmt::ProfileFormat)
    fmt.combine || error(ArgumentError("combine=false"))
    lilist, n, m, totalbytes = parse_flat(fmt.combine ? StackFrame : UInt64, data, fmt.C)
    filenamemap = Dict{Symbol,String}()
    print_flat(io, lilist, n, m, cols, filenamemap, fmt)
    Base.println(io, "Total snapshots: ", length(data))
    Base.println(io, "Total bytes: ", totalbytes)
end

function tree!(root::StackFrameTree{T}, all::Vector{Alloc}, C::Bool, recur::Symbol) where {T}
    tops = Vector{StackFrameTree{T}}()
    build = Dict{T, StackFrameTree{T}}()
    for r in all
        first = true
        nb = r.size # or 1 for counting
        root.recur = 0
        root.count += nb
        parent = root
        for i in reverse(eachindex(r.stacktrace))
            frame = r.stacktrace[i]
            key = (T === UInt64 ? ip : frame)
            if (recur === :flat && !frame.from_c) || recur === :flatc
                # see if this frame already has a parent
                this = get!(build, frame, parent)
                if this !== parent
                    # Rewind the `parent` tree back, if this exact ip (FIXME) was already present *higher* in the current tree
                    push!(tops, parent)
                    parent = this
                end
            end
            !C && frame.from_c && continue
            this = get!(StackFrameTree{T}, parent.down, key)
            if recur === :off || this.recur == 0
                this.frame = frame
                this.up = parent
                this.count += nb
                this.recur = 1
            else
                this.count_recur += 1
            end
            parent = this
        end
        parent.overhead += nb
        if recur !== :off
            # We mark all visited nodes to so we'll only count those branches
            # once for each backtrace. Reset that now for the next backtrace.
            empty!(build)
            push!(tops, parent)
            for top in tops
                while top.recur != 0
                    top.max_recur < top.recur && (top.max_recur = top.recur)
                    top.recur = 0
                    top = top.up
                end
            end
            empty!(tops)
        end
        let this = parent
            while this !== root
                this.flat_count += nb
                this = this.up
            end
        end
    end
    function cleanup!(node::StackFrameTree)
        stack = [node]
        while !isempty(stack)
            node = pop!(stack)
            node.recur = 0
            empty!(node.builder_key)
            empty!(node.builder_value)
            append!(stack, values(node.down))
        end
        nothing
    end
    cleanup!(root)
    return root
end

function tree(io::IO, data::Vector{Alloc}, cols::Int, fmt::ProfileFormat)
    fmt.combine || error(ArgumentError("combine=false"))
    if fmt.combine
        root = tree!(StackFrameTree{StackFrame}(), data, fmt.C, fmt.recur)
    else
        root = tree!(StackFrameTree{UInt64}(), data, fmt.C, fmt.recur)
    end
    print_tree(io, root, cols, fmt, false)
    Base.println(io, "Total snapshots: ", length(data))
    Base.println(io, "Total bytes: ", root.count)
end

nothing