Improve performance of depth-limited type printing in stack traces

[charleskawczynski]

Can we improve the performance of depth-limited type printing in stack traces? Here is an example where printing truncated types takes quite a while with a highly nested struct:

Base.@kwdef struct Nested{A,B}
    num::Int = 1
end
nest_val(na, nb, ::Val{1}) = Nested{na, nb}()
nest_val(na, nb, ::Val{n}) where {n} = nest_val(Nested{na, nb}, Nested{na, nb}, Val(n-1))
nest_val(na, nb, n::Int) = nest_val(na, nb, Val(n))
nest_val(n) = nest_val(1, 1, n)
foo1(t::Nested) = error("oops")
foo2(t::Nested) = foo1(t)
foo3(t::Nested) = foo2(t)
foo4(t::Nested) = foo3(t)
foo5(t::Nested) = foo4(t)
foo6(t::Nested) = foo5(t)
foo7(t::Nested) = foo6(t)
foo8(t::Nested) = foo7(t)
foo9(t::Nested) = foo8(t)
foo10(t::Nested) = foo9(t)
NV = nest_val(20); # be careful with changing 20 to a larger number
foo10(NV) # errors, takes ~45 seconds on Julia 1.10

I haven't looked at the implementation in a while, but I imagine that this could be related to #55807.

[nsajko]

Try the latest nightly, xref #55575.

[charleskawczynski]

Try the latest nightly, xref #55575.

I just tried it. It's actually slower on nightly. ~ 148 seconds

[KristofferC]

I think this is the same as #37990.

The code tries to check if this type could potentially have a type alias somewhere and in that case print the type alias. This check is very slow and it is done over and over and over.

julia/base/show.jl

Lines 729 to 754 in 1cfda3f

	function show_typealias(io::IO, name::GlobalRef, x::Type, env::SimpleVector, wheres::Vector)
	if !(get(io, :compact, false)::Bool)
	# Print module prefix unless alias is visible from module passed to
	# IOContext. If :module is not set, default to Main.
	# nothing can be used to force printing prefix.
	from = get(io, :module, Main)
	if (from === nothing || !isvisible(name.name, name.mod, from))
	show(io, name.mod)
	print(io, ".")
	end
	end
	print(io, name.name)
	isempty(env) && return
	io = IOContext(io)
	for p in wheres
	io = IOContext(io, :unionall_env => p)
	end
	orig = getfield(name.mod, name.name)
	vars = TypeVar[]
	while orig isa UnionAll
	push!(vars, orig.var)
	orig = orig.body
	end
	show_typeparams(io, env, Core.svec(vars...), wheres)
	nothing
	end

[KristofferC]

The solution here is to not form the whole type string because it is ridiculously large. Instead the printing should stop going down the nesting once we know we have gone far enough down to not make a difference to the stacktrace type printing.

[charleskawczynski]

I have something that works, and it passes the existing tests, but I'm not sure if I've correctly handled all of the edge cases.

The notable timing difference is:

@time string_type_depth_limited(typeof(nest_val(20)); maxdepth = 3)
  0.106258 seconds (191.16 k allocations: 12.450 MiB, 99.51% compilation time)
@time Base.type_depth_limit(string(typeof(nest_val(20))), _max_display_size(stdout); maxdepth = 3)
  6.357849 seconds (85.98 M allocations: 5.080 GiB, 8.09% gc time, 0.09% compilation time)

Here is my script. Should I open a PR? Or is this enough for someone else to?

#=
julia
include("performant_type_printing.jl")
=#

function _max_display_size(io::IO)
    sz = get(io, :displaysize, displaysize(io))::Tuple{Int, Int}
    return max(sz[2], 120)
end

string_type_depth_limited(@nospecialize(T), _width = nothing; maxdepth = nothing)::String =
	string_type_depth_limited(stdout, _width, T; maxdepth)

function string_type_depth_limited(io::IO, _width, @nospecialize(T); maxdepth = nothing)::String
	str_lim = _string_type_depth_limited(io, T, maxdepth)
	width = isnothing(_width) ? _max_display_size(io) : _width
	return Base.type_depth_limit(str_lim, width; maxdepth)
end

function string_typeparams(io::IO, env::Core.SimpleVector, orig::Core.SimpleVector, wheres::Vector)
    n = length(env)
    elide = length(wheres)
    function egal_var(p::TypeVar, @nospecialize o)
        return o isa TypeVar &&
            ccall(:jl_types_egal, Cint, (Any, Any), p.ub, o.ub) != 0 &&
            ccall(:jl_types_egal, Cint, (Any, Any), p.lb, o.lb) != 0
    end
    for i = n:-1:1
        p = env[i]
        if p isa TypeVar
            if i == n && egal_var(p, orig[i]) && Base.show_can_elide(p, wheres, elide, env, i)
                n -= 1
                elide -= 1
            elseif p.lb === Union{} && isgensym(p.name) && Base.show_can_elide(p, wheres, elide, env, i)
                elide -= 1
            elseif p.ub === Any && isgensym(p.name) && Base.show_can_elide(p, wheres, elide, env, i)
                elide -= 1
            end
        end
    end
    s_result = if n > 0
    	params_string = join(
    		map(1:n) do i
	            p = env[i]
	            suffix = i < n ? ", " : ""
	            if p isa TypeVar
	                if p.lb === Union{} && something(findfirst(@nospecialize(w) -> w === p, wheres), 0) > elide
	                    string("<:", p.ub, suffix)
	                elseif p.ub === Any && something(findfirst(@nospecialize(w) -> w === p, wheres), 0) > elide
	                    string(">:", p.lb, suffix)
	                else
	                    string(p, suffix)
	                end
	            else
	                string(p, suffix)
	            end
    		end
		)
    	string("{", params_string, "}")
    else
    	""
    end
    resize!(wheres, elide)
    return s_result
end

function string_type_alias_and_params(io::IO, name::Core.GlobalRef, x::Type, env::Core.SimpleVector, wheres::Vector)
	alias = if !(get(io, :compact, false)::Bool)
        # Print module prefix unless alias is visible from module passed to
        # IOContext. If :module is not set, default to Main.
        # nothing can be used to force printing prefix.
        from = get(io, :module, Main)
        if (from === nothing || !Base.isvisible(name.name, name.mod, from))
            string(name.mod, ".", name.name)
        else
        	string(name.name)
        end
    else
    	string(name.name)
    end
    if isempty(env)
		return (alias, ())
	end
    io = IOContext(io)
    for p in wheres
        io = IOContext(io, :unionall_env => p)
    end
    orig = getfield(name.mod, name.name)
    vars = TypeVar[]
    while orig isa UnionAll
        push!(vars, orig.var)
        orig = orig.body
    end
    return string(alias, string_typeparams(io, env, Core.svec(vars...), wheres))
end

function _string_type_depth_limited(io::IO, @nospecialize(T), maxdepth, depth = 0)::String
	maxdepth === nothing && return string(T)
	param_length(T) = hasproperty(T, :parameters) ? length(T.parameters) : length(T)
	if hasproperty(T, :name)
	    properT = Base.makeproper(io, T)
	    alias = Base.make_typealias(properT)

		if alias === nothing
			s = T.name.wrapper
			if depth > maxdepth
				return ""
			elseif depth == maxdepth
				return "$s{…}"
			elseif depth + 1 == maxdepth
				param_length(T) == 0 && return string(s)
				if all(x->param_length(x)==1, T.parameters)
					ps = map(T.parameters) do p
						_string_type_depth_limited(io, p, maxdepth, depth+1)
					end
					return "$s{$(join(ps, ", "))}"
				else
					return "$s{…}"
				end
			else
				param_length(T) == 0 && return string(s)
				ps = map(T.parameters) do p
					_string_type_depth_limited(io, p, maxdepth, depth+1)
				end
				return "$s{$(join(ps, ", "))}"
			end
		else
		    wheres = Base.make_wheres(io, alias[2], T)
		    return string_type_alias_and_params(io, alias[1], T, alias[2], wheres)
		end
	elseif T isa UnionAll
		return string(T)
	else
		if depth > maxdepth
			return ""
		elseif depth == maxdepth
			return "…"
		else
			return string(T)
		end
	end
end

struct F55952{T} end
Base.@kwdef struct Nested{A,B}
    num::Int = 1
end

struct Sing end
nest_val(na, nb, ::Val{1}) = Nested{na, nb}()
nest_val(na, nb, ::Val{n}) where {n} = nest_val(Nested{na, nb}, Nested{na, nb}, Val(n-1))
nest_val(na, nb, n::Int) = nest_val(na, nb, Val(n))
nest_val(n) = nest_val(1, 1, n)
foo1(t::Nested) = error("oops")
foo2(t::Nested) = foo1(t)
foo3(t::Nested) = foo2(t)
foo4(t::Nested) = foo3(t)
foo5(t::Nested) = foo4(t)
foo6(t::Nested) = foo5(t)
foo7(t::Nested) = foo6(t)
foo8(t::Nested) = foo7(t)
foo9(t::Nested) = foo8(t)
foo10(t::Nested) = foo9(t)
# foo10(nest_val(20)) # errors, takes ~45 seconds on Julia 1.10

