Merge pull request #41 from JuliaCollections/teh/more_examples

Add binary tree examples

Author: timholy

examples/binarytree_core.jl

@@ -0,0 +1,25 @@
+using AbstractTrees
+
+mutable struct BinaryNode{T}
+    data::T
+    parent::BinaryNode{T}
+    left::BinaryNode{T}
+    right::BinaryNode{T}
+
+    # Root constructor
+    BinaryNode{T}(data) where T = new{T}(data)
+    # Child node constructor
+    BinaryNode{T}(data, parent::BinaryNode{T}) where T = new{T}(data, parent)
+end
+BinaryNode(data) = BinaryNode{typeof(data)}(data)
+
+function leftchild(data, parent::BinaryNode)
+    !isdefined(parent, :left) || error("left child is already assigned")
+    node = typeof(parent)(data, parent)
+    parent.left = node
+end
+function rightchild(data, parent::BinaryNode)
+    !isdefined(parent, :right) || error("right child is already assigned")
+    node = typeof(parent)(data, parent)
+    parent.right = node
+end

examples/binarytree_easy.jl

@@ -0,0 +1,44 @@
+# This file demonstrates a relatively simple way to implement the AbstractTrees
+# interface for a classic binary tree. See "binarytree_infer.jl" for a more performant
+# (but laborious) approach.
+
+if !isdefined(@__MODULE__, :BinaryNode)
+    include("binarytree_core.jl")
+end
+
+## Things we need to define
+function AbstractTrees.children(node::BinaryNode)
+    if isdefined(node, :left)
+        if isdefined(node, :right)
+            return (node.left, node.right)
+        end
+        return (node.left,)
+    end
+    isdefined(node, :right) && return (node.right,)
+    return ()
+end
+
+## Things that make printing prettier
+AbstractTrees.printnode(io::IO, node::BinaryNode) = print(io, node.data)
+
+## Optional enhancements
+# These next two definitions allow inference of the item type in iteration.
+# (They are not sufficient to solve all internal inference issues, however.)
+Base.eltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = BinaryNode{T}
+Base.IteratorEltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = Base.HasEltype()
+
+## Let's test it. First build a tree.
+root = BinaryNode(0)
+l = leftchild(1, root)
+r = rightchild(2, root)
+lr = rightchild(3, l)
+
+print_tree(root)
+collect(PostOrderDFS(root))
+@static if isdefined(@__MODULE__, :Test)
+    @testset "binarytree_easy.jl" begin
+        @test [node.data for node in PostOrderDFS(root)] == [3, 1, 2, 0]
+        @test [node.data for node in PreOrderDFS(root)] == [0, 1, 3, 2]
+        @test [node.data for node in Leaves(root)] == [3, 2]
+    end
+end

examples/binarytree_infer.jl

