Scope macro interactions by module

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
 
 include("utils.jl")
 include("errors.jl")
@@ -17,10 +17,10 @@ using ._Utils: extract_symbol, JULIA_OK, Unknown, specializing_typeof, type_inst
 using ._Errors: TypeInstabilityError, TypeInstabilityWarning, AllowUnstableDataRace
 using ._Preferences
 using ._Printing
-using ._Interactions: MACRO_BEHAVIOR, MacroInteractions, CompatibleMacro, IncompatibleMacro, DontPropagateMacro, register_macro!, get_macro_behavior, ignore_function
+using ._Interactions: MACRO_BEHAVIOR, MacroInteractions, CompatibleMacro, IncompatibleMacro, DontPropagateMacro, 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

src/macro_interactions.jl

@@ -17,50 +17,53 @@ end
 # Macros we dont want to propagate
 const MACRO_BEHAVIOR = (;
     table=Dict([
-        Symbol("@stable") => IncompatibleMacro,             # <self>
+        (Main => Symbol("@stable")) => IncompatibleMacro,             # <self>
         # ^ We don't want to stabilize a function twice.
-        Symbol("@unstable") => IncompatibleMacro,           # <self>
+        (Main => Symbol("@unstable")) => IncompatibleMacro,           # <self>
         # ^ 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)
+    default = get(MACRO_BEHAVIOR.table, Main => ex, CompatibleMacro)
+    return get(MACRO_BEHAVIOR.table, m => ex, default)
+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,41 +77,6 @@ function combine_behavior(a::MacroInteractions, b::MacroInteractions)
     end
 end
 
-"""
-    register_macro!(macro_name::Symbol, behavior::MacroInteractions)
-
-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!(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
-
 """
     ignore_function(f)
 

src/macros.jl

@@ -3,6 +3,7 @@ module _Macros
 
 using .._Utils: JULIA_OK
 using .._Stabilization: _stable
+using .._Interactions: MACRO_BEHAVIOR, CompatibleMacro, DontPropagateMacro, IncompatibleMacro
 
 """
     @stable [options...] [code_block]
@@ -96,4 +97,46 @@ 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]
+    lock(MACRO_BEHAVIOR.lock) do
+        if haskey(MACRO_BEHAVIOR.table, __module__ => macro_name)
+            error(
+                "Macro `$macro_name` already registered in module $__module__ with behavior ($(MACRO_BEHAVIOR.table[__module__ => macro_name]).",
+            )
+        end
+        MACRO_BEHAVIOR.table[__module__ => macro_name] = behavior
+        MACRO_BEHAVIOR.table[__module__ => macro_name]
+    end
+end
+
 end

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,17 +158,17 @@ 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
     end
     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()

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,21 @@ 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
+
+    module AModule
+        using DispatchDoctor
+        @register_macro CompatibleMacro @stable
+    end
+
+    @test DDI.get_macro_behavior(Main, Symbol("@stable")) == DDI.IncompatibleMacro
+    @test DDI.get_macro_behavior(AModule, Symbol("@stable")) == DDI.CompatibleMacro
+
+    @test_throws LoadError eval(quote
+        @register_macro CompatibleMacro @a_macro
+        @register_macro IncompatibleMacro @a_macro
+    end)
+end
 
 @run_package_tests