treemap now allows for modification of children

Author: ExpandingMan

src/iteration.jl

@@ -412,7 +412,6 @@ function Base.iterate(ti::StatelessBFS, ind)
 end
 
 
-#TODO: still experimenting here
 """
     MapNode{T,C}
 
@@ -428,13 +427,13 @@ struct MapNode{T,C}
     value::T
     children::C
 
-    #TODO: need a way to modify children but hard to think of a good general way of doing it...
-    function MapNode(𝒻, node)
-        v = 𝒻(node)
-        ch = map(c -> MapNode(𝒻, c), children(node))
-        new{typeof(v),typeof(ch)}(v, ch)
+    function MapNode(f, node)
+        (v, ch) = f(node)
+        ch′ = map(c -> MapNode(f, c), ch)
+        new{typeof(v),typeof(ch′)}(v, ch′)
     end
 end
+MapNode(node) = MapNode(n -> (nodevalue(n), children(n)), nodevalue(node))
 
 children(μ::MapNode) = μ.children
 
@@ -452,31 +451,51 @@ end
 Base.show(io::IO, ::MIME"text/plain", μ::MapNode) = print_tree(io, μ)
 
 
-#TODO: still experimenting here
 """
     treemap(f, node)
 
-Apply the function `f` to every node in the tree with root `node`.  `node` must satisfy the AbstractTrees interface.
-Instead of returning the result of `f(n)` directly the result will be a tree of [`MapNode`](@ref) objects isomorphic
-to the original tree but with values equal to the corresponding `f(n)`.
+Apply the function `f` to every node in the tree with root `node`, where `f` is a function that returns `(v, ch)` where
+`v` is a new value for the node (i.e. as returned by [`nodevalue`](@ref) and `ch` is the new children of the node.
+`f` will be called recursively so that all the children returned by `f` for a parent node will again be called for
+each child.  This means that to maintain the structure of a tree but merely map to new values, one should define
+`f = node -> (g(node), children(node))` for some function `g` which returns a value for the node.
+
+The nodes of the output tree will all be represented by [`MapNode`](@ref) objects wrapping the returned values.
+This is necessary in order to guarantee that the output types can describe any tree topology.
 
 Note that in most common cases tree nodes are of a type which depends on their connectedness and the function
 `f` should take this into account.  For example the tree `[1, [2, 3]]` contains integer leaves but two
 `Vector` nodes.  Therefore, the `f` in `treemap(f, [1, [2,3]])` must be a function that is valid for either
 `Int` or `Vector`.  Alternatively, to only operate on leaves do `map(𝒻, Leaves(itr))`.
 
-## Example
+It's very easy to write an `f` that makes `treemap` stack-overflow.  To avoid this, ensure that `f` eventually
+termiantes, i.e. that sometimes it returns empty `children`.  For example, if `f(n) = (nothing, [0; children(n)])` will
+stack-overflow because every node will have at least 1 child.
+
+## Examples
 ```julia
 julia> t = [1, [2, 3]];
 
-julia> f(x) = x isa AbstractArray ? nothing : x + 1;
+julia> f(n) = n isa AbstractArray ? (nothing, children(n)) : (n+1, children(n))
 
 julia> treemap(f, t)
-MapNode{Nothing, MapNode}(nothing)
-├─ MapNode{Int64, MapNode{Union{}}}(2)
-└─ MapNode{Nothing, MapNode{Int64, MapNode{Union{}}}}(nothing)
-   ├─ MapNode{Int64, MapNode{Union{}}}(3)
-   └─ MapNode{Int64, MapNode{Union{}}}(4)
+nothing
+├─ 2
+└─ nothing
+   ├─ 3
+   └─ 4
+
+julia> g(n) = isempty(children(n)) ? (nodevalue(n), []) : (nodevalue(n), [0; children(n)])
+g (generic function with 1 method)
+
+julia> treemap(g, t)
+Any[1, [2, 3]]
+├─ 0
+├─ 1
+└─ [2, 3]
+   ├─ 0
+   ├─ 2
+   └─ 3
 ```
 """
 treemap(f, node) = MapNode(f, node)

test/trees.jl

@@ -85,21 +85,15 @@ include(joinpath(@__DIR__,"examples","onetree.jl"))
     @test nodevalue.(collect(PostOrderDFS(n))) == [0,4,3,2]
 end
 
-#=
-@testset "treemap!" begin
-    # Test modification while iterating over PreOrderDFS
+@testset "treemap" begin
     a = [1,[2,[3]]]
-    b = treemap!(PreOrderDFS(a)) do node
-        !isa(node, Vector) && return node
-        ret = pushfirst!(copy(node),0)
-        # And just for good measure stomp over the old node to make sure nothing
-        # is cached.
-        empty!(node)
-        ret
-    end
-    @test b == Any[0,1,Any[0,2,[0,3]]]
+    f = n -> n isa AbstractArray ? (nothing, children(n)) : (n+1, children(n))
+    b = treemap(f, a)
+    @test nodevalue.(PreOrderDFS(b)) == [nothing, 2, nothing, 3, nothing, 4]
+    g = n -> isempty(children(n)) ? (nodevalue(n), ()) : (nothing, [0; children(n)])
+    b = treemap(g, a)
+    @test nodevalue.(PostOrderDFS(b)) == [0, 1, 0, 2, 0, 3, nothing, nothing, nothing]
 end
-=#
 
 #=
 struct IntTree