add function fleaves (#73)

* define fflatten

* improve docstring

* add to docs

* cl/fflatten

Author: CarloLucibello

docs/src/api.md

@@ -25,4 +25,5 @@ Functors.IterateWalk
 ```@docs
 Functors.fmapstructure
 Functors.fcollect
+Functors.fleaves
 ```

src/Functors.jl

@@ -1,6 +1,6 @@
 module Functors
 
-export @functor, @flexiblefunctor, fmap, fmapstructure, fcollect, execute
+export @functor, @flexiblefunctor, fmap, fmapstructure, fcollect, execute, fleaves
 
 include("functor.jl")
 include("walks.jl")
@@ -303,4 +303,36 @@ julia> fcollect(m, exclude = v -> Functors.isleaf(v))
 """
 fcollect
 
+
+
+"""
+    fleaves(x; exclude = isleaf)
+
+Traverse `x` by recursing each child of `x` as defined by [`functor`](@ref)
+and collecting the leaves into a flat array, 
+ordered by a breadth-first traversal of `x`, respecting the iteration order of `children` calls.
+
+The `exclude` function is used to determine whether to recurse into a node, therefore
+identifying the leaves as the nodes for which `exclude` returns `true`.
+
+See also [`fcollect`](@ref) for a similar function that collects all nodes instead.
+
+# Examples
+
+```jldoctest
+julia> struct Bar; x; end
+
+julia> @functor Bar
+
+julia> struct TypeWithNoChildren; x; y; end
+
+julia> m = (a=Bar([1,2,3]), b=TypeWithNoChildren(4, 5))
+
+julia> fleaves(m)
+2-element Vector{Any}:
+ [1, 2, 3]
+ TypeWithNoChildren(:a, :b)
+"""
+fleaves
+
 end # module

src/maps.jl

@@ -15,3 +15,6 @@ fmapstructure(f, x; kwargs...) = fmap(f, x; walk = StructuralWalk(), kwargs...)
 
 fcollect(x; exclude = v -> false) =
   execute(ExcludeWalk(CollectWalk(), _ -> nothing, exclude), x)
+
+fleaves(x; exclude = isleaf) =
+  execute(ExcludeWalk(FlattenWalk(), x -> [x], exclude), x)

src/walks.jl

@@ -234,7 +234,17 @@ julia> collect(zipped_iter)
 struct IterateWalk <: AbstractWalk end
 
 function (walk::IterateWalk)(recurse, x, ys...)
-  func, _ = functor(x)
-  yfuncs = map(y -> functor(typeof(x), y)[1], ys)
-  return Iterators.flatten(_map(recurse, func, yfuncs...))
+  x_children = children(x)
+  ys_children = map(children, ys)
+  return Iterators.flatten(_map(recurse, x_children, ys_children...))
 end
+
+struct FlattenWalk <: AbstractWalk end
+
+function (walk::FlattenWalk)(recurse, x, ys...)
+  x_children = _values(children(x))
+  ys_children = map(children, ys)
+  res = _map(recurse, x_children, ys_children...)
+  return reduce(vcat, _values(res))
+end
+

test/basics.jl

@@ -392,3 +392,16 @@ end
 @testset "Deprecated first-arg walk API to fmap" begin
   @test (@test_deprecated fmap(Functors.DefaultWalk(), nothing, (1, 2, 3))) == (1, 2, 3)
 end
+
+@testset "fleaves" begin
+  x = (1, (2, 3), (a=4, b=(5, 6), c=7));
+  @test fleaves(x) == [1, 2, 3, 4, 5, 6, 7]
+  @test fleaves(x, exclude=x -> Functors.isleaf(x) || (x isa NamedTuple)) == [1, 2, 3, (a = 4, b = (5, 6), c = 7)]
+
+  x = Dict("a" => Foo(1, 2), "b" => Bar([1,2,3]))
+  xflat = fleaves(x)
+  # @test xflat== [1, 2, [1, 2, 3]] # cannot guarantee ordering with Dict
+  @test xflat isa Vector
+  @test length(xflat) == 3
+  @test 1 ∈ xflat && 2 ∈ xflat && [1, 2, 3] ∈ xflat
+end