diff --git a/src/DispatchDoctor.jl b/src/DispatchDoctor.jl index fc02912..f429420 100644 --- a/src/DispatchDoctor.jl +++ b/src/DispatchDoctor.jl @@ -1,6 +1,6 @@ module DispatchDoctor -export @stable, @unstable, allow_unstable, TypeInstabilityError, register_macro! +export @stable, @unstable, @register_macro, allow_unstable, TypeInstabilityError, register_macro! include("utils.jl") include("errors.jl") @@ -20,7 +20,7 @@ using ._Printing using ._Interactions: MACRO_BEHAVIOR, MacroInteractions, CompatibleMacro, IncompatibleMacro, DontPropagateMacro, register_macro!, get_macro_behavior, ignore_function using ._RuntimeChecks: INSTABILITY_CHECK_ENABLED, allow_unstable, is_precompiling using ._Stabilization: _stable, _stabilize_all, _stabilize_fnc, _stabilize_module -using ._Macros: @stable, @unstable +using ._Macros: @stable, @unstable, @register_macro #! format: on end diff --git a/src/macro_interactions.jl b/src/macro_interactions.jl index df42f65..8752efa 100644 --- a/src/macro_interactions.jl +++ b/src/macro_interactions.jl @@ -17,50 +17,58 @@ end # Macros we dont want to propagate const MACRO_BEHAVIOR = (; table=Dict([ - Symbol("@stable") => IncompatibleMacro, # + (Main => Symbol("@stable")) => IncompatibleMacro, # # ^ We don't want to stabilize a function twice. - Symbol("@unstable") => IncompatibleMacro, # + (Main => Symbol("@unstable")) => IncompatibleMacro, # # ^ This is the purpose of `@unstable` - Symbol("@doc") => DontPropagateMacro, # Core + (Main => Symbol("@doc")) => DontPropagateMacro, # Core # ^ Base.@__doc__ takes care of this. - Symbol("@assume_effects") => IncompatibleMacro, # Base + (Main => Symbol("@assume_effects")) => IncompatibleMacro, # Base # ^ Some effects are incompatible, like # :nothrow, so this requires much more # work to get working. TODO. - Symbol("@enum") => IncompatibleMacro, # Base + (Main => Symbol("@enum")) => IncompatibleMacro, # Base # ^ TODO. Seems to interact. - Symbol("@eval") => IncompatibleMacro, # Base + (Main => Symbol("@eval")) => IncompatibleMacro, # Base # ^ Too much flexibility to apply, # and user could always use `@eval` # inside function. - Symbol("@deprecate") => IncompatibleMacro, # Base + (Main => Symbol("@deprecate")) => IncompatibleMacro, # Base # ^ TODO. Seems to interact. - Symbol("@generated") => IncompatibleMacro, # Base + (Main => Symbol("@generated")) => IncompatibleMacro, # Base # ^ In principle this is compatible but # needs additional logic to work. - Symbol("@kwdef") => IncompatibleMacro, # Base + (Main => Symbol("@kwdef")) => IncompatibleMacro, # Base # ^ TODO. Seems to interact. - Symbol("@pure") => IncompatibleMacro, # Base + (Main => Symbol("@pure")) => IncompatibleMacro, # Base # ^ See `@assume_effects`. - Symbol("@everywhere") => DontPropagateMacro, # Distributed + (Main => Symbol("@everywhere")) => DontPropagateMacro, # Distributed # ^ Prefer to have block passed to workers # only a single time. And `@everywhere` # works with blocks of code, so it is # fine. - Symbol("@model") => IncompatibleMacro, # Turing + (Main => Symbol("@model")) => IncompatibleMacro, # Turing # ^ Fairly common macro used to define # probabilistic models. The syntax is # incompatible with `@stable`. - Symbol("@capture") => IncompatibleMacro, # MacroTools + (Main => Symbol("@capture")) => IncompatibleMacro, # MacroTools # ^ Similar to `@model`. ]), lock=Threads.SpinLock(), ) -get_macro_behavior(_) = CompatibleMacro -get_macro_behavior(ex::Symbol) = get(MACRO_BEHAVIOR.table, ex, CompatibleMacro) -get_macro_behavior(ex::QuoteNode) = get_macro_behavior(ex.value) -function get_macro_behavior(ex::Expr) - parts = map(get_macro_behavior, ex.args) +get_macro_behavior(_, _) = CompatibleMacro +function get_macro_behavior(m::Module, ex::Symbol) + while m != Main + if haskey(MACRO_BEHAVIOR.table, m => ex) break + else m = parentmodule(m) + end + end + + return get(MACRO_BEHAVIOR.table, m => ex, CompatibleMacro) +end +get_macro_behavior(m::Module, ex::QuoteNode) = get_macro_behavior(m, ex.value) +function get_macro_behavior(m::Module, ex::Expr) + parts = map(arg -> get_macro_behavior(m, arg), ex.args) return reduce(combine_behavior, parts; init=CompatibleMacro) end @@ -74,6 +82,18 @@ function combine_behavior(a::MacroInteractions, b::MacroInteractions) end end +function _register_macro!(m::Module, macro_name::Symbol, behavior::MacroInteractions) + lock(MACRO_BEHAVIOR.lock) do + if haskey(MACRO_BEHAVIOR.table, m => macro_name) + error( + "Macro `$macro_name` already registered in module $m with behavior ($(MACRO_BEHAVIOR.table[m => macro_name]).", + ) + end + MACRO_BEHAVIOR.table[m => macro_name] = behavior + MACRO_BEHAVIOR.table[m => macro_name] + end +end + """ register_macro!(macro_name::Symbol, behavior::MacroInteractions) @@ -97,17 +117,7 @@ using DispatchDoctor: register_macro!, IncompatibleMacro register_macro!(Symbol("@mymacro"), IncompatibleMacro) ``` """ -function register_macro!(macro_name::Symbol, behavior::MacroInteractions) - lock(MACRO_BEHAVIOR.lock) do - if haskey(MACRO_BEHAVIOR.table, macro_name) - error( - "Macro `$macro_name` already registered with behavior $(MACRO_BEHAVIOR.table[macro_name]).", - ) - end - MACRO_BEHAVIOR.table[macro_name] = behavior - MACRO_BEHAVIOR.table[macro_name] - end -end +register_macro!(macro_name::Symbol, behavior::MacroInteractions) = _register_macro!(Main, macro_name, behavior) """ ignore_function(f) diff --git a/src/macros.jl b/src/macros.jl index 5a9bb45..5e7015b 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -3,6 +3,7 @@ module _Macros using .._Utils: JULIA_OK using .._Stabilization: _stable +using .._Interactions: CompatibleMacro, DontPropagateMacro, IncompatibleMacro, _register_macro! """ @stable [options...] [code_block] @@ -96,4 +97,39 @@ macro unstable(fex) return esc(fex) end +""" + @register_macro(behavior, macro_name) + +Register a macro with a specified behavior in the `MACRO_BEHAVIOR` list. + +This function adds a new macro and its associated behavior to the global list that +tracks how macros should be treated when encountered during the stabilization +process. The behavior can be one of `CompatibleMacro`, `IncompatibleMacro`, or `DontPropagateMacro`, +which influences how the `@stable` macro interacts with the registered macro. + +The default behavior for `@stable` is to assume `CompatibleMacro` unless explicitly declared. + +# Arguments +- `macro_name::Symbol`: The symbol representing the macro to register. +- `behavior::MacroInteractions`: The behavior to associate with the macro, which dictates how it should be handled. + +# Examples +```julia +using DispatchDoctor: @register_macro, IncompatibleMacro + +@register_macro IncompatibleMacro @mymacro +``` +""" +macro register_macro(behavior_name, macro_call) + behavior = + if behavior_name == :CompatibleMacro CompatibleMacro + elseif behavior_name == :DontPropagateMacro DontPropagateMacro + elseif behavior_name == :IncompatibleMacro IncompatibleMacro + else error("$behavior_name is not a valid macro interaction") + end + macro_name = macro_call.args[1] + + _register_macro!(__module__, macro_name, behavior) +end + end diff --git a/src/stabilization.jl b/src/stabilization.jl index 6c0917a..ddd2efe 100644 --- a/src/stabilization.jl +++ b/src/stabilization.jl @@ -30,7 +30,8 @@ function _stable(args...; calling_module, source_info, kws...) if options.mode in ("error", "warn") out, metadata = _stabilize_all( ex, - DownwardMetadata(); + DownwardMetadata(), + calling_module; source_info, kws..., options.mode, @@ -79,13 +80,13 @@ function UpwardMetadata(downward_metadata::DownwardMetadata; matching_function:: ) end -function _stabilize_all(ex, downward_metadata::DownwardMetadata; kws...) +function _stabilize_all(ex, downward_metadata::DownwardMetadata, calling_module; kws...) return ex, UpwardMetadata(downward_metadata) end -function _stabilize_all(ex::Expr, downward_metadata::DownwardMetadata; kws...) +function _stabilize_all(ex::Expr, downward_metadata::DownwardMetadata, calling_module; kws...) #! format: off if ex.head == :macrocall - macro_behavior = get_macro_behavior(ex.args[1]) + macro_behavior = get_macro_behavior(calling_module, ex.args[1]) if macro_behavior == IncompatibleMacro return ex, UpwardMetadata(downward_metadata) elseif macro_behavior == CompatibleMacro @@ -98,7 +99,7 @@ function _stabilize_all(ex::Expr, downward_metadata::DownwardMetadata; kws...) push!(macro_keys, my_key) new_downward_metadata = DownwardMetadata(; macros_to_use, macro_keys) - inner_ex, upward_metadata = _stabilize_all(ex.args[end], new_downward_metadata; kws...) + inner_ex, upward_metadata = _stabilize_all(ex.args[end], new_downward_metadata, calling_module; kws...) if isempty(upward_metadata.unused_macros) # It has been applied! So we just return the inner part @@ -120,7 +121,7 @@ function _stabilize_all(ex::Expr, downward_metadata::DownwardMetadata; kws...) @assert macro_behavior == DontPropagateMacro # Apply to last argument only - inner_ex, upward_metadata = _stabilize_all(ex.args[end], downward_metadata; kws...) + inner_ex, upward_metadata = _stabilize_all(ex.args[end], downward_metadata, calling_module; kws...) new_ex = Expr(:macrocall, ex.args[1:end-1]..., inner_ex) return new_ex, upward_metadata end @@ -134,7 +135,7 @@ function _stabilize_all(ex::Expr, downward_metadata::DownwardMetadata; kws...) # Incompatible with two functions return ex, UpwardMetadata(downward_metadata) elseif ex.head == :module - return _stabilize_module(ex, downward_metadata; kws...) + return _stabilize_module(ex, downward_metadata, calling_module; kws...) elseif ex.head == :call && ex.args[1] == Symbol("include") && length(ex.args) == 2 # We can't track the matches in includes, so just assume # there are some matches. TODO: However, this is not a great solution. @@ -146,7 +147,7 @@ function _stabilize_all(ex::Expr, downward_metadata::DownwardMetadata; kws...) # TODO: Should report `isdef` to MacroTools as not capturing all cases return _stabilize_fnc(ex, downward_metadata; kws...) else - stabilized_args = map(e -> _stabilize_all(e, DownwardMetadata(); kws...), ex.args) + stabilized_args = map(e -> _stabilize_all(e, DownwardMetadata(), calling_module; kws...), ex.args) merged_upward_metadata = reduce(merge, map(last, stabilized_args); init=UpwardMetadata()) new_ex = Expr(ex.head, map(first, stabilized_args)...) return new_ex, UpwardMetadata(downward_metadata; matching_function=merged_upward_metadata.matching_function) @@ -157,7 +158,7 @@ end function _stabilizing_include(m::Module, path; kws...) inner = let kws = kws (ex,) -> let - new_ex, upward_metadata = _stabilize_all(ex, DownwardMetadata(); kws...) + new_ex, upward_metadata = _stabilize_all(ex, DownwardMetadata(), m; kws...) @assert isempty(upward_metadata.unused_macros) new_ex end @@ -165,9 +166,9 @@ function _stabilizing_include(m::Module, path; kws...) return m.include(inner, path) end -function _stabilize_module(ex, downward_metadata; kws...) +function _stabilize_module(ex, downward_metadata, calling_module; kws...) stabilized_args = map( - e -> _stabilize_all(e, DownwardMetadata(); kws...), ex.args[3].args + e -> _stabilize_all(e, DownwardMetadata(), calling_module; kws...), ex.args[3].args ) merged_upward_metadata = reduce( merge, map(last, stabilized_args); init=UpwardMetadata() diff --git a/test/unittests.jl b/test/unittests.jl index 9d0df62..4f980d8 100644 --- a/test/unittests.jl +++ b/test/unittests.jl @@ -508,7 +508,7 @@ end g(x, y) = x > 0 ? y : 0.0 end), DispatchDoctor._Stabilization.DownwardMetadata(), - ), + Main), ) # First, we capture `f` using postwalk and `@capture` @@ -564,7 +564,7 @@ end using DispatchDoctor: _stabilize_fnc, _stabilize_all ex = _stabilize_all( - :(function donothing end), DispatchDoctor._Stabilization.DownwardMetadata() + :(function donothing end), DispatchDoctor._Stabilization.DownwardMetadata(), Main )[1] @test ex == :(function donothing end) @@ -591,6 +591,7 @@ end return ex end), DispatchDoctor._Stabilization.DownwardMetadata(), + Main ) # Should skip the internal function @@ -831,6 +832,7 @@ end return x > 0 ? x : 0.0 end), DispatchDoctor._Stabilization.DownwardMetadata(), + Main ) JULIA_OK && @test occursin("propagate_inbounds", string(ex)) end @@ -840,16 +842,14 @@ end macro mymacro(ex) return esc(ex) end - if !haskey(DispatchDoctor.MACRO_BEHAVIOR.table, Symbol("@mymacro")) - register_macro!(Symbol("@mymacro"), DispatchDoctor.IncompatibleMacro) - end - @test DispatchDoctor.get_macro_behavior(:(@mymacro x = 1)) == + @register_macro IncompatibleMacro @mymacro + @test DispatchDoctor.get_macro_behavior(@__MODULE__, :(@mymacro x = 1)) == DispatchDoctor.IncompatibleMacro # Trying to register again should fail with a useful error if VERSION >= v"1.9" - @test_throws "Macro `@mymacro` already registered" register_macro!( - Symbol("@mymacro"), DispatchDoctor.IncompatibleMacro + @test_throws "Macro `@mymacro` already registered" eval( + :(@register_macro IncompatibleMacro @mymacro) ) end @@ -873,18 +873,12 @@ end macro dontpropagatemacro(ex) return esc(ex) end - if !haskey(DDI.MACRO_BEHAVIOR.table, Symbol("@compatiblemacro")) - register_macro!(Symbol("@compatiblemacro"), DDI.CompatibleMacro) - end - if !haskey(DDI.MACRO_BEHAVIOR.table, Symbol("@incompatiblemacro")) - register_macro!(Symbol("@incompatiblemacro"), DDI.IncompatibleMacro) - end - if !haskey(DDI.MACRO_BEHAVIOR.table, Symbol("@dontpropagatemacro")) - register_macro!(Symbol("@dontpropagatemacro"), DDI.DontPropagateMacro) - end - @test DDI.get_macro_behavior(:(@compatiblemacro true x = 1)) == DDI.CompatibleMacro - @test DDI.get_macro_behavior(:(@incompatiblemacro x = 1)) == DDI.IncompatibleMacro - @test DDI.get_macro_behavior(:(@dontpropagatemacro x = 1)) == DDI.DontPropagateMacro + @register_macro CompatibleMacro @compatiblemacro + @register_macro IncompatibleMacro @incompatiblemacro + @register_macro DontPropagateMacro @dontpropagatemacro + @test DDI.get_macro_behavior(@__MODULE__, :(@compatiblemacro true x = 1)) == DDI.CompatibleMacro + @test DDI.get_macro_behavior(@__MODULE__, :(@incompatiblemacro x = 1)) == DDI.IncompatibleMacro + @test DDI.get_macro_behavior(@__MODULE__, :(@dontpropagatemacro x = 1)) == DDI.DontPropagateMacro @test DDI.combine_behavior(DDI.CompatibleMacro, DDI.CompatibleMacro) == DDI.CompatibleMacro @@ -920,7 +914,7 @@ end end if DispatchDoctor.JULIA_OK new_ex, upward_metadata = DispatchDoctor._stabilize_all( - ex, DispatchDoctor._Stabilization.DownwardMetadata() + ex, DispatchDoctor._Stabilization.DownwardMetadata(), @__MODULE__ ) # We should expect: # 1. All of the `@dontpropagatemacro`'s to be on the outside of the block. @@ -965,6 +959,7 @@ end return x > 0 ? x : 0.0 end), downward_metadata, + Main ) @test upward_metadata.matching_function @test isempty(upward_metadata.unused_macros) @@ -1236,5 +1231,30 @@ end # julia process with things like --code-coverage disabled. # See https://discourse.julialang.org/t/improving-speed-of-runtime-dispatch-detector/114697/14?u=milescranmer end +@testitem "Macros with same name" begin + using DispatchDoctor: _Interactions as DDI + @register_macro IncompatibleMacro @new_macro + + module AModule + using DispatchDoctor + @register_macro CompatibleMacro @new_macro + module BModule end + end + + module CModule + module DModule + using DispatchDoctor + @register_macro DontPropagateMacro @new_macro + end + end + + @test DDI.get_macro_behavior(@__MODULE__, Symbol("@new_macro")) == DDI.IncompatibleMacro + @test DDI.get_macro_behavior(AModule, Symbol("@new_macro")) == DDI.CompatibleMacro + @test DDI.get_macro_behavior(AModule.BModule, Symbol("@new_macro")) == DDI.CompatibleMacro + @test DDI.get_macro_behavior(CModule, Symbol("@new_macro")) == DDI.IncompatibleMacro + @test DDI.get_macro_behavior(CModule.DModule, Symbol("@new_macro")) == DDI.DontPropagateMacro + + @test_throws LoadError eval(:(@register_macro CompatibleMacro @new_macro)) +end @run_package_tests