@time string_type_depth_limited(typeof(nest_val(20)); maxdepth = 3) # 0.106258 seconds
@time Base.type_depth_limit(string(typeof(nest_val(20))), _max_display_size(stdout); maxdepth = 3) # 6.3 seconds

Base_type_depth_limit(@nospecialize(T), width; maxdepth = nothing) =
	Base.type_depth_limit(string(T), width; maxdepth)

struct F49231{a,b,c,d,e,f,g} end
(::F49231)(a,b,c) = error("oops")

using Test
@testset "Performant type printing" begin
	@testset "type_depth_limit" begin
		    for tdl in (Base_type_depth_limit, string_type_depth_limited)
		    T = typeof(view([1, 2, 3], 1:2))
		    @test tdl(T, 0, maxdepth = 1) == "SubArray{…}"
		    @test tdl(T, 0, maxdepth = 2) == "SubArray{$Int, 1, Vector{…}, Tuple{…}, true}"
		    @test tdl(T, 0, maxdepth = 3) == "SubArray{$Int, 1, Vector{$Int}, Tuple{UnitRange{…}}, true}"
		    @test tdl(T, 0, maxdepth = 4) == "SubArray{$Int, 1, Vector{$Int}, Tuple{UnitRange{$Int}}, true}"
		    @test tdl(T, 3) == "SubArray{…}"
		    @test tdl(T, 44) == "SubArray{…}"
		    @test tdl(T, 45) == "SubArray{$Int, 1, Vector{…}, Tuple{…}, true}"
		    @test tdl(T, 59) == "SubArray{$Int, 1, Vector{…}, Tuple{…}, true}"
		    @test tdl(T, 60) == "SubArray{$Int, 1, Vector{$Int}, Tuple{UnitRange{…}}, true}"
		    @test tdl(T, 100) == "SubArray{$Int, 1, Vector{$Int}, Tuple{UnitRange{$Int}}, true}"

		    T = Vector{V} where V<:AbstractVector{T} where T<:Real
		    @test tdl(T, 0, maxdepth = 1) == "Vector{…} where {…}"
		    @test tdl(T, 0, maxdepth = 2) == "Vector{V} where {T<:Real, V<:AbstractVector{…}}"
		    @test tdl(T, 0, maxdepth = 3) == "Vector{V} where {T<:Real, V<:AbstractVector{T}}"
		    @test tdl(T, 20) == "Vector{…} where {…}"
		    @test tdl(T, 46) == "Vector{…} where {…}"
		    @test tdl(T, 47) == "Vector{V} where {T<:Real, V<:AbstractVector{T}}"

		    T = F49231{Vector,Val{('}','}')},Vector{Vector{Vector{Vector}}},Tuple{Int,Int,Int,Int,Int,Int,Int},Int,Int,Int}
		    @test tdl(T, 105) == "F49231{Vector, Val{('}', '}')}, Vector{Vector{Vector{Vector}}}, NTuple{7, $Int}, $Int, $Int, $Int}"
		    @test tdl(T, 85) == "F49231{Vector, Val{…}, Vector{…}, NTuple{…}, $Int, $Int, $Int}"
		end
	end
end

[charleskawczynski]

cc @JeffBezanson ?

[charleskawczynski]

I think I might be missing recursion on the type parameters, e.g., string(">:", p.lb, suffix), which should probably need to become calls to _string_type_depth_limited. But I think that might be all that's needed.

Most of the functions above are small modifications of existing functions where, instead of these methods printing, they return strings.

[charleskawczynski]

Also, I left the width parameter, to preserve that behavior-- I'm not really clear on what purpose it serves.

[charleskawczynski]

To give some context as to what I did:

• I reused code from base/show.jl to return strings instead of printing one
• the top-level method that I wrote, string_type_depth_limited (right below), still calls Base.type_depth_limit to preserve the behavior that the width parameter enforces-- although I admit I don't fully understand why that is needed. This solution, while not optimal, still significantly reduces the worst-case scenario, which is pretty bad (based on the example given).

function string_type_depth_limited(io::IO, _width, @nospecialize(T); maxdepth = nothing)::String
	str_lim = _string_type_depth_limited(io, T, maxdepth)
	width = isnothing(_width) ? _max_display_size(io) : _width
	return Base.type_depth_limit(str_lim, width; maxdepth)
end