Defining traits on functions using macros

Hello,

I am trying to use traits on functions and I would like to define the trait automatically based on existing methods.
I have several traits like EntityFeature and FrameFeature and a function featuretype(::typeof(myfun)) = EntityFeature() to get the trait and use it for dispatch.
The trait is assigned depending of existing methods of myfun.
Now what I would like to automatically define the featuretype function when I define myfun.
I have written a macro that works like this:

macro feature(fun)
    f = eval(fun)  # evaluate in the scope of this module
    T = fun.args[1].args[1] # hacky ? get function name
    if hasmethod(f, Tuple{Roadway, Entity})
        @eval featuretype(::typeof($T)) = EntityFeature()
    elseif hasmethod(f, Tuple{Roadway, Frame, Entity})
        @eval featuretype(::typeof($T)) = FrameFeature()
    elseif hasmethod(f, Tuple{Roadway, Vector{<:Frame}, Any})
        @eval featuretype(::typeof($T)) = TemporalFeature()
    else
        error(""""
              unsupported feature type
              try adding a new feature trait or implement one of the existing feature type supported
              """)
    end
    return 
end

The intended usage is

@feature function posgx(roadway::Roadway, veh::Entity) 
    return posg(veh).x 
end

The macro should define the function featuretype(::typeof(posgx)) = EntityFeature()
It works but I am having the following concerns:

  • if I try to document one of the @feature functions with docstrings it breaks.
  • it only evaluate the function within the scope of the module where the macro is defined
    I would really appreciate some feedback and help on this, thanks!

I would like to avoid using hasmethod at run time.

Here are some references:

check out Announcing Traits.jl - a revival of Julia traits

and
This comment

1 Like

Hello and welcome to discourse!

As mentioned in the previous answer, for this particular use case, you should probably try to re-use existing features: if you can avoid having to write and maintain such a macro yourself, it will probably be better in the long run.

That being said, for the sake of learning, there are few issues with your macro. A macro should be used to take as input some code that you write (i.e. syntax), and output new code that you don’t want to bother writing (new syntax). Here your macro does not output anything: it only has side effects. Using eval in a macro usually indicates that something should be fixed.

Here is a (simplified) version of what you were trying to achieve, but rewritten in such a way that it does not have any side effect. Instead, it generates some code that will itself have side effects (e.g define new methods).

# Let's put the macro in a new module, so that we can check that there
# are no issues related to module boundaries.
module Feature
export @feature

using MacroTools

macro feature(defun)
    # Use pattern matching features from MacroTools to get the
    # function name
    @capture defun function funname_(args__)
        body_
    end

    quote
        # start with outputting the method definition unchanged
        Base.@__doc__ $defun

        # generate and output some code that will conditionally
        # define new methods
        if hasmethod($funname, Tuple{Int, Float64})
            @doc "docstring for featuretype"
            featuretype(::typeof($funname)) = 1
        end
    end |> esc # escape everything: since this is destined to be
    #            used at top level, we should not have hygiene
    #            problems here
end

end # module

This macro in itself has no side effect. We can analyze what it does by checking the code it generates for a given input. Just to be clear, instead of conditionally generating method definitions, this macro generates code that will conditionally define the required methods:

julia> using .Feature
julia> @macroexpand @feature function posgx(roadway::Int, veh::Float64)
           "OK"
       end
# I cleaned up the output a bit
quote
    function posgx(roadway::Int, veh::Float64)
        "OK"
    end
    if hasmethod(posgx, Tuple{Int, Float64})
        begin
            featuretype(::typeof(posgx)) = begin
                    1
            end
        end
    end
end

Since we’re happy with the generated code, we can try it in real conditions:

julia> "docstring for posgx"
       @feature function posgx(roadway::Int, veh::Float64)
           "OK"
       end
featuretype

# The given method has been defined
julia> posgx(42, 3.)
"OK"

# The trait helper method has also been defined
julia> featuretype(posgx)
1

# And is documented
help?> featuretype
search: featuretype

  docstring for featuretype
6 Likes

Thanks!
I was able to remove the macro and have only one implementation of featuretype using the Tricks.jl package and static_hasmethod

function featuretype(f::Function)
    if static_hasmethod(f, Tuple{Roadway, Entity})
        return EntityFeature()
    elseif static_hasmethod(f, Tuple{Roadway, Frame, Entity})
        return FrameFeature()
    elseif static_hasmethod(f, Tuple{Roadway, Vector{<:Frame}, Entity})
        return TemporalFeature()
    else
        error(""""
              unsupported feature type
              try adding a new feature trait or implement one of the existing feature type supported
              """)
    end
    return 
end

Thank you this was exactly what I wanted to originally do. I did not know how to use the MacroTools package! This is very nice!

If featuretype is defined with the Feature module. Shouldn’t the macro define Feature.featuretype(...)?

Just to check that I understand correctly, to have posgx documented correctly I should add
Base.@__doc__ $defun before $defun is that correct?

The code that defines featuretype is generated by the macro, and then evaluated in the module where the macro is used. So, no: both the new method definition and featuretype are defined in the current module (as they would if you had just written the generated code by hand) , and neither has to be qualified with the Feature module name. That module only defines the @feature macro itself.

Yes, this is correct. I edited my previous post to account for this, thanks!

1 Like

How can I check that static_hasmethod actually works? i.e. that the result featuretype will be inferred at compile time.

@oxinabox would be the one to answer that

