MasonProtter · December 9, 2025 18:19 · Dec 9, 2025
diff --git a/contextual.jl b/contextual.jl
@@ -0,0 +1,177 @@
+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
No results found