@@ -0,0 +1,71 @@
+# This file illustrates how to create inferrable tree-iteration methods in circumstances
+# where the children are not naturally indexable.
+# See "binarytree_easy.jl" for a simpler approach.
+
+if !isdefined(@__MODULE__, :BinaryNode)
+    include("binarytree_core.jl")
+end
+
+## Enhancement of the "native" binary tree
+# You might have the methods below even if you weren't trying to support AbstractTrees.
+
+# Implement iteration over the immediate children of a node
+function Base.iterate(node::BinaryNode)
+    isdefined(node, :left) && return (node.left, false)
+    isdefined(node, :right) && return (node.right, true)
+    return nothing
+end
+function Base.iterate(node::BinaryNode, state::Bool)
+    state && return nothing
+    isdefined(node, :right) && return (node.right, true)
+    return nothing
+end
+Base.IteratorSize(::Type{BinaryNode{T}}) where T = Base.SizeUnknown()
+Base.eltype(::Type{BinaryNode{T}}) where T = BinaryNode{T}
+
+## Things we need to define to leverage the native iterator over children
+## for the purposes of AbstractTrees.
+# Set the traits of this kind of tree
+Base.eltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = BinaryNode{T}
+Base.IteratorEltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = Base.HasEltype()
+AbstractTrees.parentlinks(::Type{BinaryNode{T}}) where T = AbstractTrees.StoredParents()
+AbstractTrees.siblinglinks(::Type{BinaryNode{T}}) where T = AbstractTrees.StoredSiblings()
+# Use the native iteration for the children
+AbstractTrees.children(node::BinaryNode) = node
+
+Base.parent(root::BinaryNode, node::BinaryNode) = isdefined(node, :parent) ? node.parent : nothing
+
+function AbstractTrees.nextsibling(tree::BinaryNode, child::BinaryNode)
+    isdefined(child, :parent) || return nothing
+    p = child.parent
+    if isdefined(p, :right)
+        child === p.right && return nothing
+        return p.right
+    end
+    return nothing
+end
+
+# We also need `pairs` to return something sensible.
+# If you don't like integer keys, you could do, e.g.,
+#   Base.pairs(node::BinaryNode) = BinaryNodePairs(node)
+# and have its iteration return, e.g., `:left=>node.left` and `:right=>node.right` when defined.
+# But the following is easy:
+Base.pairs(node::BinaryNode) = enumerate(node)
+
+
+AbstractTrees.printnode(io::IO, node::BinaryNode) = print(io, node.data)
+
+root = BinaryNode(0)
+l = leftchild(1, root)
+r = rightchild(2, root)
+lr = rightchild(3, l)
+
+print_tree(root)
+collect(PostOrderDFS(root))
+@static if isdefined(@__MODULE__, :Test)
+    @testset "binarytree_infer.jl" begin
+        @test @inferred(map(x->x.data, PostOrderDFS(root))) == [3, 1, 2, 0]
+        @test @inferred(map(x->x.data, PreOrderDFS(root))) == [0, 1, 3, 2]
+        @test @inferred(map(x->x.data, Leaves(root))) == [3, 2]
+    end
+end

examples/fstree.jl

@@ -24,6 +24,6 @@ end
 printnode(io::IO, d::Directory) = print(io, basename(d.path))
 printnode(io::IO, f::File) = print(io, basename(f.path))
 
-dirpath = realpath(joinpath(dirname(pathof(AbstractTrees)),".."))
+dirpath = realpath(dirname(@__DIR__))
 d = Directory(dirpath)
 print_tree(d)

src/AbstractTrees.jl

@@ -23,7 +23,8 @@ include("implicitstacks.jl")
     children(x)
 
 Return the immediate children of node `x`. You should specialize this method
-for custom tree structures.
+for custom tree structures. It should return an iterable object for which an
+appropriate implementation of `Base.pairs` is available.
 
 # Example
 
@@ -201,7 +202,7 @@ show(io::IO, tree::Tree) = print_tree(io, tree.x)
 
 mutable struct AnnotationNode{T}
     val::T
-    children::Array{AnnotationNode{T}}
+    children::Vector{AnnotationNode{T}}
 end
 
 children(x::AnnotationNode) = x.children
@@ -317,12 +318,11 @@ Any[1,Any[2,3]]
 
 we will get [1,2,3,Any[2,3],Any[1,Any[2,3]]]
 """
-struct PostOrderDFS <: TreeIterator{Any}
-    tree::Any
-    PostOrderDFS(x::Any) = new(x)
+struct PostOrderDFS{T} <: TreeIterator{T}
+    tree::T
 end
 PostOrderDFS(tree::Tree) = PostOrderDFS(tree.x)
-IteratorSize(::Type{PostOrderDFS}) = SizeUnknown()
+IteratorSize(::Type{PostOrderDFS{T}}) where T = SizeUnknown()
 
 """
 Iterator to visit the nodes of a tree, guaranteeing that parents
@@ -378,16 +378,8 @@ parentstate(tree, state) = parentstate(tree, state, treekind(tree))
 update_state!(old_state, cs, idx) = next(cs, idx)[1]
 
 
-struct ImplicitRootState
-end
 getindex(x::AbstractArray, ::ImplicitRootState) = x
 
