Created
December 9, 2025 18:19
-
-
Save MasonProtter/835175d189c96a21de0f6b51afde4e02 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| module Contextual | |
| using Core: MethodMatch, MethodInstance, CodeInstance, Compiler | |
| using Core.Compiler: | |
| _methods_by_ftype, | |
| InferenceParams, | |
| get_world_counter, | |
| tls_world_age, | |
| InferenceResult, | |
| typeinf, | |
| InternalMethodTable, | |
| InferenceState, | |
| NativeInterpreter, | |
| AbstractInterpreter, | |
| OptimizationParams, | |
| MethodTableView, | |
| OverlayMethodTable, | |
| WorldRange, | |
| finish, | |
| InferenceResult, | |
| OptimizationState, | |
| OptimizationParams, | |
| optimize, | |
| store_backedges, | |
| finish!, | |
| CodeInfo, | |
| isexpr, | |
| argextype, | |
| singleton_type, | |
| Builtin, | |
| GlobalRef, | |
| CodeInstance, | |
| ir_to_codeinf!, | |
| typeinf_ext_toplevel, | |
| IRCode, | |
| argextype, | |
| widenconst | |
| const CC = Core.Compiler | |
| using Base: | |
| specialize_method, | |
| isexpr | |
| # include("code_cache.jl") | |
| abstract type ContextOwner end | |
| struct ContextualInterpreter{Owner} <: AbstractInterpreter | |
| # Cache of inference results for this particular interpreter | |
| inf_cache::Vector{InferenceResult} | |
| # The world age we're working inside of | |
| world::UInt | |
| # Parameters for inference and optimization | |
| inf_params::InferenceParams | |
| opt_params::OptimizationParams | |
| codegen_cache::IdDict{CodeInstance, CodeInfo} | |
| function ContextualInterpreter{Owner}(world::UInt, | |
| ip::InferenceParams, | |
| op::OptimizationParams) where {Owner} | |
| @assert world <= Base.get_world_counter() | |
| return new{Owner}( | |
| # Initially empty cache | |
| InferenceResult[], | |
| # world age counter | |
| world, | |
| # parameters for inference and optimization | |
| ip, | |
| op, | |
| IdDict{CodeInstance, CodeInfo}() | |
| ) | |
| end | |
| end | |
| function ContextualInterpreter{Owner}(; | |
| world=Base.get_world_counter(), | |
| inf_params=InferenceParams(), | |
| opt_params=OptimizationParams()) where {Owner} | |
| ContextualInterpreter{Owner}(world, inf_params, opt_params) | |
| end | |
| Core.Compiler.InferenceParams(interp::ContextualInterpreter) = interp.inf_params | |
| Core.Compiler.OptimizationParams(interp::ContextualInterpreter) = interp.opt_params | |
| Core.Compiler.get_inference_world(interp::ContextualInterpreter) = interp.world | |
| Core.Compiler.get_inference_cache(interp::ContextualInterpreter) = interp.inf_cache | |
| Core.Compiler.cache_owner(interp::ContextualInterpreter{Owner}) where {Owner} = Owner.instance | |
| Core.Compiler.codegen_cache(interp::ContextualInterpreter) = interp.codegen_cache | |
| function generated_ci_in_absint_body(world::UInt, lnn, this, f, owner, args) | |
| sig = Type{Tuple{f, args.parameters...}} | |
| # Core.println(sig) | |
| sig isa Type{<:Type{<:Tuple}} || error() | |
| tt = sig.parameters[1] | |
| interp = ContextualInterpreter{owner.parameters[1]}(; world) | |
| match, valid_worlds = Core.Compiler.findsup(tt, Core.Compiler.method_table(interp)) | |
| if match === nothing | |
| error(lazy"Unable to find matching $tt") | |
| end | |
| mi = specialize_method(match.method, match.spec_types, match.sparams)::MethodInstance | |
| cinst = Core.Compiler.typeinf_ext_toplevel(interp, mi, Compiler.SOURCE_MODE_ABI) | |
| ci = expr_to_codeinfo(@__MODULE__(), [Symbol("#self#"), :f, :owner, :args], [], (), :(return $cinst)) | |
| matches = Base._methods_by_ftype(sig, -1, world) | |
| if !isnothing(matches) | |
| ci.edges = Core.MethodInstance[] | |
| for match in Base._methods_by_ftype(sig, -1, world) | |
| mi = Base.specialize_method(match) | |
| push!(ci.edges, mi) | |
| end | |
| end | |
| return ci | |
| end | |
| function expr_to_codeinfo(m::Module, argnames, spnames, sp, e::Expr) | |
| lam = Expr(:lambda, argnames, | |
| Expr(Symbol("scope-block"), | |
| Expr(:block, | |
| Expr(:return, | |
| Expr(:block, | |
| e, | |
| ))))) | |
| ex = if spnames === nothing || isempty(spnames) | |
| lam | |
| else | |
| Expr(Symbol("with-static-parameters"), lam, spnames...) | |
| end | |
| ci = Base.generated_body_to_codeinfo(ex, @__MODULE__(), false) | |
| @assert ci isa Core.CodeInfo "Failed to create a CodeInfo from the given expression. This might mean it contains a closure or comprehension?\n Offending expression: $e" | |
| ci | |
| end | |
| function refresh_generated_ci_in_absint() | |
| @eval function generated_ci_in_absint(f, owner, args) | |
| $(Expr(:meta, :generated_only)) | |
| $(Expr(:meta, :generated, generated_ci_in_absint_body)) | |
| end | |
| end | |
| refresh_generated_ci_in_absint() | |
| macro context(_s::Symbol) | |
| s = esc(_s) | |
| s_method_table = Symbol(_s, :_METHOD_TALBE) | |
| s_code_cache = esc(Symbol(_s, :_CODE_CACHE)) | |
| quote | |
| struct $s <: ContextOwner end | |
| Base.Experimental.@MethodTable($s_method_table) | |
| $Core.Compiler.method_table(::ContextualInterpreter{$s}) = $s_method_table | |
| function $s(f, args...) | |
| cinst = generated_ci_in_absint(f, $s, args) | |
| invoke(f, cinst, args...) | |
| end | |
| end | |
| end | |
| @noinline function Core.OptimizedGenerics.CompilerPlugins.typeinf(::Owner, mi::MethodInstance, source_mode::UInt8) where {Owner <: ContextOwner} | |
| # Base.invoke_in_world(which(Core.OptimizedGenerics.CompilerPlugins.typeinf, Tuple{ContextOwner, MethodInstance, UInt8}).primary_world, | |
| Compiler.typeinf_ext_toplevel(ContextualInterpreter{Owner}(; world=Base.tls_world_age()), | |
| mi, source_mode) | |
| end | |
| @noinline function Core.OptimizedGenerics.CompilerPlugins.typeinf_edge(::Owner, mi::MethodInstance, parent_frame::Compiler.InferenceState, world::UInt, source_mode::UInt8) where {Owner <: ContextOwner} | |
| # TODO: This isn't quite right, we're just sketching things for now | |
| error() | |
| interp = ContextualInterpreter{Owner}(; world) | |
| Compiler.typeinf_edge(interp, mi.def, mi.specTypes, Core.svec(), | |
| parent_frame, false, false) | |
| end | |
| macro contextual(ctx, fdef) | |
| ex = esc(:($Base.Experimental.@overlay $get_method_table($ctx) $fdef)) | |
| end | |
| export @context, @contextual | |
| end # module Contextual |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment