add @leaf macro (#55)

* add `@leaf` macro

* move code around

* fix test, add docstring

* add `@leaf` to documentation

* simplify macro

Author: CarloLucibello

docs/src/api.md

@@ -1,6 +1,7 @@
 ```@docs
 Functors.fmap
 Functors.@functor
+Functors.@leaf
 ```
 
 ```@docs

src/Functors.jl

@@ -282,26 +282,26 @@ julia> struct Bar; x; end
 
 julia> @functor Bar
 
-julia> struct NoChildren; x; y; end
+julia> struct TypeWithNoChildren; x; y; end
 
-julia> m = Foo(Bar([1,2,3]), NoChildren(:a, :b))
-Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+julia> m = Foo(Bar([1,2,3]), TypeWithNoChildren(:a, :b))
+Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))
 
 julia> fcollect(m)
 4-element Vector{Any}:
- Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+ Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))
  Bar([1, 2, 3])
  [1, 2, 3]
- NoChildren(:a, :b)
+ TypeWithNoChildren(:a, :b)
 
 julia> fcollect(m, exclude = v -> v isa Bar)
 2-element Vector{Any}:
- Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
- NoChildren(:a, :b)
+ Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))
+ TypeWithNoChildren(:a, :b)
 
 julia> fcollect(m, exclude = v -> Functors.isleaf(v))
 2-element Vector{Any}:
- Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+ Foo(Bar([1, 2, 3]), TypeWithNoChildren(:a, :b))
  Bar([1, 2, 3])
 ```
 """

src/functor.jl

@@ -1,13 +1,25 @@
+function functor end
 
-functor(T, x) = (), _ -> x
+const NoChildren = Tuple{}
+
+"""
+    @leaf T
+
+Define [`functor`](@ref) for the type `T` so that  `isleaf(x::T) == true`.
+"""
+macro leaf(T)
+  :($Functors.functor(::Type{<:$(esc(T))}, x) = ($Functors.NoChildren(), _ -> x))
+end
+
+@leaf Any # every type is a leaf by default
 functor(x) = functor(typeof(x), x)
 
 functor(::Type{<:Tuple}, x) = x, identity
 functor(::Type{<:NamedTuple{L}}, x) where L = NamedTuple{L}(map(s -> getproperty(x, s), L)), identity
 functor(::Type{<:Dict}, x) = Dict(k => x[k] for k in keys(x)), identity
 
 functor(::Type{<:AbstractArray}, x) = x, identity
-functor(::Type{<:AbstractArray{<:Number}}, x) = (), _ -> x
+@leaf AbstractArray{<:Number}
 
 function makefunctor(m::Module, T, fs = fieldnames(T))
   yᵢ = 0
@@ -31,7 +43,7 @@ macro functor(args...)
   functorm(args...)
 end
 
-isleaf(@nospecialize(x)) = children(x) === ()
+isleaf(@nospecialize(x)) = children(x) === NoChildren()
 
 children(x) = functor(x)[1]
 

test/basics.jl

@@ -24,7 +24,7 @@ struct NoChild{T}; x::T; end
   has_children = Foo(1, 2)
   @test Functors.isleaf(no_children)
   @test !Functors.isleaf(has_children)
-  @test Functors.children(no_children) == ()
+  @test Functors.children(no_children) === Functors.NoChildren()
   @test Functors.children(has_children) == (x=1, y=2)
 end
 
@@ -352,3 +352,14 @@ end
     @test fmap(+, m1, n1) == Dict("x" => [5, 7], "y" => Dict(:a=>3.1, :b=>4.2))
   end
 end
+
+@testset "@leaf" begin
+  struct A; x; end
+  @functor A
+  a = A(1)
+  @test Functors.children(a) === (x = 1,)
+  Functors.@leaf A
+  children, re = Functors.functor(a)
+  @test children == Functors.NoChildren() 
+  @test re(children) === a
+end