@code_warntype or Test.@inferred can help you with such questions. For example:

using Tricks

function featuretype_dyn(f::Function)
    if hasmethod(f, Tuple{Int, Float64})
        return Val(1)
    end
    Val(2)
end

function featuretype_sta(f::Function)
    if static_hasmethod(f, Tuple{Int, Float64})
        return Val(1)
    end
    Val(2)
end

We can see that the dynamic version does not infer:

julia> using Test
julia> @inferred featuretype_dyn(foo)
ERROR: return type Val{1} does not match inferred return type Union{Val{1}, Val{2}}
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] top-level scope at REPL[6]:1

whereas the static one does:

julia> @inferred featuretype_sta(foo)
Val{1}()

Like @ffevotte says, checking if Tricks.jl works is done with @code_typed or @code_warntype.
(@code_llvm and @code_native can also work but generally one of the two flavors of typed is what you want)

I would not use @inferred as I think it leads to less clarity on what is going an, and you don’t really need to do this in your tests.
Nothing you do can change if Tricks.jl works, at least not in a way that can be detected by @inferred.
Tricks can’t generate dynamic code like calls to hasmethod, it can only generate IR that contains basically constant literals.

The thing that I am fairly sure can’t happen (but can happen in many generalizations of Trick’s trick).

Say you have a function
B which calls static_hasmethod(C,...).
for C initially not having that method, and so static_hasmethod returns constant false…

Then when C gains a method, then static_hasmethod(::typeof(C),...) will be redefined to return true.

If the call from B to static_hasmethod is static, then this will intern trigger B to be recompiled to use the new definition of static_hasmethod(::typeof(C),...) (which returns true).
Example of such a static use of static_hasmethod:

B() = static_hasmethod(C, Tuple{String}) ? "Yes" : "No"

This is the case that will fully compile out (see demo below.)

Conversely, If the call from B to static_hasmethod is dynamic, then the redefinition of static_hasmethod(::typeof(C),...) will not trigger a recompilation of B to use the new definition (that always returns true)
but this doesn’t matter because it is still a dynamic call, so it always calls the current definition.
This case will not compile away (see demo below).
But it also will not give an incorrect answer because it will call the current definition
Further is is still about 9x faster than has_method
Example of a dynamic call to static_hasmethod:

B() = static_hasmethod(C, Tuple{rand((Int, Bool, String, Char, Float32))})

The thing to worry about would be if a dynamic call to to a functiom could be inlined, without attaching the edges (recompilation hooks) that static calls attach.
But I am pretty sure we never do that (dynamic calls never get inlined),
and further that doing it wouldn’t make sense anyway.

Static Use of static_hasmethod demo:

julia> using Tricks

julia> B() = static_hasmethod(C, Tuple{String}) ? "Yes" : "No"
B (generic function with 1 method)

julia> @code_typed B()  # No function at all
CodeInfo(
1 ─ %1 = Main.static_hasmethod(Main.C, Tuple{String})::Any
└──      goto #3 if not %1
2 ─      return "Yes"
3 ─      return "No"
) => String

julia> function C end
C (generic function with 0 methods)

julia> @code_typed B()  # no matched method
CodeInfo(
1 ─     goto #3 if not false
2 ─     nothing::Nothing
3 ┄     return "No"
) => String

julia> C(::String) = 1
C (generic function with 1 method)

julia> @code_typed B()   # now there is a matched method
CodeInfo(
1 ─     return "Yes"
) => String

example of a dynamic call to static_hasmethod

julia> using Tricks

julia> B() = static_hasmethod(C, Tuple{rand((String, String))}) ? "Yes" : "No"
B (generic function with 1 method)

julia> @code_typed B()  # function not defined
CodeInfo(
1 ─ %1 = Core.tuple(Main.String, Main.String)::Core.Compiler.Const((String, String), false)
│   %2 = invoke Main.rand(%1::Tuple{DataType,DataType})::DataType
│   %3 = Core.apply_type(Main.Tuple, %2)::Type{var"#s65"} where var"#s65"<:Tuple
│   %4 = Main.static_hasmethod(Main.C, %3)::Any
└──      goto #3 if not %4
2 ─      return "Yes"
3 ─      return "No"
) => String

julia> function C end
C (generic function with 0 methods)

julia> @code_typed B()   # function does not have the method
CodeInfo(
1 ─ %1 = Core.tuple(Main.String, Main.String)::Core.Compiler.Const((String, String), false)
│   %2 = invoke Main.rand(%1::Tuple{DataType,DataType})::DataType
│   %3 = Core.apply_type(Main.Tuple, %2)::Type{var"#s65"} where var"#s65"<:Tuple
│   %4 = Main.static_hasmethod(Main.C, %3)::Any
└──      goto #3 if not %4
2 ─      return "Yes"
3 ─      return "No"
) => String

julia> @code_typed B()  # function does have the method
CodeInfo(
1 ─ %1 = Core.tuple(Main.String, Main.String)::Core.Compiler.Const((String, String), false)
│   %2 = invoke Main.rand(%1::Tuple{DataType,DataType})::DataType
│   %3 = Core.apply_type(Main.Tuple, %2)::Type{var"#s65"} where var"#s65"<:Tuple
│   %4 = Main.static_hasmethod(Main.C, %3)::Any
└──      goto #3 if not %4
2 ─      return "Yes"
3 ─      return "No"
) => String

Would be useful if @jameson and @NHDaly can verify I have said no lies here.

3 Likes