-"""
-Trees must override with method if the state of the root is not the same as the
-tree itself (e.g. IndexedTrees should always override this method).
-"""
-rootstate(x) = ImplicitRootState()
-
 function firststate(ti::PreOrderDFS{T}) where T
     if isa(parentlinks(ti.tree), StoredParents) &&
             isa(siblinglinks(ti.tree), SiblingLinks)
@@ -415,6 +407,7 @@ relative_state(tree, parentstate, childstate::ImplicitIndexStack) =
     childstate.stack[end]
 relative_state(tree, parentstate, childstate::ImplicitNodeStack) =
     relative_state(tree, parentstate, childstate.idx_stack)
+
 function nextsibling(tree, state)
     ps = parentstate(tree, state)
     cs = childstates(tree, ps)
@@ -435,7 +428,7 @@ function nextsibling(node, ::StoredParents, ::ImplicitSiblings, ::RegularTree)
     last_was_node && return nothing
     error("Tree inconsistency: node not a child of parent")
 end
-nextsibling(node, ::Any, ::StoredSiblings, ::Any) = error("Trees with explicit siblings must override the `prevsibling` method explicitly")
+nextsibling(node, ::Any, ::StoredSiblings, ::Any) = error("Trees with explicit siblings must override the `nextsibling` method explicitly")
 nextsibling(node) = nextsibling(node, parentlinks(node), siblinglinks(node), treekind(node))
 
 function prevsibling(node, ::StoredParents, ::ImplicitSiblings, ::RegularTree)
@@ -465,6 +458,8 @@ end
 children(tree::Subtree) = children(tree.tree, tree.state)
 nodetype(tree::Subtree) = nodetype(tree.tree)
 idxtype(tree::Subtree) = idxtype(tree.tree)
+rootstate(tree::Subtree) = tree.state
+parentlinks(::Type{Subtree{T,S}}) where {T,S} = parentlinks(T)
 
 joinstate(tree, a, b) = b
 
@@ -495,8 +490,11 @@ function stepstate(ti::TreeIterator, state)
     nothing
 end
 
-getnode(tree, ns) = isa(treekind(tree), IndexedTree) ? tree[ns] : ns
 getnode(tree::AbstractShadowTree, ns::ImplicitNodeStack) = tree[ns.idx_stack.stack]
+getnode(tree, ns) = getnode(tree, ns, treekind(tree))
+getnode(tree, ns, ::IndexedTree) = tree[ns]
+getnode(tree, ns, ::RegularTree) = ns
+getnode(tree, ::ImplicitRootState, ::RegularTree) = tree
 
 function iterate(ti::TreeIterator)
     state = firststate(ti)

src/traits.jl

@@ -26,6 +26,17 @@ struct ImplicitSiblings <: SiblingLinks; end
 siblinglinks(::Type) = ImplicitSiblings()
 siblinglinks(tree) = siblinglinks(typeof(tree))
 
+struct ImplicitRootState
+end
+
+"""
+    state = rootstate(tree)
+
+Trees must override with method if the state of the root is not the same as the
+tree itself (e.g. IndexedTrees should always override this method).
+"""
+rootstate(x) = ImplicitRootState()
+
 
 abstract type TreeKind end
 struct RegularTree <: TreeKind; end
@@ -34,12 +45,10 @@ struct IndexedTree <: TreeKind; end
 treekind(tree::Type) = RegularTree()
 treekind(tree) = treekind(typeof(tree))
 children(tree, node, ::RegularTree) = children(node)
+children(tree, ::ImplicitRootState, ::RegularTree) = children(tree)
 children(tree, node, ::IndexedTree) = (tree[y] for y in childindices(tree, node))
 children(tree, node) = children(tree, node, treekind(tree))
 
-function rootstate()
-end
-
 childindices(tree, node) =
   tree == node ? childindices(tree, rootstate(tree)) :
   error("Must implement childindices(tree, node)")

test/runtests.jl

@@ -2,17 +2,16 @@ using AbstractTrees
 using Test
 import Base: ==
 
