add fcollect

src/Functors.jl

@@ -2,7 +2,7 @@ module Functors
 
 using MacroTools
 
-export @functor, fmap
+export @functor, fmap, fcollect
 
 include("functor.jl")
 

src/functor.jl

@@ -29,7 +29,20 @@ macro functor(args...)
   functorm(args...)
 end
 
-isleaf(x) = functor(x)[1] === ()
+"""
+    isleaf(x)
+
+Return true if `x` has no [`children`](@ref) according to [`functor`](@ref).
+"""
+isleaf(x) = children(x) === ()
+
+"""
+    children(x)
+
+Return the children of `x` as defined by [`functor`](@ref).
+Equivalent to `functor(x)[1]`.
+"""
+children(x) = functor(x)[1]
 
 function fmap1(f, x)
   func, re = functor(x)
@@ -42,3 +55,58 @@ function fmap(f, x; cache = IdDict())
   haskey(cache, x) && return cache[x]
   cache[x] = isleaf(x) ? f(x) : fmap1(x -> fmap(f, x, cache = cache), x)
 end
+
+"""
+    fcollect(x; exclude = v -> false)
+
+Traverse `x` by recursing each child of `x` as defined by [`functor`](@ref)
+and collecting the results into a flat array.
+
+Doesn't recurse inside branches rooted at nodes `v`
+for which `exclude(v) == true`.
+In such cases, the root `v` is also excluded from the result.
+By default, `exclude` always yields `false`. 
+
+See also [`children`](@ref).
+
+# Examples
+
+```jldoctest
+julia> struct Foo; x; y; end
+
+julia> @functor Foo
+
+julia> struct Bar; x; end
+
+julia> @functor Bar
+
+julia> struct NoChildren; x; y; end 
+
+julia> m = Foo(Bar([1,2,3]), NoChildren(:a, :b))
+
+julia> fcollect(m)
+4-element Vector{Any}:
+ Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+ Bar([1, 2, 3])
+ [1, 2, 3]
+ NoChildren(: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)
+
+julia> fcollect(m, exclude = v -> Functors.isleaf(v))
+2-element Vector{Any}:
+ Foo(Bar([1, 2, 3]), NoChildren(:a, :b))
+ Bar([1, 2, 3])
+```
+"""
+function fcollect(x; cache = [], exclude = v -> false)
+  x in cache && return cache
+  if !exclude(x)
+    push!(cache, x)
+    foreach(y -> fcollect(y; cache = cache, exclude = exclude), children(x))
+  end
+  return cache
+end

test/basics.jl

@@ -1,19 +1,29 @@
 using Functors, Test
 
-@testset "Nested" begin
-  struct Foo
-    x
-    y
-  end
+struct Foo
+  x
+  y
+end
+@functor Foo
 
-  @functor Foo
+struct Bar
+  x
+end
+@functor Bar
 
-  struct Bar
-    x
-  end
+struct Baz
+  x
+  y
+  z
+end
+@functor Baz (y,)
 
-  @functor Bar
+struct NoChildren 
+  x
+  y
+end
 
+@testset "Nested" begin
   model = Bar(Foo(1, [1, 2, 3]))
 
   model′ = fmap(float, model)
@@ -22,17 +32,27 @@ using Functors, Test
   @test model′.x.y isa Vector{Float64}
 end
 
-@testset "Property list" begin
-  struct Baz
-    x
-    y
-    z
-  end
-
-  @functor Baz (y,)
-
+@testset "Property list" begin  
   model = Baz(1, 2, 3)
   model′ = fmap(x -> 2x, model)
 
   @test (model′.x, model′.y, model′.z) == (1, 4, 3)
 end
+
+@testset "fcollect" begin
+  m1 = [1, 2, 3]
+  m2 = 1
+  m3 = Foo(m1, m2)
+  m4 = Bar(m3)
+  @test all(fcollect(m4) .=== [m4, m3, m1, m2])
+  @test all(fcollect(m4, exclude = x -> x isa Array) .=== [m4, m3, m2])
+  @test all(fcollect(m4, exclude = x -> x isa Foo) .=== [m4])
+
+  m1 = [1, 2, 3]
+  m2 = Bar(m1)
+  m0 = NoChildren(:a, :b)
+  m3 = Foo(m2, m0)
+  m4 = Bar(m3)
+  println(fcollect(m4))
+  @test all(fcollect(m4) .=== [m4, m3, m2, m1, m0])
+end