improve inference of `keys(::Dict{String})` to fix invalidations

[ranocha]

Can we improve inference of the following snippet?

julia> code_warntype(keys, (Dict{String},))
MethodInstance for keys(::Dict{String})
  from keys(a::AbstractDict) in Base at abstractdict.jl:105
Arguments
  #self#::Core.Const(keys)
  a::Dict{String}
Body::Base.KeySet{String, _A} where _A<:(AbstractDict{String})
1 ─ %1 = Base.KeySet(a)::Base.KeySet{String, _A} where _A<:(AbstractDict{String})
└──      return %1

I am using Julia v1.8.2. I think the return type should be inferred more concretely, with _A<:(Dict{String}) and not _A<:(AbstractDict{String}). This would fix a big amount of invalidations, e.g.,

julia> import Pkg; Pkg.activate(temp=true); Pkg.add("DataStructures")

julia> using SnoopCompileCore; invalidations = @snoopr(using DataStructures); using SnoopCompile

julia> length(uinvalidated(invalidations))
482

julia> trees = invalidation_trees(invalidations)
...
 inserting iterate(v::Union{Base.KeySet{<:Any, <:SwissDict}, Base.ValueIterator{<:SwissDict}}, state) in DataStructures at ~/.julia/packages/DataStructures/59MD0/src/swiss_dict.jl:646 invalidated:
...

julia> ascend(root)
Choose a call for analysis (q to quit):
     iterate(::Base.KeySet{String, _A} where _A<:(Dict{String}), ::Int64)
       iterate(::T) where T<:(Base.KeySet{String, _A} where _A<:(Dict{String}))
         copyto!(::Vector{String}, ::Base.KeySet{String, _A} where _A<:(AbstractDict{String}))
           _collect(::UnitRange{Int64}, ::Base.KeySet{String, _A} where _A<:(AbstractDict{String}), ::Base.HasElt
             collect(::Base.KeySet{String, _A} where _A<:(AbstractDict{String}))
 >             #_print#10(::Int64, ::Bool, ::Bool, ::Function, ::typeof(TOML.Internals.Printer._print), ::Nothing
...

The relevant part is

julia/stdlib/TOML/src/print.jl

Lines 101 to 104 in d498d36

	akeys = keys(a)
	if sorted
	akeys = sort!(collect(akeys); by)
	end

In the first line, akeys is inferred as Base.KeySet{String, _A} where _A<:(AbstractDict{String}). If inference gave Dict{String} here, the invalidations would be gone. I also don't see a nice way to fix these invalidations since _print accepts a::AbstractDict and some comments suggest that AbstractDict is really required, not just Dict

julia/stdlib/TOML/src/print.jl

Line 146 in d498d36

# Use runtime dispatch here since the type of value seems not to be enforced other than as AbstractDict

[KristofferC]

I also don't see a nice way to fix these invalidations

Hiding the keys call from inference could work around it I guess?

[ranocha]

Could you please give a hint how to do this?

[KristofferC]

invokelatest (as is done in other places in this file)

[ranocha]

Just to confirm I understand correctly: You propose to use

akeys = @invokelatest keys(a)

instead, correct? Then, akeys would be inferred as Any and we would get a lot more dynamic dispatch here, which could reduce runtime performance (not critical here) and lead to more invalidations (would need to be checked). That sounds like a reasonable workaround for now if we don't get new invalidations from it.

[Keno]

invokelatest is pretty expensive. Use inferencebarrier to hide things from inference.

[Keno]

Regardless, the original request is also reasonable. We should be able to squeeze a bit more precision out of apply_type_tfunc.

[Keno]

I think the commented out code here would do it actually:

julia/base/compiler/tfuncs.jl

Line 1530 in f927d25

# These blocks improve type info but make compilation a bit slower.

That code has been commented out for many years, so probably worth revisiting the benchmark (probably needs to be updated though)

[ranocha]

It looks like this part hasn't been touched in >5 years, so yes, I think it would be great to revisit it 👍

[ranocha]

It looks like the parts have been added in #20115 - but I cannot really find good benchmarks there.

[aviatesk]

Now we have a proper benchmark for Julia level compilation, so we can just comment in that code and see what "@-nanosoldier runbenchmarks("inference", vs=":master")" says.

[aviatesk]

@ranocha do you want to open a PR on this?