`KeyError: key :S not found` in `map_ssas_to_source`

[jishnub]

julia> using FillArrays, InfiniteArrays

julia> @descend Matrix(Zeros(∞,∞))
ERROR: KeyError: key :S not found
Stacktrace:
  [1] getindex
    @ ./dict.jl:484 [inlined]
  [2] map_ssas_to_source(src::Core.CodeInfo, rootnode::JuliaSyntax.TreeNode{JuliaSyntax.SyntaxData}, Δline::Int64)
    @ TypedSyntax ~/.julia/packages/TypedSyntax/JHruX/src/node.jl:657
  [3] tsn_and_mappings(m::Method, src::Core.CodeInfo, rt::Any, sourcetext::SubString{String}, lineno::Int64; warn::Bool, strip_macros::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ TypedSyntax ~/.julia/packages/TypedSyntax/JHruX/src/node.jl:47
  [4] tsn_and_mappings(m::Method, src::Core.CodeInfo, rt::Any; warn::Bool, strip_macros::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ TypedSyntax ~/.julia/packages/TypedSyntax/JHruX/src/node.jl:34
  [5] tsn_and_mappings
    @ ~/.julia/packages/TypedSyntax/JHruX/src/node.jl:28 [inlined]
  [6] #get_typed_sourcetext#32
    @ ~/Dropbox/JuliaPackages/Cthulhu.jl/src/reflection.jl:350 [inlined]
  [7] get_typed_sourcetext
    @ ~/Dropbox/JuliaPackages/Cthulhu.jl/src/reflection.jl:348 [inlined]
  [8] find_callsites(interp::Cthulhu.CthulhuInterpreter, CI::Core.CodeInfo, stmt_infos::Vector{Core.Compiler.CallInfo}, mi::Core.MethodInstance, slottypes::Vector{Any}, optimize::Bool, annotate_source::Bool)
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/reflection.jl:28
  [9] _descend(term::REPL.Terminals.TTYTerminal, interp::Cthulhu.CthulhuInterpreter, curs::Cthulhu.CthulhuCursor; override::Nothing, debuginfo::Symbol, optimize::Bool, interruptexc::Bool, iswarn::Bool, hide_type_stable::Bool, verbose::Nothing, remarks::Bool, with_effects::Bool, inline_cost::Bool, type_annotations::Bool, annotate_source::Bool)
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/Cthulhu.jl:497
 [10] _descend(term::REPL.Terminals.TTYTerminal, interp::Cthulhu.CthulhuInterpreter, mi::Core.MethodInstance; kwargs::Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:iswarn,), Tuple{Bool}}})
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/Cthulhu.jl:766
 [11] _descend(term::REPL.Terminals.TTYTerminal, args::Any; interp::Core.Compiler.NativeInterpreter, kwargs::Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:iswarn,), Tuple{Bool}}})
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/Cthulhu.jl:782
 [12] __descend_with_error_handling(args::Any; terminal::Any, kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}})
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/Cthulhu.jl:240
 [13] _descend_with_error_handling(f::Any, argtypes::Any; kwargs::Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:iswarn,), Tuple{Bool}}})
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/Cthulhu.jl:229
 [14] descend_code_typed(::Any, ::Vararg{Any}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/Cthulhu.jl:174
 [15] descend_code_typed(::Any, ::Any)
    @ Cthulhu ~/Dropbox/JuliaPackages/Cthulhu.jl/src/Cthulhu.jl:174
 [16] top-level scope
    @ REPL[8]:1

[timholy]

I recognize this is one of the CodeTracking bugs identified in #401, where it finds the wrong method:

julia> mi = badmis[end-7]
MethodInstance for (Vector)(::Vector{S}) where S

julia> mi.def
(Array{T, N} where T)(x::AbstractArray{S, N}) where {S, N}
     @ Core boot.jl:498

julia> definition(String, mi.def)
("Array{T,1}() where {T} = Array{T,1}(undef, 0)", 496)

Clearly not the same method. There's work on improving CodeTracking going on in JuliaDebug/CodeTracking.jl#108

[NHDaly]

Would it be helpful to have another MRE?

[NHDaly]

This MRE passes on julia 1.11, but fails on 1.10.

Error first, then the source file at the end:

(@v1.10) pkg> st
Status `~/.julia/environments/v1.10/Project.toml`
  [69d22d85] About v1.0.1
  [1520ce14] AbstractTrees v0.4.5
⌃ [6e4b80f9] BenchmarkTools v1.5.0
  [a2441757] Coverage v1.6.1
  [f68482b8] Cthulhu v2.16.2
  [31a5f54b] Debugger v0.7.10
  [ab62b9b5] DeepDiffs v1.2.0
⌃ [fb4d412d] FixedPointDecimals v0.5.3
  [c27321d9] Glob v1.3.1
  [92ed2492] HeapSnapshotUtils v0.1.0 `https://github.com/RelationalAI/HeapSnapshotUtils.jl#main`
  [7ec9b9c5] Humanize v1.0.0
⌃ [5903a43b] Infiltrator v1.8.3
  [70703baa] JuliaSyntax v0.4.10
  [1fcbbee2] LookingGlass v0.3.3 `~/.julia/dev/LookingGlass`
  [bdcacae8] LoopVectorization v0.12.171
  [1914dd2f] MacroTools v0.5.15
  [85b6ec6f] MethodAnalysis v0.4.13
⌃ [e4faabce] PProf v3.1.3
⌃ [14b8a8f1] PkgTemplates v0.7.52
⌃ [c46f51b8] ProfileView v1.8.0
  [92933f4c] ProgressMeter v1.10.2
  [9c30249a] RAI v0.2.9
  [817f1d60] ReTestItems v1.29.0
⌃ [295af30f] Revise v3.7.1
  [aa65fe97] SnoopCompile v3.0.2 `~/work/jl_depots/raicode2/dev/SnoopCompile`
  [e2b509da] SnoopCompileCore v3.0.0 `~/.julia/dev/SnoopCompile/SnoopCompileCore`
  [ac92255e] Speculator v0.2.0
  [1e6cf692] TestEnv v1.102.0 `~/work/jl_depots/raicode2/dev/TestEnv`
⌃ [e689c965] Tracy v0.1.3
Info Packages marked with ⌃ have new versions available and may be upgradable.

(@v1.10) pkg> add MacroTools

julia> Revise.includet("Onions/src/cthluhu-mre.jl")

julia> @descend cmp(RelNumber(1), RelNumber(2))
cmp(a::RelNumber, b::RelNumber) @ Main ~/Documents/play/rel-interpreter/Onions/src/cthluhu-mre.jl:222
222 function Base.cmp::Core.Const(cmp)(a::RelNumber::RelNumber, b::RelNumber::RelNumber)::Int64
223     return peel(b::RelNumber)::Int64 do bv; cmp_inner(a, bv) ; end
224 end
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [w]arn, [h]ide type-stable statements, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native, [j]ump to source always.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
 • peel(b::RelNumber)
   ↩
ERROR: KeyError: key Symbol("#53#T") not found
Stacktrace:
  [1] getindex
    @ ./dict.jl:498 [inlined]
  [2] map_ssas_to_source(src::Core.CodeInfo, mi::Core.MethodInstance, rootnode::JuliaSyntax.SyntaxNode, Δline::Int64)
    @ TypedSyntax ~/.julia/packages/TypedSyntax/eS4sW/src/node.jl:816
  [3] tsn_and_mappings(mi::Core.MethodInstance, src::Core.CodeInfo, rt::Any, sourcetext::SubString{…}, lineno::Int64; warn::Bool, strip_macros::Bool, kwargs::@Kwargs{})
    @ TypedSyntax ~/.julia/packages/TypedSyntax/eS4sW/src/node.jl:54
  [4] tsn_and_mappings(mi::Core.MethodInstance, src::Core.CodeInfo, rt::Any; warn::Bool, strip_macros::Bool, kwargs::@Kwargs{})
    @ TypedSyntax ~/.julia/packages/TypedSyntax/eS4sW/src/node.jl:39
  [5] tsn_and_mappings
    @ ~/.julia/packages/TypedSyntax/eS4sW/src/node.jl:32 [inlined]
  [6] #get_typed_sourcetext#23
    @ ~/.julia/packages/Cthulhu/dysgf/src/reflection.jl:377 [inlined]
  [7] get_typed_sourcetext
    @ ~/.julia/packages/Cthulhu/dysgf/src/reflection.jl:376 [inlined]
  [8] find_callsites(interp::Cthulhu.CthulhuInterpreter, CI::Core.CodeInfo, stmt_infos::Vector{…}, mi::Core.MethodInstance, slottypes::Vector{…}, optimize::Bool, annotate_source::Bool, pc2excts::Nothing)
    @ Cthulhu ~/.julia/packages/Cthulhu/dysgf/src/reflection.jl:33
  [9] 
    @ Cthulhu ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:571
 [10] 
    @ Cthulhu ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:722
 [11] _descend
    @ ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:472 [inlined]
 [12] _descend
    @ ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:875 [inlined]
 [13] #_descend#128
    @ ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:891 [inlined]
 [14] __descend_with_error_handling(args::Any; terminal::Any, kwargs...)
    @ Cthulhu ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:275
 [15] _descend_with_error_handling(f::Any, argtypes::Any; kwargs::@Kwargs{iswarn::Bool})
    @ Cthulhu ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:264
 [16] descend(::Any, ::Vararg{Any}; kwargs::@Kwargs{})
    @ Cthulhu ~/.julia/packages/Cthulhu/dysgf/src/Cthulhu.jl:318
 [17] top-level scope
    @ REPL[6]:1
Some type information was truncated. Use `show(err)` to see complete types.

Whereas on 1.11, this works fine:

julia> @descend cmp(RelNumber(1), RelNumber(2))
cmp(a::RelNumber, b::RelNumber) @ Main ~/Documents/play/rel-interpreter/Onions/src/cthluhu-mre.jl:222
222 function Base.cmp::Core.Const(cmp)(a::RelNumber::RelNumber, b::RelNumber::RelNumber)::Int64
223     return peel(b::RelNumber)::Int64 do bv; cmp_inner(a, bv) ; end
224 end
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [w]arn, [h]ide type-stable statements, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native, [j]ump to source always.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
 • peel(b::RelNumber)
   ↩
peel(func::var"#120#F", x::var"#53#T") where {var"#120#F", var"#53#T"<:RelNumber} @ Main ~/Documents/play/rel-interpreter/Onions/src/cthluhu-mre.jl:108
108 function (Onions.peel(
109             func::var"#10#11"{RelNumber}::F,
110             x::RelNumber::T
111         ) where {F, T<:$(esc(namify(T)))})::Int64
112             tag::UInt8 = getfield(x::RelNumber, :tag)::UInt8
113             return $(ifelse(
114                 [(:(tag === $(UInt8(i))), :(func(unsafe_getproperty(x, Val($(QuoteNode(f)))))))
115                  for (i, f) in enumerate(namify.(fields))],
116                 # It is _super_ perf-sensitive to outline the assertion-failure here.
117                 # From 35μs to 5μs for peel benchmark. In any case, this is expected to
118                 # be unreachable for a well-formed union instance.
119                 :(tag_not_found(T, tag::UInt8)::Bool)
120             ))
121         end
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [w]arn, [h]ide type-stable statements, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native, [j]ump to source always.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
 • %9 = < concrete eval > Val(::Core.Const(:i8))::Core.Const(Val{:i8}())
   %10 = unsafe_getproperty(::RelNumber,::Val{:i8})::Int8
   %11 = #10(::Int8)::Int64
   %19 = < concrete eval > Val(::Core.Const(:i16))::Core.Const(Val{:i16}())
   %20 = unsafe_getproperty(::RelNumber,::Val{:i16})::Int16
   %21 = #10(::Int16)::Int64
   %29 = < concrete eval > Val(::Core.Const(:i32))::Core.Const(Val{:i32}())
   %30 = unsafe_getproperty(::RelNumber,::Val{:i32})::Int32
   %31 = #10(::Int32)::Int64
v  %39 = < concrete eval > Val(::Core.Const(:i64))::Core.Const(Val{:i64}())

(@v1.11) pkg> st
Status `~/.julia/environments/v1.11/Project.toml`
  [6e4b80f9] BenchmarkTools v1.6.0
⌃ [f68482b8] Cthulhu v2.16.1
  [31a5f54b] Debugger v0.7.10
  [1914dd2f] MacroTools v0.5.15
⌃ [e4faabce] PProf v3.1.3
  [817f1d60] ReTestItems v1.29.0
  [295af30f] Revise v3.7.2
Info Packages marked with ⌃ have new versions available and may be upgradable.

MRE file for the above:

module Onions

export @union

isonion(::Type) = false

using MacroTools
using MacroTools: @capture, @q

function fieldtype(ex)
    @capture(ex, _::T_ | _)
    return something(T, Any)
end

params(ex) = isexpr(ex, :curly) ? ex.args[2:end] : []

function freetypevars(ex)
    if isexpr(ex, Symbol)
        [ex]
    elseif !isexpr(ex)
        []
    elseif isexpr(ex, :(.))
        freetypevars(ex.args[1])
    elseif isexpr(ex, :curly, :tuple)
        reduce(vcat, freetypevars.(ex.args))
    else
        error("unrecognised type expression $ex")
    end
end

function ifelse(clauses, default=nothing)
    return foldr(((cond, body), els) -> Expr(:if, cond, body, els), clauses; init=default)
end

@noinline throw_not_set(::Type{T}, f) where {T} =
    error("Field $f is not active for type $(T)")
@noinline tag_not_found(::Type{T}, tag) where {T} =
    error("Malformed instance of union type $(T): Unexpected tag value: $(tag)")

function field end
function fields end
function fieldidx end
function unsafe_getproperty end
function _ueq end

# Peel the field with the given type. Implemented by the user.
# E.g. peel_as(Int, x) gets the Int field of x
# TODO (azreika): bit of a hack to make migration smoother
function peel_as end

macro union(ex)
    @capture(ex, struct T_ <: ParentType_
        fields__
    end | struct T_
        fields__
    end) || error("@union struct ...")
    types = Dict(namify(f) => fieldtype(f) for f in fields)
    fields = namify.(fields)
    if isnothing(ParentType)
        ParentType = Any
    end
    @assert length(fields) >= 1 || error("Union must have at least one field")
    forbidden = [namify(T), namify.(params(T))...]
    unboxed = filter(x -> isdisjoint(freetypevars(types[x]), forbidden), fields)
    type = @q struct $(esc(T)) <: $(esc(ParentType))
        tag::UInt8
        bits::Storage
        ptrs::P($(esc(namify(T))), $(esc.(params(T))...))
        $([
            @q function $(esc(T))(::Val{$(QuoteNode(f))}, value) where {$(esc.(params(T))...)}
                f_val = convert($(esc(types[f])), value)
                if $(f in unboxed) && isbitstype($(esc(types[f])))
                    new($(UInt8(i)), f_val, nothing)
                else
                    new($(UInt8(i)), nothing, f_val)
                end
            end
            for (i, f) in enumerate(fields)
        ]...)
    end
    :(let
        types = Dict($([:($(QuoteNode(f)) => $(esc(types[f]))) for f in unboxed]...))
        unboxed = [f for (f, T) in types if isbitstype(T)]
        Storage = Union{Nothing, [types[f] for f in unboxed]...}
        P($(esc(namify(T))), $(esc.(params(T))...)) = Union{Nothing,$([:($(QuoteNode(f)) in unboxed ? Union{} : $(esc(types[f]))) for f in fields]...)}
        $type
        $(esc(T))(; kw...) where {$(esc.(params(T))...)} = $(esc(T))(Val(only(kw)[1]), only(kw)[2])
        Onions.isonion(::Type{<:$(esc(namify(T)))}) = true
        Onions.fields(x::$(esc(T))) where {$(esc.(params(T))...)} = ($(QuoteNode.(fields)...),)
        Onions.fieldidx(::Type{<:$(esc(T))}, f::Symbol) where {$(esc.(params(T))...)} =
            # NOTE: This lookup in the NamedTuple compiles away when f is const-propped in
            # getproperty, and it's a constant-time lookup if it's dynamic.
            getfield(($([:($f = $(UInt8(i))) for (i, f) in enumerate(namify.(fields))]...),), f)
        $([
            :(function Onions.unsafe_getproperty(x::$(esc(T)), ::Val{$(QuoteNode(f))}) where {$(esc.(params(T))...)}
                $(if f in unboxed
                    :(if isbitstype($(esc(types[f])))
                        return getfield(x, :bits)::$(esc(types[f]))
                    end)
                end)
                getfield(x, :ptrs)::$(esc(types[f]))
            end)
            for f in fields
        ]...)
        # NOTE: Even though the following two functions share most of their implementation,
        # for some reason julia doesn't infer them well if you try to factor out a shared
        # impl. So we duplicate the code here a touch.
        function Onions.peel(
            func::F,
            x::T
        ) where {F, T<:$(esc(namify(T)))}
            tag = getfield(x, :tag)
            return $(ifelse(
                [(:(tag === $(UInt8(i))), :(func(unsafe_getproperty(x, Val($(QuoteNode(f)))))))
                 for (i, f) in enumerate(namify.(fields))],
                # It is _super_ perf-sensitive to outline the assertion-failure here.
                # From 35μs to 5μs for peel benchmark. In any case, this is expected to
                # be unreachable for a well-formed union instance.
                :(tag_not_found(T, tag))
            ))
        end
        function Onions._ueq(a::T, b::T) where {T<:$(esc(namify(T)))}
            getfield(a, :tag) == getfield(b, :tag) || return false
            tag = getfield(a, :tag)
            return $(ifelse(
                [(:(tag === $(UInt8(i))), :(
                        f = Val($(QuoteNode(f)));
                        unsafe_getproperty(a, f) == unsafe_getproperty(b, f)
                    ))
                 for (i, f) in enumerate(namify.(fields))],
                # It is _super_ perf-sensitive to outline the assertion-failure here.
                # From 35μs to 5μs for peel benchmark. In any case, this is expected to
                # be unreachable for a well-formed union instance.
                :(tag_not_found(T, tag))
            ))
        end

        function Base.getproperty(x::$(esc(T)), f::Symbol) where {$(esc.(params(T))...)}
            # Comparing ints is cheaper than looking up the symbol for the tag
            fieldidx(typeof(x), f) == getfield(x, :tag) ||
                Onions.throw_not_set(typeof(x), f)
            return unsafe_getproperty(x, Val(f))
        end
        Base.:(==)(a::$(esc(namify(T))), b::$(esc(namify(T)))) = ueq(a, b)
        function Base.show_default(io::IO, x::$(esc(T))) where {$(esc.(params(T))...)}
            show(io, typeof(x))
            print(io, "($(field(x)) = ")
            show(io, peel(x))
            print(io, ")")
            return
        end
        nothing
    end)
end

function aligned_type(align)
    if align == 1
        UInt8
    elseif align == 2
        UInt16
    elseif align == 4
        UInt32
    elseif align == 8
        UInt64
    elseif align == 16
        UInt128
    else
        error("Unsupported alignment: $align")
    end
end

peel(x) = peel(identity, x)

# User-overridable endpoint for implementing custom `==` for `@union` types.
ueq(a::T, b::T) where T = _ueq(a, b)

field(x) = fields(x)[getfield(x, :tag)]

function isfield(x, f)
    @assert f in fields(x)
    return field(x) === f
end


end  # module Onions

using .Onions
using .Onions: peel, field, isfield

@union struct RelNumber
    i8::Int8
    i16::Int16
    i32::Int32
    i64::Int64

    u8::UInt8
    u16::UInt16
    u32::UInt32
    u64::UInt64

    f16::Float16
    f32::Float32
    f64::Float64
end
RelNumber(x::Int8) = RelNumber(i8=x)
RelNumber(x::Int16) = RelNumber(i16=x)
RelNumber(x::Int32) = RelNumber(i32=x)
RelNumber(x::Int64) = RelNumber(i64=x)

RelNumber(x::UInt8) = RelNumber(u8=x)
RelNumber(x::UInt16) = RelNumber(u16=x)
RelNumber(x::UInt32) = RelNumber(u32=x)
RelNumber(x::UInt64) = RelNumber(u64=x)

RelNumber(x::Float16) = RelNumber(f16=x)
RelNumber(x::Float32) = RelNumber(f32=x)
RelNumber(x::Float64) = RelNumber(f64=x)

Base.:(<)(a::RelNumber, b::RelNumber) = cmp(a, b) == -1
Base.isless(a::RelNumber, b::RelNumber) = <(a,b)

function Base.cmp(a::RelNumber, b::RelNumber)
    return peel(b) do bv; cmp_inner(a, bv) ; end
end
@noinline cmp_inner(a::RelNumber, bv) = peel(a) do av; cmp(av, bv) ; end