#####
##### Compilation context
#####

# User-extended context allows parametrization of the pipeline through
# our subtype of AbstractInterpreter
abstract type CompilationContext end

struct NoContext <: CompilationContext end

@doc(
"""
    abstract type CompilationContext end

Parametrize the Mixtape pipeline by inheriting from `CompilationContext`. Similar to the context objects in [Cassette.jl](https://julia.mit.edu/Cassette.jl/stable/contextualdispatch.html). By using the interface methods [`transform`](@ref) and [`optimize!`](@ref) -- the user can control different parts of the compilation pipeline.
""", CompilationContext)

transform(ctx::CompilationContext, b) = b
transform(ctx::CompilationContext, b, sig) = transform(ctx, b)

@doc(
"""
    transform(ctx::CompilationContext, b::Core.CodeInfo)::Core.CodeInfo
    transform(ctx::CompilationContext, b::Core.CodeInfo, sig::Tuple)::Core.CodeInfo

User-defined transform which operates on lowered `Core.CodeInfo`. There's two versions: (1) ignores the signature of the current method body under consideration and (2) provides the signature as `sig`.

Transforms might typically follow a simple "swap" format using `CodeInfoTools.Builder`:

```julia
function transform(::MyCtx, src)
    b = CodeInfoTools.Builder(b)
    for (k, st) in b
        b[k] = swap(st))
    end
    return CodeInfoTools.finish(b)
end
```

but more advanced formats are possible. For further utilities, please see [CodeInfoTools.jl](https://github.com/JuliaCompilerPlugins/CodeInfoTools.jl).
""", transform)

optimize!(ctx::CompilationContext, b) = julia_passes!(b)

@doc(
"""
    optimize!(ctx::CompilationContext, b::OptimizationBundle)::Core.Compiler.IRCode

User-defined transform which operates on inferred IR provided by an [`OptimizationBundle`](@ref) instance.

The fallback implementation is:
```julia
optimize!(ctx::CompilationContext, b::OptimizationBundle) = julia_passes!(b)
```
which runs a set of standard (and required) Julia passes to the lowered and inferred `Core.Compiler.IRCode`.

!!! warning
    If you overload this method, _you are responsible for the optimization pass_! This means that you should know what you're doing (in general), and you will likely also want to call `julia_passes!` yourself. Be aware of this -- or else you'll receive verification errors on `Core.Compiler.IRCode`.
""", optimize!)

allow(f::C, args...) where {C <: CompilationContext} = false
function allow(ctx::CompilationContext, mod::Module, fn, args...)
    return allow(ctx, mod) || allow(ctx, fn, args...)
end

@doc(
"""
    allow(f::CompilationContext, args...)::Bool

Determines whether the user-defined [`transform`](@ref) and [`optimize!`](@ref) are allowed to look at a lowered `Core.CodeInfo` or `Core.Compiler.IRCode` instance.

The user is allowed to greenlight modules:

```julia
allow(::MyCtx, m::Module) == m == SomeModule
```

or even specific signatures

```julia
allow(::MyCtx, fn::typeof(rand), args...) = true
```
""", allow)
