Compiler/extras/CompilerDevTools/src/CompilerDevTools.jl - master - githubmirror/julia

Compiler/extras/CompilerDevTools/src/CompilerDevTools.jl
81 строка4 KB
Gabriel Baraldi
Use dictionary instead of linear search for the local cache. (#60795)
07 фев 2026, 18:08
07 фев 2026, 18:08a508609
module CompilerDevTools

using Compiler
using Compiler: argextype, widenconst
using Core.IR
using Base: isexpr

mutable struct SplitCacheOwner end

struct SplitCacheInterp <: Compiler.AbstractInterpreter
    world::UInt
    owner::SplitCacheOwner
    inf_params::Compiler.InferenceParams
    opt_params::Compiler.OptimizationParams
    inf_cache::Compiler.InferenceCache
    codegen_cache::IdDict{CodeInstance,CodeInfo}
    function SplitCacheInterp(;
        world::UInt = Base.get_world_counter(),
        owner::SplitCacheOwner = SplitCacheOwner(),
        inf_params::Compiler.InferenceParams = Compiler.InferenceParams(),
        opt_params::Compiler.OptimizationParams = Compiler.OptimizationParams(),
        inf_cache::Compiler.InferenceCache = Compiler.InferenceCache())
        new(world, owner, inf_params, opt_params, inf_cache, IdDict{CodeInstance,CodeInfo}())
    end
end

Compiler.InferenceParams(interp::SplitCacheInterp) = interp.inf_params
Compiler.OptimizationParams(interp::SplitCacheInterp) = interp.opt_params
Compiler.get_inference_world(interp::SplitCacheInterp) = interp.world
Compiler.get_inference_cache(interp::SplitCacheInterp) = interp.inf_cache
Compiler.cache_owner(interp::SplitCacheInterp) = interp.owner
Compiler.codegen_cache(interp::SplitCacheInterp) = interp.codegen_cache

import Core.OptimizedGenerics.CompilerPlugins: typeinf, typeinf_edge
@eval @noinline typeinf(owner::SplitCacheOwner, mi::MethodInstance, source_mode::UInt8) =
    Base.invoke_in_world(which(typeinf, Tuple{SplitCacheOwner, MethodInstance, UInt8}).primary_world, Compiler.typeinf_ext_toplevel, SplitCacheInterp(; world=Base.tls_world_age(), owner), mi, source_mode)

@eval @noinline function typeinf_edge(owner::SplitCacheOwner, mi::MethodInstance, parent_frame::Compiler.InferenceState, world::UInt, source_mode::UInt8)
    # TODO: This isn't quite right, we're just sketching things for now
    interp = SplitCacheInterp(; world, owner)
    Compiler.typeinf_edge(interp, mi.def, mi.specTypes, Core.svec(), parent_frame, false, false)
end

function lookup_method_instance(f, args...)
    @ccall jl_method_lookup(Any[f, args...]::Ptr{Any}, (1+length(args))::Csize_t, Base.tls_world_age()::Csize_t)::Ref{Core.MethodInstance}
end

function Compiler.transform_result_for_cache(interp::SplitCacheInterp, result::Compiler.InferenceResult, edges::Compiler.SimpleVector)
    opt = result.src::Compiler.OptimizationState
    ir = opt.optresult.ir::Compiler.IRCode
    override = with_new_compiler
    for inst in ir.stmts
        stmt = inst[:stmt]
        isexpr(stmt, :call) || continue
        f = stmt.args[1]
        f === override && continue
        T = widenconst(argextype(f, ir))
        T <: Core.Builtin && continue
        insert!(stmt.args, 1, override)
        insert!(stmt.args, 3, interp.owner)
    end
    @invoke Compiler.transform_result_for_cache(interp::Compiler.AbstractInterpreter, result::Compiler.InferenceResult, edges::Compiler.SimpleVector)
end

with_new_compiler(f, args...; owner::SplitCacheOwner = SplitCacheOwner()) = with_new_compiler(f, owner, args...)

function with_new_compiler(f, owner::SplitCacheOwner, args...)
    # We try to avoid introducing `with_new_compiler` in the first place,
    # but if we can't see the type, it's still possible to end up with a
    # builtin here - simply forward to the ordinary builtin call.
    isa(f, Core.Builtin) && return f(args...)
    mi = lookup_method_instance(f, args...)
    new_compiler_ci = Core.OptimizedGenerics.CompilerPlugins.typeinf(
        owner, mi, Compiler.SOURCE_MODE_ABI
    )
    invoke(f, new_compiler_ci, args...)
end

export with_new_compiler

end