-if !isdefined(@__MODULE__, :isfirstrun)
-    const isfirstrun = Ref(true)
-end
-if isfirstrun[] && VERSION >= v"1.1.0-DEV.838" # requires https://github.com/JuliaLang/julia/pull/30291
-    @test isempty(detect_ambiguities(AbstractTrees, Base, Core))
-    isfirstrun[] = false
+if VERSION >= v"1.1.0-DEV.838" # requires https://github.com/JuliaLang/julia/pull/30291
+    @testset "Ambiguities" begin
+        @test isempty(detect_ambiguities(AbstractTrees, Base, Core))
+    end
 end
 
 AbstractTrees.children(x::Array) = x
 tree = Any[1,Any[2,3]]
 
+@testset "Array" begin
 io = IOBuffer()
 print_tree(io, tree)
 @test String(take!(io)) == "Array{Any,1}\n├─ 1\n└─ Array{Any,1}\n   ├─ 2\n   └─ 3\n"
@@ -23,6 +22,7 @@ print_tree(io, tree)
 
 tree2 = Any[Any[1,2],Any[3,4]]
 @test collect(PreOrderDFS(tree2)) == Any[tree2,Any[1,2],1,2,Any[3,4],3,4]
+end
 
 """
     A tree in which every node has 0 or 1 children
@@ -44,12 +44,16 @@ Base.eltype(::Type{<:TreeIterator{OneTree}}) = Int
 Base.IteratorEltype(::Type{<:TreeIterator{OneTree}}) = Base.HasEltype()
 
 ot = OneTree([2,3,4,0])
+
+@testset "OneTree" begin
+io = IOBuffer()
 print_tree(io, ot)
 @test String(take!(io)) == "2\n└─ 3\n   └─ 4\n      └─ 0\n"
 @test @inferred(collect(Leaves(ot))) == [0]
 @test eltype(collect(Leaves(ot))) === Int
 @test collect(PreOrderDFS(ot)) == [2,3,4,0]
 @test collect(PostOrderDFS(ot)) == [0,4,3,2]
+end
 
 """
     Stores an explicit parent for some other kind of tree
@@ -72,6 +76,8 @@ AbstractTrees.printnode(io::IO, t::ParentTree) =
     AbstractTrees.printnode(io::IO, t[AbstractTrees.rootstate(t)])
 
 pt = ParentTree(ot,[0,1,2,3])
+@testset "ParentTree" begin
+io = IOBuffer()
 print_tree(io, pt)
 @test String(take!(io)) == "2\n└─ 3\n   └─ 4\n      └─ 0\n"
 @test collect(Leaves(pt)) == [0]
@@ -88,7 +94,8 @@ b = treemap!(PreOrderDFS(a)) do node
     empty!(node)
     ret
 end
-@assert b == Any[0,1,Any[0,2,[0,3]]]
+@test b == Any[0,1,Any[0,2,[0,3]]]
+end
 
 struct IntTree
     num::Int
@@ -102,10 +109,12 @@ Base.eltype(::Type{<:TreeIterator{IntTree}}) = IntTree
 Base.IteratorEltype(::Type{<:TreeIterator{IntTree}}) = Base.HasEltype()
 AbstractTrees.nodetype(::IntTree) = IntTree
 iter = Leaves(itree)
+@testset "IntTree" begin
 @test @inferred(first(iter)) == IntTree(2, IntTree[])
 val, state = iterate(iter)
 @test Base.return_types(iterate, Tuple{typeof(iter), typeof(state)}) ==
     [Union{Nothing, Tuple{IntTree,typeof(state)}}]
+end
 
 #=
 @test treemap(PostOrderDFS(tree)) do ind, x, children
@@ -115,3 +124,17 @@ end == IntTree(6,[IntTree(1,IntTree[]),IntTree(5,[IntTree(2,IntTree[]),IntTree(3
 =#
 
 @test collect(PostOrderDFS([])) == Any[[]]
+
+@testset "Examples" begin
+# Ensure the examples run
+exampledir = joinpath(dirname(@__DIR__), "examples")
+examples = readdir(exampledir)
+mktemp() do filename, io
+    redirect_stdout(io) do
+        for ex in examples
+            haskey(ENV, "CI") && Sys.isapple() && ex == "fstree.jl" && continue
+            include(joinpath(exampledir, ex))
+        end
+    end
+end
+end  # @testset "Examples"