├── docs ├── .gitignore ├── Project.toml ├── make.jl └── src │ ├── api.md │ └── index.md ├── .gitignore ├── .github └── workflows │ └── TagBot.yml ├── Project.toml ├── README.md ├── NEWS.md ├── examples ├── fstree.jl ├── binarytree_core.jl ├── binarytree_easy.jl └── binarytree_infer.jl ├── .travis.yml ├── LICENSE.md ├── src ├── AbstractTrees.jl ├── traits.jl ├── indexing.jl ├── implicitstacks.jl ├── printing.jl └── iteration.jl └── test └── runtests.jl /docs/.gitignore: -------------------------------------------------------------------------------- 1 | build 2 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.jl.cov 2 | *.jl.mem 3 | Manifest.toml 4 | -------------------------------------------------------------------------------- /docs/Project.toml: -------------------------------------------------------------------------------- 1 | [deps] 2 | Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" 3 | 4 | [compat] 5 | Documenter = "0.24" 6 | -------------------------------------------------------------------------------- /.github/workflows/TagBot.yml: -------------------------------------------------------------------------------- 1 | name: TagBot 2 | on: 3 | schedule: 4 | - cron: 0 * * * * 5 | jobs: 6 | TagBot: 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: JuliaRegistries/TagBot@v1 10 | with: 11 | token: ${{ secrets.GITHUB_TOKEN }} 12 | -------------------------------------------------------------------------------- /Project.toml: -------------------------------------------------------------------------------- 1 | name = "AbstractTrees" 2 | uuid = "1520ce14-60c1-5f80-bbc7-55ef81b5835c" 3 | version = "0.3.3" 4 | authors = ["Keno Fischer "] 5 | 6 | [deps] 7 | Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" 8 | 9 | [compat] 10 | julia = "0.7, 1" 11 | 12 | [extras] 13 | Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" 14 | 15 | [targets] 16 | test = ["Test"] 17 | -------------------------------------------------------------------------------- /docs/make.jl: -------------------------------------------------------------------------------- 1 | using Documenter, AbstractTrees 2 | 3 | makedocs( 4 | modules = [AbstractTrees], 5 | sitename = "AbstractTrees.jl", 6 | authors = "Keno Fischer", 7 | format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"), 8 | pages = [ 9 | "index.md", 10 | "api.md", 11 | ], 12 | ) 13 | 14 | deploydocs( 15 | repo = "github.com/JuliaCollections/AbstractTrees.jl.git", 16 | ) 17 | -------------------------------------------------------------------------------- /docs/src/api.md: -------------------------------------------------------------------------------- 1 | # API 2 | 3 | 4 | ```@docs 5 | AnnotationNode 6 | AbstractTrees.ImplicitParents 7 | AbstractTrees.ImplicitSiblings 8 | Leaves 9 | AbstractTrees.ParentLinks 10 | PostOrderDFS 11 | PreOrderDFS 12 | ShadowTree 13 | AbstractTrees.SiblingLinks 14 | StatelessBFS 15 | AbstractTrees.StoredParents 16 | AbstractTrees.StoredSiblings 17 | Tree 18 | TreeCharSet 19 | TreeIterator 20 | children 21 | AbstractTrees.nodetype 22 | AbstractTrees.parentlinks 23 | print_tree 24 | AbstractTrees.printnode 25 | AbstractTrees.siblinglinks 26 | treemap 27 | treemap! 28 | ``` 29 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AbstractTrees 2 | 3 | [![Build Status](https://travis-ci.com/JuliaCollections/AbstractTrees.jl.svg?branch=master)](https://travis-ci.com/JuliaCollections/AbstractTrees.jl) 4 | [![codecov](https://codecov.io/gh/JuliaCollections/AbstractTrees.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/JuliaCollections/AbstractTrees.jl) 5 | [![](https://img.shields.io/badge/docs-stable-blue.svg)](https://JuliaCollections.github.io/AbstractTrees.jl/stable) 6 | 7 | This package provides several utilities for working with tree-like data structures. 8 | See the [documentation](https://JuliaCollections.github.io/AbstractTrees.jl/stable) 9 | for details, 10 | and [NEWS.md](NEWS.md) for information about important changes. 11 | -------------------------------------------------------------------------------- /NEWS.md: -------------------------------------------------------------------------------- 1 | # NEWS.md 2 | 3 | ## Breaking changes in v0.3 4 | 5 | - `getindex(::Any, ::ImplicitRootState)` is no longer defined; packages 6 | that used this method will now throw a `MethodError`. To circumvent this, 7 | define `Base.getindex(x::MyTreeType, ::AbstractTrees.ImplicitRootState) = x`. 8 | - By default, the iterators in this package now have the 9 | `Base.IteratorEltype` trait set to `Base.EltypeUnknown()`. 10 | This generally produces "narrower" (more concretely-typed) arrays when 11 | used in conjunction with `collect`. 12 | However, if you define the method `Base.eltype(::Type{<:TreeIterator{MyTreeType}})`, 13 | you should also set `Base.IteratorEltype(::Type{<:TreeIterator{MyTreeType}}) = Base.HasEltype()`. 14 | -------------------------------------------------------------------------------- /examples/fstree.jl: -------------------------------------------------------------------------------- 1 | using AbstractTrees 2 | import AbstractTrees: children, printnode 3 | 4 | struct File 5 | path::String 6 | end 7 | 8 | children(f::File) = () 9 | 10 | struct Directory 11 | path::String 12 | end 13 | 14 | function children(d::Directory) 15 | contents = readdir(d.path) 16 | children = Vector{Union{Directory,File}}(undef,length(contents)) 17 | for (i,c) in enumerate(contents) 18 | path = joinpath(d.path,c) 19 | children[i] = isdir(path) ? Directory(path) : File(path) 20 | end 21 | return children 22 | end 23 | 24 | printnode(io::IO, d::Directory) = print(io, basename(d.path)) 25 | printnode(io::IO, f::File) = print(io, basename(f.path)) 26 | 27 | dirpath = realpath(dirname(@__DIR__)) 28 | d = Directory(dirpath) 29 | print_tree(d) 30 | -------------------------------------------------------------------------------- /examples/binarytree_core.jl: -------------------------------------------------------------------------------- 1 | using AbstractTrees 2 | 3 | mutable struct BinaryNode{T} 4 | data::T 5 | parent::BinaryNode{T} 6 | left::BinaryNode{T} 7 | right::BinaryNode{T} 8 | 9 | # Root constructor 10 | BinaryNode{T}(data) where T = new{T}(data) 11 | # Child node constructor 12 | BinaryNode{T}(data, parent::BinaryNode{T}) where T = new{T}(data, parent) 13 | end 14 | BinaryNode(data) = BinaryNode{typeof(data)}(data) 15 | 16 | function leftchild(data, parent::BinaryNode) 17 | !isdefined(parent, :left) || error("left child is already assigned") 18 | node = typeof(parent)(data, parent) 19 | parent.left = node 20 | end 21 | function rightchild(data, parent::BinaryNode) 22 | !isdefined(parent, :right) || error("right child is already assigned") 23 | node = typeof(parent)(data, parent) 24 | parent.right = node 25 | end 26 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: julia 2 | os: 3 | - linux 4 | - osx 5 | julia: 6 | - 0.7 7 | - 1.0 8 | - 1 9 | - nightly 10 | 11 | notifications: 12 | email: false 13 | # uncomment the following lines to override the default test script 14 | #script: 15 | # - if [[ -a .git/shallow ]]; then git fetch --unshallow; fi 16 | # - julia --check-bounds=yes -e 'Pkg.clone(pwd()); Pkg.build("AbstractTrees"); Pkg.test("AbstractTrees"; coverage=true)' 17 | 18 | after_success: 19 | - julia -e 'using Pkg; Pkg.add("Coverage"); using Coverage; Codecov.submit(process_folder())' 20 | 21 | jobs: 22 | include: 23 | - stage: "Documentation" 24 | julia: 1.0 25 | os: linux 26 | script: 27 | - julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); 28 | Pkg.instantiate()' 29 | - julia --project=docs/ docs/make.jl 30 | after_success: skip 31 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | The AbstractTrees.jl package is licensed under the MIT "Expat" License: 2 | 3 | > Copyright (c) 2015: Keno Fischer. 4 | > 5 | > Permission is hereby granted, free of charge, to any person obtaining 6 | > a copy of this software and associated documentation files (the 7 | > "Software"), to deal in the Software without restriction, including 8 | > without limitation the rights to use, copy, modify, merge, publish, 9 | > distribute, sublicense, and/or sell copies of the Software, and to 10 | > permit persons to whom the Software is furnished to do so, subject to 11 | > the following conditions: 12 | > 13 | > The above copyright notice and this permission notice shall be 14 | > included in all copies or substantial portions of the Software. 15 | > 16 | > THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | > EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | > MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 19 | > IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY 20 | > CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, 21 | > TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 22 | > SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 23 | -------------------------------------------------------------------------------- /examples/binarytree_easy.jl: -------------------------------------------------------------------------------- 1 | # This file demonstrates a relatively simple way to implement the AbstractTrees 2 | # interface for a classic binary tree. See "binarytree_infer.jl" for a more performant 3 | # (but laborious) approach. 4 | 5 | if !isdefined(@__MODULE__, :BinaryNode) 6 | include("binarytree_core.jl") 7 | end 8 | 9 | ## Things we need to define 10 | function AbstractTrees.children(node::BinaryNode) 11 | if isdefined(node, :left) 12 | if isdefined(node, :right) 13 | return (node.left, node.right) 14 | end 15 | return (node.left,) 16 | end 17 | isdefined(node, :right) && return (node.right,) 18 | return () 19 | end 20 | 21 | ## Things that make printing prettier 22 | AbstractTrees.printnode(io::IO, node::BinaryNode) = print(io, node.data) 23 | 24 | ## Optional enhancements 25 | # These next two definitions allow inference of the item type in iteration. 26 | # (They are not sufficient to solve all internal inference issues, however.) 27 | Base.eltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = BinaryNode{T} 28 | Base.IteratorEltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = Base.HasEltype() 29 | 30 | ## Let's test it. First build a tree. 31 | root = BinaryNode(0) 32 | l = leftchild(1, root) 33 | r = rightchild(2, root) 34 | lr = rightchild(3, l) 35 | 36 | print_tree(root) 37 | collect(PostOrderDFS(root)) 38 | @static if isdefined(@__MODULE__, :Test) 39 | @testset "binarytree_easy.jl" begin 40 | @test [node.data for node in PostOrderDFS(root)] == [3, 1, 2, 0] 41 | @test [node.data for node in PreOrderDFS(root)] == [0, 1, 3, 2] 42 | @test [node.data for node in Leaves(root)] == [3, 2] 43 | end 44 | end 45 | -------------------------------------------------------------------------------- /docs/src/index.md: -------------------------------------------------------------------------------- 1 | # AbstractTrees.jl 2 | 3 | 4 | This package provides several utilities for working with tree-like data 5 | structures. Most importantly, it defines the [`children`](@ref) method that any 6 | package that contains such a data structure may import and extend in order to 7 | take advantage of any generic tree algorithm in this package (or other packages 8 | compatible with this package). 9 | 10 | 11 | ## API overview 12 | 13 | * [`print_tree`](@ref) pretty prints an arbitrary tree data structure. 14 | * [`Tree`](@ref) is a simple wrapper around an arbitrary object that allows tree-indexing into that object (i.e. indexing with collections of indices specifying the child index at every level). 15 | * [`ShadowTree`](@ref) is a tree object that combines two trees of equal structure into a single tree (indexing always produces another `ShadowTree`, but `setindex!` with tuples is allowed). Useful for adding annotation nodes to other trees without modifying that tree structure itself. 16 | * [`Leaves`](@ref) is an iterator to visit the leaves of a tree in order. 17 | * [`PostOrderDFS`](@ref) is a depth-first search (i.e. will visit node's children before it's lexicographically following siblings) that guarantees to visit children before their parents. 18 | * [`PreOrderDFS`](@ref) is same as `PostOrderDFS` but visits parents before their children. 19 | * [`StatelessBFS`](@ref) iterates over a tree level-by-level, but does not keep state (causing this to be $O(n^2)$, but able to handle changing trees). 20 | * [`treemap`](@ref) maps each node of a tree to obtain a new tree. 21 | * [`treemap!`](@ref) maps each node of a tree in place. 22 | 23 | 24 | ## Traits 25 | 26 | * [`AbstractTrees.nodetype(tree)`](@ref) can be defined to make iteration inferable. 27 | * [`AbstractTrees.ParentLinks`](@ref) can be defined to return [`AbstractTrees.StoredParents()`](@ref) if a tree type stores explicit links to a parent; [`AbstractTrees.SiblingLinks`](@ref), when set to [`AbstractTrees.StoredSiblings()`](@ref), serves the same role for siblings. See their docstrings for more information. 28 | 29 | ## Examples 30 | 31 | The [examples folder](https://github.com/JuliaCollections/AbstractTrees.jl/tree/master/examples) contains a number of usage examples of varying complexity. 32 | -------------------------------------------------------------------------------- /src/AbstractTrees.jl: -------------------------------------------------------------------------------- 1 | """ 2 | This package is intended to provide an abstract interface for working 3 | with tree structures. 4 | Though the package itself is not particularly sophisticated, it defines 5 | the interface that can be used by other packages to talk about trees. 6 | """ 7 | module AbstractTrees 8 | 9 | export print_tree, TreeCharSet, TreeIterator, Leaves, PostOrderDFS, Tree, 10 | AnnotationNode, StatelessBFS, treemap, treemap!, PreOrderDFS, 11 | ShadowTree, children 12 | 13 | import Base: getindex, setindex!, iterate, nextind, print, show, 14 | eltype, IteratorSize, IteratorEltype, length, push!, pop! 15 | using Base: SizeUnknown, EltypeUnknown 16 | using Markdown 17 | 18 | 19 | abstract type AbstractShadowTree end 20 | 21 | 22 | """ 23 | children(x) 24 | 25 | Return the immediate children of node `x`. You should specialize this method 26 | for custom tree structures. It should return an iterable object for which an 27 | appropriate implementation of `Base.pairs` is available. 28 | 29 | The default behavior is to assume that if an object is iterable, iterating over 30 | it gives its children. If an object is not iterable, assume it does not have 31 | children. 32 | 33 | # Example 34 | 35 | ``` 36 | struct MyNode{T} 37 | data::T 38 | children::Vector{MyNode{T}} 39 | end 40 | AbstractTrees.children(node::MyNode) = node.children 41 | ``` 42 | """ 43 | children(x) = Base.isiterable(typeof(x)) ? x : () 44 | 45 | has_children(x) = children(x) !== () 46 | 47 | 48 | include("traits.jl") 49 | include("implicitstacks.jl") 50 | include("printing.jl") 51 | include("indexing.jl") 52 | include("iteration.jl") 53 | 54 | 55 | ## Special cases 56 | 57 | # Types which are iterable but shouldn't be considered tree-iterable 58 | children(x::Number) = () 59 | children(x::Char) = () 60 | children(x::Task) = () 61 | children(x::AbstractString) = () 62 | 63 | # Define this here, there isn't really a good canonical package to define this 64 | # elsewhere 65 | children(x::Expr) = x.args 66 | 67 | # For AbstractDicts 68 | 69 | printnode(io::IO, kv::Pair{K,V}) where {K,V} = printnode(io,kv[1]) 70 | children(kv::Pair{K,V}) where {K,V} = (kv[2],) 71 | 72 | # For potentially-large containers, just show the type 73 | 74 | printnode(io::IO, ::T) where T <: Union{AbstractArray, AbstractDict} = print(io, T) 75 | 76 | 77 | end # module 78 | -------------------------------------------------------------------------------- /src/traits.jl: -------------------------------------------------------------------------------- 1 | abstract type ParentLinks end 2 | 3 | """ 4 | Indicates that this tree stores parent links explicitly. The implementation 5 | is responsible for defining the parentind function to expose this 6 | information. 7 | """ 8 | struct StoredParents <: ParentLinks; end 9 | struct ImplicitParents <: ParentLinks; end 10 | 11 | parentlinks(::Type) = ImplicitParents() 12 | parentlinks(tree) = parentlinks(typeof(tree)) 13 | 14 | abstract type SiblingLinks end 15 | 16 | """ 17 | Indicates that this tree stores sibling links explicitly, or can compute them 18 | quickly (e.g. because the tree has a (small) fixed branching ratio, so the 19 | current index of a node can be determined by quick linear search). The 20 | implementation is responsible for defining the relative_state function 21 | to expose this information. 22 | """ 23 | struct StoredSiblings <: SiblingLinks; end 24 | struct ImplicitSiblings <: SiblingLinks; end 25 | 26 | siblinglinks(::Type) = ImplicitSiblings() 27 | siblinglinks(tree) = siblinglinks(typeof(tree)) 28 | 29 | struct ImplicitRootState 30 | end 31 | 32 | """ 33 | state = rootstate(tree) 34 | 35 | Trees must override with method if the state of the root is not the same as the 36 | tree itself (e.g. IndexedTrees should always override this method). 37 | """ 38 | rootstate(x) = ImplicitRootState() 39 | 40 | 41 | abstract type TreeKind end 42 | struct RegularTree <: TreeKind; end 43 | struct IndexedTree <: TreeKind; end 44 | 45 | treekind(tree::Type) = RegularTree() 46 | treekind(tree) = treekind(typeof(tree)) 47 | children(tree, node, ::RegularTree) = children(node) 48 | children(tree, ::ImplicitRootState, ::RegularTree) = children(tree) 49 | children(tree, node, ::IndexedTree) = (tree[y] for y in childindices(tree, node)) 50 | children(tree, node) = children(tree, node, treekind(tree)) 51 | 52 | childindices(tree, node) = 53 | tree == node ? childindices(tree, rootstate(tree)) : 54 | error("Must implement childindices(tree, node)") 55 | function childindices() 56 | end 57 | 58 | function parentind 59 | end 60 | 61 | """ 62 | nodetype(tree) 63 | 64 | A trait function, defined on the tree object, specifying the types of the nodes. 65 | The default is `Any`. When applicable, define this trait to make iteration inferrable. 66 | 67 | # Example 68 | 69 | ``` 70 | struct IntTree 71 | num::Int 72 | children::Vector{IntTree} 73 | end 74 | AbstractTrees.children(itree::IntTree) = itree.children 75 | AbstractTrees.nodetype(::IntTree) = IntTree 76 | ``` 77 | 78 | This suffices to make iteration over, e.g., `Leaves(itree::IntTree)` inferrable. 79 | """ 80 | nodetype(tree) = Any 81 | idxtype(tree) = Int 82 | -------------------------------------------------------------------------------- /src/indexing.jl: -------------------------------------------------------------------------------- 1 | # Wrappers which allow for tree-like indexing into objects 2 | 3 | 4 | struct Tree 5 | x::Any 6 | end 7 | Tree(x::Tree) = x 8 | Tree(x::AbstractShadowTree) = x 9 | show(io::IO, tree::Tree) = print_tree(io, tree.x) 10 | 11 | mutable struct AnnotationNode{T} 12 | val::T 13 | children::Vector{AnnotationNode{T}} 14 | end 15 | 16 | children(x::AnnotationNode) = x.children 17 | printnode(io::IO, x::AnnotationNode) = print(io, x.val) 18 | 19 | struct ShadowTree <: AbstractShadowTree 20 | tree::Tree 21 | shadow::Tree 22 | ShadowTree(x::Tree,y::Tree) = new(x,y) 23 | ShadowTree(x,y) = ShadowTree(Tree(x),Tree(y)) 24 | end 25 | first_tree(x::ShadowTree) = x.tree 26 | second_tree(x::ShadowTree) = x.shadow 27 | 28 | function zip_min(c1, c2) 29 | n1, n2 = length(c1), length(c2) 30 | if n1 < n2 31 | c2 = Iterators.take(c2,n1) 32 | elseif n2 < n1 33 | c1 = Iterators.take(c1,n2) 34 | end 35 | zip(c1, c2) 36 | end 37 | 38 | make_zip(x::AbstractShadowTree) = zip_min(children(x.tree.x), children(x.shadow.x)) 39 | 40 | function children(x::AbstractShadowTree) 41 | map(res->typeof(x)(res[1], res[2]),make_zip(x)) 42 | end 43 | 44 | iterate(x::AbstractShadowTree, state...) = iterate(make_zip(x), state...) 45 | 46 | function make_annotations(cb, tree, parent, s) 47 | s = cb(tree, parent, s) 48 | AnnotationNode{Any}(s, AnnotationNode{Any}[make_annotations(cb, child, tree, s) for child in children(tree)]) 49 | end 50 | 51 | function getindex(tree::Tree, indices) 52 | node = tree.x 53 | for idx in indices 54 | node = children(node)[idx] 55 | end 56 | node 57 | end 58 | getindex(tree::Tree, indices::T) where {T<:ImplicitNodeStack} = 59 | getindex(tree, indices.idx_stack.stack) 60 | 61 | 62 | function getindexhighest(tree::Tree, indices) 63 | node = tree.x 64 | for (i,idx) in enumerate(indices) 65 | cs = children(node) 66 | if idx > length(cs) 67 | return (indices[1:i-1],node) 68 | end 69 | node = children(node)[idx] 70 | end 71 | (indices, node) 72 | end 73 | 74 | function setindex!(tree::Tree, val, indices) 75 | setindex!(children(getindex(tree,indices[1:end-1])),val,indices[end]) 76 | end 77 | setindex!(tree::Tree, val, indices::T) where {T<:ImplicitNodeStack} = 78 | setindex!(tree, val, indices.idx_stack.stack) 79 | 80 | function getindex(tree::AbstractShadowTree, indices) 81 | typeof(tree)(Tree(first_tree(tree))[indices],Tree(second_tree(tree))[indices]) 82 | end 83 | 84 | function setindex!(tree::AbstractShadowTree, val, indices) 85 | setindex!(Tree(first_tree(tree)), first_tree(val), indices) 86 | setindex!(Tree(second_tree(tree)), second_tree(val), indices) 87 | end 88 | 89 | function setindex!(tree::AbstractShadowTree, val::Tuple, indices) 90 | setindex!(Tree(first_tree(tree)), val[1], indices) 91 | setindex!(Tree(second_tree(tree)), val[2], indices) 92 | end 93 | -------------------------------------------------------------------------------- /examples/binarytree_infer.jl: -------------------------------------------------------------------------------- 1 | # This file illustrates how to create inferrable tree-iteration methods in circumstances 2 | # where the children are not naturally indexable. 3 | # See "binarytree_easy.jl" for a simpler approach. 4 | 5 | if !isdefined(@__MODULE__, :BinaryNode) 6 | include("binarytree_core.jl") 7 | end 8 | 9 | ## Enhancement of the "native" binary tree 10 | # You might have the methods below even if you weren't trying to support AbstractTrees. 11 | 12 | # Implement iteration over the immediate children of a node 13 | function Base.iterate(node::BinaryNode) 14 | isdefined(node, :left) && return (node.left, false) 15 | isdefined(node, :right) && return (node.right, true) 16 | return nothing 17 | end 18 | function Base.iterate(node::BinaryNode, state::Bool) 19 | state && return nothing 20 | isdefined(node, :right) && return (node.right, true) 21 | return nothing 22 | end 23 | Base.IteratorSize(::Type{BinaryNode{T}}) where T = Base.SizeUnknown() 24 | Base.eltype(::Type{BinaryNode{T}}) where T = BinaryNode{T} 25 | 26 | ## Things we need to define to leverage the native iterator over children 27 | ## for the purposes of AbstractTrees. 28 | # Set the traits of this kind of tree 29 | Base.eltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = BinaryNode{T} 30 | Base.IteratorEltype(::Type{<:TreeIterator{BinaryNode{T}}}) where T = Base.HasEltype() 31 | AbstractTrees.parentlinks(::Type{BinaryNode{T}}) where T = AbstractTrees.StoredParents() 32 | AbstractTrees.siblinglinks(::Type{BinaryNode{T}}) where T = AbstractTrees.StoredSiblings() 33 | # Use the native iteration for the children 34 | AbstractTrees.children(node::BinaryNode) = node 35 | 36 | Base.parent(root::BinaryNode, node::BinaryNode) = isdefined(node, :parent) ? node.parent : nothing 37 | 38 | function AbstractTrees.nextsibling(tree::BinaryNode, child::BinaryNode) 39 | isdefined(child, :parent) || return nothing 40 | p = child.parent 41 | if isdefined(p, :right) 42 | child === p.right && return nothing 43 | return p.right 44 | end 45 | return nothing 46 | end 47 | 48 | # We also need `pairs` to return something sensible. 49 | # If you don't like integer keys, you could do, e.g., 50 | # Base.pairs(node::BinaryNode) = BinaryNodePairs(node) 51 | # and have its iteration return, e.g., `:left=>node.left` and `:right=>node.right` when defined. 52 | # But the following is easy: 53 | Base.pairs(node::BinaryNode) = enumerate(node) 54 | 55 | 56 | AbstractTrees.printnode(io::IO, node::BinaryNode) = print(io, node.data) 57 | 58 | root = BinaryNode(0) 59 | l = leftchild(1, root) 60 | r = rightchild(2, root) 61 | lr = rightchild(3, l) 62 | 63 | print_tree(root) 64 | collect(PostOrderDFS(root)) 65 | @static if isdefined(@__MODULE__, :Test) 66 | @testset "binarytree_infer.jl" begin 67 | @test @inferred(map(x->x.data, PostOrderDFS(root))) == [3, 1, 2, 0] 68 | @test @inferred(map(x->x.data, PreOrderDFS(root))) == [0, 1, 3, 2] 69 | @test @inferred(map(x->x.data, Leaves(root))) == [3, 2] 70 | end 71 | end 72 | -------------------------------------------------------------------------------- /src/implicitstacks.jl: -------------------------------------------------------------------------------- 1 | abstract type ImplicitStack end 2 | struct ImplicitIndexStack{S} <: ImplicitStack 3 | stack::Vector{S} 4 | end 5 | Base.copy(s::ImplicitIndexStack) = typeof(s)(copy(s.stack)) 6 | """ 7 | Keeps a stack of nodes and their corresponding indices. Note that the last node 8 | is not explicitly stored in the node_stack, such that length(node_stack) == 9 | length(idx_stack)-1 (unless we're at the root in which case both are empty) 10 | """ 11 | struct ImplicitNodeStack{T, S} <: ImplicitStack 12 | node_stack::Vector{T} 13 | idx_stack::ImplicitIndexStack{S} 14 | end 15 | Base.copy(s::ImplicitNodeStack) = typeof(s)(copy(s.node_stack), copy(s.idx_stack)) 16 | Base.isempty(s::ImplicitNodeStack) = isempty(s.idx_stack) 17 | Base.isempty(s::ImplicitIndexStack) = isempty(s.stack) 18 | getnode(tree, ns::ImplicitIndexStack) = isempty(ns.stack) ? tree[rootstate(tree)] : 19 | (@assert isa(treekind(tree), IndexedTree); tree[ns.stack[end]]) 20 | function getnode(tree, stack::ImplicitNodeStack) 21 | isempty(stack.node_stack) ? 22 | (isempty(stack.idx_stack) ? tree : children(tree)[stack.idx_stack.stack[end]]) : 23 | children(stack.node_stack[end])[stack.idx_stack.stack[end]] 24 | end 25 | struct ImplicitChildStates{T, S} 26 | tree::T 27 | stack::S 28 | end 29 | Base.IteratorSize(::Type{T}) where {T<:ImplicitChildStates} = Base.SizeUnknown() 30 | children(states::ImplicitChildStates) = children(states.tree, states.stack) 31 | 32 | parentstate(tree, state::ImplicitNodeStack) = 33 | ImplicitNodeStack(state.node_stack[1:end-1], parentstate(tree, state.idx_stack)) 34 | parentstate(tree, state::ImplicitIndexStack) = ImplicitIndexStack(state.stack[1:end-1]) 35 | childstates(tree, state) = childstates(tree, state, treekind(tree)) 36 | childstates(tree, state::ImplicitStack) = ImplicitChildStates(tree, state) 37 | children(tree, state::ImplicitNodeStack) = children(tree, getnode(tree, state)) 38 | children(tree, state::ImplicitIndexStack) = isempty(state) ? 39 | children(tree) : children(tree, tree[state.stack[end]]) 40 | 41 | childstates(s::ImplicitChildStates) = isempty(s.stack) ? 42 | childstates(s.tree, s.tree) : childstates(s.tree, 43 | isa(s.stack, ImplicitNodeStack) ? getnode(s.tree, s.stack) : s.stack.stack[end]) 44 | nextind(s::ImplicitChildStates, ind) = nextind(childstates(s), ind) 45 | function iterate(s::ImplicitChildStates, ind=1) 46 | y = iterate(childstates(s), ind) 47 | y === nothing && return nothing 48 | _, ni = y 49 | (update_state!(s.tree, copy(s.stack), childstates(s), ind, treekind(s.tree)), ni) 50 | end 51 | 52 | update_state!(tree, ns, cs, idx, _) = update_state!(tree, ns, cs, idx) 53 | update_state!(tree, ns::ImplicitIndexStack, cs, idx) = (push!(ns.stack, idx); ns) 54 | update_state!(tree, ns::ImplicitIndexStack, cs, idx, ::IndexedTree) = (push!(ns.stack, iterate(cs, idx)[1]); ns) 55 | function update_state!(tree, ns::ImplicitNodeStack, cs, idx) 56 | !isempty(ns) && push!(ns.node_stack, getnode(tree, ns)) 57 | update_state!(tree, ns.idx_stack, cs, idx) 58 | ns 59 | end 60 | 61 | function joinstate(tree, state::ImplicitNodeStack, new_state::ImplicitNodeStack) 62 | isempty(new_state) && return state 63 | ImplicitNodeStack([push!(copy(state.node_stack), getnode(tree, state)); new_state.node_stack], 64 | joinstate(tree, state.idx_stack, new_state.idx_stack)) 65 | end 66 | joinstate(tree, state::ImplicitIndexStack, new_state::ImplicitIndexStack) = 67 | ImplicitIndexStack([state.stack; new_state.stack]) 68 | 69 | isroot(tree, state::ImplicitStack) = isempty(state) 70 | -------------------------------------------------------------------------------- /src/printing.jl: -------------------------------------------------------------------------------- 1 | """ 2 | print_tree(tree, maxdepth=5; kwargs...) 3 | print_tree(io, tree, maxdepth=5; kwargs...) 4 | print_tree(f::Function, io, tree, maxdepth=5; kwargs...) 5 | 6 | # Usage 7 | Prints an ASCII formatted representation of the `tree` to the given `io` object. 8 | By default all children will be printed up to a maximum level of 5, though this 9 | value can be overriden by the `maxdepth` parameter. Nodes that are truncated are 10 | indicated by a vertical ellipsis below the truncated node, this indication can be 11 | turned off by providing `indicate_truncation=false` as a kwarg. The charset to use in 12 | printing can be customized using the `charset` keyword argument. 13 | You can control the printing of individual nodes by passing a function `f(io, node)`; 14 | the default is [`AbstractTrees.printnode`](@ref). 15 | 16 | # Examples 17 | ```julia 18 | julia> print_tree(stdout, Dict("a"=>"b","b"=>['c','d'])) 19 | Dict{String,Any}("b"=>['c','d'],"a"=>"b") 20 | ├─ b 21 | │ ├─ c 22 | │ └─ d 23 | └─ a 24 | └─ b 25 | 26 | julia> print_tree(stdout, '0'=>'1'=>'2'=>'3', 2) 27 | '0' 28 | └─ '1' 29 | └─ '2' 30 | ⋮ 31 | 32 | julia> print_tree(stdout, Dict("a"=>"b","b"=>['c','d']); 33 | charset = TreeCharSet('+','\\','|',"--","⋮")) 34 | Dict{String,Any}("b"=>['c','d'],"a"=>"b") 35 | +-- b 36 | | +-- c 37 | | \\-- d 38 | \\-- a 39 | \\-- b 40 | ``` 41 | 42 | """ 43 | print_tree 44 | 45 | 46 | """ 47 | printnode(io::IO, node) 48 | 49 | Print a single node. The default is to show a compact representation of `node`. 50 | Override this if you want nodes printed in a custom way in [`print_tree`](@ref), 51 | or if you want your print function to print part of the tree by default. 52 | 53 | # Example 54 | 55 | ``` 56 | struct MyNode{T} 57 | data::T 58 | children::Vector{MyNode{T}} 59 | end 60 | AbstractTrees.printnode(io::IO, node::MyNode) = print(io, "MyNode(\$(node.data))") 61 | ``` 62 | """ 63 | printnode(io::IO, node) = show(IOContext(io, :compact => true), node) 64 | 65 | 66 | struct TreeCharSet 67 | mid 68 | terminator 69 | skip 70 | dash 71 | trunc 72 | end 73 | 74 | # Default charset 75 | TreeCharSet() = TreeCharSet('├','└','│','─','⋮') 76 | TreeCharSet(mid, term, skip, dash) = TreeCharSet(mid, term, skip, dash, '⋮') 77 | 78 | 79 | function print_prefix(io, depth, charset, active_levels) 80 | for current_depth in 0:(depth-1) 81 | if current_depth in active_levels 82 | print(io,charset.skip," "^(textwidth(charset.dash)+1)) 83 | else 84 | print(io," "^(textwidth(charset.skip)+textwidth(charset.dash)+1)) 85 | end 86 | end 87 | end 88 | 89 | function _print_tree(printnode::Function, io::IO, tree, maxdepth = 5; indicate_truncation = true, 90 | depth = 0, active_levels = Int[], charset = TreeCharSet(), withinds = false, 91 | inds = [], from = nothing, to = nothing, roottree = tree) 92 | nodebuf = IOBuffer() 93 | isa(io, IOContext) && (nodebuf = IOContext(nodebuf, io)) 94 | if withinds 95 | printnode(nodebuf, tree, inds) 96 | else 97 | tree != roottree && isa(treekind(roottree), IndexedTree) ? 98 | printnode(nodebuf, roottree[tree]) : 99 | printnode(nodebuf, tree) 100 | end 101 | str = String(take!(isa(nodebuf, IOContext) ? nodebuf.io : nodebuf)) 102 | for (i,line) in enumerate(split(str, '\n')) 103 | i != 1 && print_prefix(io, depth, charset, active_levels) 104 | println(io, line) 105 | end 106 | c = isa(treekind(roottree), IndexedTree) ? 107 | childindices(roottree, tree) : children(roottree, tree) 108 | if c !== () 109 | if depth < maxdepth 110 | s = Iterators.Stateful(from === nothing ? pairs(c) : Iterators.Rest(pairs(c), from)) 111 | while !isempty(s) 112 | ind, child = popfirst!(s) 113 | ind === to && break 114 | active = false 115 | child_active_levels = active_levels 116 | print_prefix(io, depth, charset, active_levels) 117 | if isempty(s) 118 | print(io, charset.terminator) 119 | else 120 | print(io, charset.mid) 121 | child_active_levels = push!(copy(active_levels), depth) 122 | end 123 | print(io, charset.dash, ' ') 124 | print_tree(printnode, io, child, maxdepth; 125 | indicate_truncation=indicate_truncation, depth = depth + 1, 126 | active_levels = child_active_levels, charset = charset, withinds=withinds, 127 | inds = withinds ? [inds; ind] : [], roottree = roottree) 128 | end 129 | elseif indicate_truncation 130 | print_prefix(io, depth, charset, active_levels) 131 | println(io, charset.trunc) 132 | print_prefix(io, depth, charset, active_levels) 133 | println(io) 134 | end 135 | end 136 | end 137 | print_tree(f::Function, io::IO, tree, args...; kwargs...) = _print_tree(f, io, tree, args...; kwargs...) 138 | print_tree(io::IO, tree, args...; kwargs...) = print_tree(printnode, io, tree, args...; kwargs...) 139 | print_tree(tree, args...; kwargs...) = print_tree(stdout::IO, tree, args...; kwargs...) 140 | -------------------------------------------------------------------------------- /test/runtests.jl: -------------------------------------------------------------------------------- 1 | using AbstractTrees 2 | using Test 3 | import Base: == 4 | 5 | 6 | if VERSION >= v"1.1.0-DEV.838" # requires https://github.com/JuliaLang/julia/pull/30291 7 | @testset "Ambiguities" begin 8 | @test isempty(detect_ambiguities(AbstractTrees, Base, Core)) 9 | end 10 | end 11 | 12 | 13 | @testset "Array" begin 14 | tree = Any[1,Any[2,3]] 15 | io = IOBuffer() 16 | print_tree(io, tree) 17 | @test String(take!(io)) == "Array{Any,1}\n├─ 1\n└─ Array{Any,1}\n ├─ 2\n └─ 3\n" 18 | @test collect(Leaves(tree)) == [1,2,3] 19 | @test collect(Leaves(tree)) isa Vector{Int} 20 | @test collect(PostOrderDFS(tree)) == Any[1,2,3,Any[2,3],Any[1,Any[2,3]]] 21 | @test collect(StatelessBFS(tree)) == Any[Any[1,Any[2,3]],1,Any[2,3],2,3] 22 | 23 | tree2 = Any[Any[1,2],Any[3,4]] 24 | @test collect(PreOrderDFS(tree2)) == Any[tree2,Any[1,2],1,2,Any[3,4],3,4] 25 | end 26 | 27 | 28 | """ 29 | A tree in which every node has 0 or 1 children 30 | """ 31 | struct OneTree 32 | nodes::Vector{Int} 33 | end 34 | AbstractTrees.treekind(::Type{OneTree}) = AbstractTrees.IndexedTree() 35 | AbstractTrees.siblinglinks(::Type{OneTree}) = AbstractTrees.StoredSiblings() 36 | AbstractTrees.relative_state(t::OneTree, _, __::Int) = 1 37 | Base.getindex(t::OneTree, idx) = t.nodes[idx] 38 | AbstractTrees.childindices(tree::OneTree, node::Int) = 39 | (ret = (node == 0 || tree[node] == 0) ? () : (tree[node],)) 40 | AbstractTrees.children(tree::OneTree) = AbstractTrees.children(tree, tree) 41 | AbstractTrees.rootstate(tree::OneTree) = 1 42 | AbstractTrees.printnode(io::IO, t::OneTree) = 43 | AbstractTrees.printnode(io::IO, t[AbstractTrees.rootstate(t)]) 44 | Base.eltype(::Type{<:TreeIterator{OneTree}}) = Int 45 | Base.IteratorEltype(::Type{<:TreeIterator{OneTree}}) = Base.HasEltype() 46 | 47 | @testset "OneTree" begin 48 | ot = OneTree([2,3,4,0]) 49 | io = IOBuffer() 50 | print_tree(io, ot) 51 | @test String(take!(io)) == "2\n└─ 3\n └─ 4\n └─ 0\n" 52 | @test @inferred(collect(Leaves(ot))) == [0] 53 | @test eltype(collect(Leaves(ot))) === Int 54 | @test collect(PreOrderDFS(ot)) == [2,3,4,0] 55 | @test collect(PostOrderDFS(ot)) == [0,4,3,2] 56 | end 57 | 58 | 59 | """ 60 | Stores an explicit parent for some other kind of tree 61 | """ 62 | struct ParentTree{T} 63 | tree::T 64 | parents::Vector{Int} 65 | end 66 | AbstractTrees.treekind(::Type{ParentTree{T}}) where {T} = AbstractTrees.treekind(T) 67 | AbstractTrees.parentlinks(::Type{ParentTree{T}}) where {T} = AbstractTrees.StoredParents() 68 | AbstractTrees.siblinglinks(::Type{ParentTree{T}}) where {T} = AbstractTrees.siblinglinks(T) 69 | AbstractTrees.relative_state(t::ParentTree, x, r::Int) = 70 | AbstractTrees.relative_state(t.tree, x, r) 71 | Base.getindex(t::ParentTree, idx) = t.tree[idx] 72 | AbstractTrees.childindices(tree::ParentTree, node::Int) = AbstractTrees.childindices(tree.tree, node) 73 | AbstractTrees.children(tree::ParentTree) = AbstractTrees.children(tree, tree) 74 | AbstractTrees.rootstate(tree::ParentTree) = AbstractTrees.rootstate(tree.tree) 75 | AbstractTrees.parentind(tree::ParentTree, node::Int) = tree.parents[node] 76 | AbstractTrees.printnode(io::IO, t::ParentTree) = 77 | AbstractTrees.printnode(io::IO, t[AbstractTrees.rootstate(t)]) 78 | 79 | @testset "ParentTree" begin 80 | ot = OneTree([2,3,4,0]) 81 | pt = ParentTree(ot,[0,1,2,3]) 82 | io = IOBuffer() 83 | print_tree(io, pt) 84 | @test String(take!(io)) == "2\n└─ 3\n └─ 4\n └─ 0\n" 85 | @test collect(Leaves(pt)) == [0] 86 | @test collect(PreOrderDFS(pt)) == [2,3,4,0] 87 | @test collect(PostOrderDFS(pt)) == [0,4,3,2] 88 | 89 | # Test modification while iterating over PreOrderDFS 90 | a = [1,[2,[3]]] 91 | b = treemap!(PreOrderDFS(a)) do node 92 | !isa(node, Vector) && return node 93 | ret = pushfirst!(copy(node),0) 94 | # And just for good measure stomp over the old node to make sure nothing 95 | # is cached. 96 | empty!(node) 97 | ret 98 | end 99 | @test b == Any[0,1,Any[0,2,[0,3]]] 100 | end 101 | 102 | 103 | struct IntTree 104 | num::Int 105 | children::Vector{IntTree} 106 | end 107 | ==(x::IntTree,y::IntTree) = x.num == y.num && x.children == y.children 108 | AbstractTrees.children(itree::IntTree) = itree.children 109 | Base.eltype(::Type{<:TreeIterator{IntTree}}) = IntTree 110 | Base.IteratorEltype(::Type{<:TreeIterator{IntTree}}) = Base.HasEltype() 111 | AbstractTrees.nodetype(::IntTree) = IntTree 112 | 113 | @testset "IntTree" begin 114 | itree = IntTree(1, [IntTree(2, IntTree[])]) 115 | iter = Leaves(itree) 116 | @test @inferred(first(iter)) == IntTree(2, IntTree[]) 117 | val, state = iterate(iter) 118 | @test Base.return_types(iterate, Tuple{typeof(iter), typeof(state)}) == 119 | [Union{Nothing, Tuple{IntTree,typeof(state)}}] 120 | end 121 | 122 | #= 123 | @test treemap(PostOrderDFS(tree)) do ind, x, children 124 | IntTree(isa(x,Int) ? x : mapreduce(x->x.num,+,0,children), 125 | isempty(children) ? IntTree[] : children) 126 | end == IntTree(6,[IntTree(1,IntTree[]),IntTree(5,[IntTree(2,IntTree[]),IntTree(3,IntTree[])])]) 127 | =# 128 | 129 | @test collect(PostOrderDFS([])) == Any[[]] 130 | 131 | @testset "Examples" begin 132 | # Ensure the examples run 133 | exampledir = joinpath(dirname(@__DIR__), "examples") 134 | examples = readdir(exampledir) 135 | mktemp() do filename, io 136 | redirect_stdout(io) do 137 | for ex in examples 138 | haskey(ENV, "CI") && Sys.isapple() && ex == "fstree.jl" && continue 139 | include(joinpath(exampledir, ex)) 140 | end 141 | end 142 | end 143 | end 144 | 145 | 146 | struct Num{I} end 147 | Num(I::Int) = Num{I}() 148 | Base.show(io::IO, ::Num{I}) where {I} = print(io, I) 149 | AbstractTrees.children(x::Num{I}) where {I} = (Num(I+1), Num(I+1)) 150 | 151 | struct SingleChildInfiniteDepth end 152 | AbstractTrees.children(::SingleChildInfiniteDepth) = (SingleChildInfiniteDepth(),) 153 | 154 | @testset "Test print_tree truncation" begin 155 | 156 | # test that `print_tree(headnode, maxdepth)` truncates the output at right depth 157 | # julia > print_tree(Num(0), 3) 158 | # 0 159 | # ├─ 1 160 | # │ ├─ 2 161 | # │ │ ├─ 3 162 | # │ │ └─ 3 163 | # │ └─ 2 164 | # │ ├─ 3 165 | # │ └─ 3 166 | # └─ 1 167 | # ├─ 2 168 | # │ ├─ 3 169 | # │ └─ 3 170 | # └─ 2 171 | # ├─ 3 172 | # └─ 3 173 | # 174 | 175 | for maxdepth in [3,5,8] 176 | buffer = IOBuffer() 177 | print_tree(buffer, Num(0), maxdepth) 178 | ptxt = String(take!(buffer)) 179 | n1 = sum([1 for c in ptxt if c=="$(maxdepth-1)"[1]]) 180 | n2 = sum([1 for c in ptxt if c=="$maxdepth"[1]]) 181 | n3 = sum([1 for c in ptxt if c=="$(maxdepth+1)"[1]]) 182 | @test n1==2^(maxdepth-1) 183 | @test n2==2^maxdepth 184 | @test n3==0 185 | end 186 | 187 | # test that `print_tree(headnode)` prints truncation characters under each 188 | # node at the default maxdepth level = 5 189 | truncation_char = AbstractTrees.TreeCharSet().trunc 190 | buffer = IOBuffer() 191 | print_tree(buffer, Num(0)) 192 | ptxt = String(take!(buffer)) 193 | n1 = sum([1 for c in ptxt if c=='5']) 194 | n2 = sum([1 for c in ptxt if c=='6']) 195 | n3 = sum([1 for c in ptxt if c==truncation_char]) 196 | @test n1==2^5 197 | @test n2==0 198 | @test n3==2^5 199 | lines = split(ptxt, '\n') 200 | for i in 1:length(lines) 201 | if ~isempty(lines[i]) && lines[i][end] == '5' 202 | @test lines[i+1][end] == truncation_char 203 | end 204 | end 205 | 206 | # test correct number of lines printed 1 207 | buffer = IOBuffer() 208 | print_tree(buffer, SingleChildInfiniteDepth()) 209 | ptxt = String(take!(buffer)) 210 | numlines = sum([1 for c in split(ptxt, '\n') if ~isempty(strip(c))]) 211 | @test numlines == 7 # 1 (head node) + 5 (default depth) + 1 (truncation char) 212 | 213 | # test correct number of lines printed 2 214 | buffer = IOBuffer() 215 | print_tree(buffer, SingleChildInfiniteDepth(), 3) 216 | ptxt = String(take!(buffer)) 217 | numlines = sum([1 for c in split(ptxt, '\n') if ~isempty(strip(c))]) 218 | @test numlines == 5 # 1 (head node) + 3 (depth) + 1 (truncation char) 219 | 220 | # test correct number of lines printed 3 221 | buffer = IOBuffer() 222 | print_tree(buffer, SingleChildInfiniteDepth(), 3, indicate_truncation=false) 223 | ptxt = String(take!(buffer)) 224 | numlines = sum([1 for c in split(ptxt, '\n') if ~isempty(strip(c))]) 225 | @test numlines == 4 # 1 (head node) + 3 (depth) 226 | end 227 | -------------------------------------------------------------------------------- /src/iteration.jl: -------------------------------------------------------------------------------- 1 | # Utilities for tree traversal and iteration 2 | 3 | 4 | abstract type TreeIterator{T} end 5 | IteratorEltype(::Type{<:TreeIterator}) = EltypeUnknown() 6 | 7 | """ 8 | Iterator to visit the leaves of a tree, e.g. for the tree 9 | 10 | ``` 11 | Any[1,Any[2,3]] 12 | ├─ 1 13 | └─ Any[2,3] 14 | ├─ 2 15 | └─ 3 16 | ``` 17 | 18 | we will get `[1,2,3]`. 19 | """ 20 | struct Leaves{T} <: TreeIterator{T} 21 | tree::T 22 | end 23 | IteratorSize(::Type{Leaves{T}}) where {T} = SizeUnknown() 24 | 25 | """ 26 | Iterator to visit the nodes of a tree, guaranteeing that children 27 | will be visited before their parents. 28 | 29 | e.g. for the tree 30 | 31 | ``` 32 | Any[1,Any[2,3]] 33 | ├─ 1 34 | └─ Any[2,3] 35 | ├─ 2 36 | └─ 3 37 | ``` 38 | 39 | we will get `[1, 2, 3, [2, 3], [1, [2, 3]]]`. 40 | """ 41 | struct PostOrderDFS{T} <: TreeIterator{T} 42 | tree::T 43 | end 44 | PostOrderDFS(tree::Tree) = PostOrderDFS(tree.x) 45 | IteratorSize(::Type{PostOrderDFS{T}}) where T = SizeUnknown() 46 | 47 | """ 48 | Iterator to visit the nodes of a tree, guaranteeing that parents 49 | will be visited before their children. 50 | 51 | Optionally takes a filter function that determines whether the iterator 52 | should continue iterating over a node's children (if it has any) or should 53 | consider that node a leaf. 54 | 55 | e.g. for the tree 56 | 57 | ``` 58 | Any[Any[1,2],Any[3,4]] 59 | ├─ Any[1,2] 60 | | ├─ 1 61 | | └─ 2 62 | └─ Any[3,4] 63 | ├─ 3 64 | └─ 4 65 | ``` 66 | 67 | we will get `[[[1, 2], [3, 4]], [1, 2], 1, 2, [3, 4], 3, 4]`. 68 | 69 | # Invalidation 70 | Modifying the underlying tree while iterating over it, is allowed, however, 71 | if parents and sibling links are not explicitly stored, the identify of any 72 | parent of the last obtained node does not change (i.e. mutation is allowed, 73 | replacing nodes is not). 74 | 75 | """ 76 | struct PreOrderDFS{T} <: TreeIterator{T} 77 | tree::T 78 | filter::Function 79 | PreOrderDFS{T}(tree,filter::Function=(args...)->true) where {T} = new{T}(tree,filter) 80 | end 81 | PreOrderDFS(tree::T,filter::Function=(args...)->true) where {T} = PreOrderDFS{T}(tree,filter) 82 | PreOrderDFS(tree::Tree,filter::Function=(args...)->true) = PreOrderDFS(tree.x,filter) 83 | IteratorSize(::Type{PreOrderDFS{T}}) where {T} = SizeUnknown() 84 | 85 | # State depends on what kind of tree we have: 86 | # - Parents/Siblings are not stored: 87 | # - RegularTree: ImplicitNodeStack 88 | # - IndexedTree: ImplicitIndexStack 89 | # - Parents/Siblings are stored: 90 | # - RegularTree: Nodes 91 | # - IndexedTree: Indices 92 | # 93 | childstates(tree, state, ::IndexedTree) = childindices(tree, state) 94 | childstates(tree, state, ::RegularTree) = children(tree, state) 95 | parentstate(tree, state, ::IndexedTree) = parentind(tree, state) 96 | parentstate(tree, state, ::RegularTree) = parent(tree, state) 97 | 98 | parentstate(tree, state) = parentstate(tree, state, treekind(tree)) 99 | 100 | update_state!(old_state, cs, idx) = next(cs, idx)[1] 101 | 102 | 103 | getindex(x::AbstractArray, ::ImplicitRootState) = x 104 | 105 | function firststate(ti::PreOrderDFS{T}) where T 106 | if isa(parentlinks(ti.tree), StoredParents) && 107 | isa(siblinglinks(ti.tree), SiblingLinks) 108 | rootstate(ti.tree) 109 | else 110 | state = ImplicitIndexStack(idxtype(ti.tree)[]) 111 | if !isa(treekind(typeof(ti.tree)), IndexedTree) 112 | state = ImplicitNodeStack(nodetype(ti.tree)[], state) 113 | end 114 | state 115 | end 116 | end 117 | function firststate(ti::Union{Leaves, PostOrderDFS}) 118 | state = firststate(PreOrderDFS(ti.tree)) 119 | while true 120 | css = childstates(ti.tree, state) 121 | isempty(css) && break 122 | state = first(css) 123 | end 124 | state 125 | end 126 | 127 | nextind(::Base.Generator, idx) = idx + 1 128 | relative_state(tree, parentstate, childstate::ImplicitIndexStack) = 129 | childstate.stack[end] 130 | relative_state(tree, parentstate, childstate::ImplicitNodeStack) = 131 | relative_state(tree, parentstate, childstate.idx_stack) 132 | 133 | function nextsibling(tree, state) 134 | ps = parentstate(tree, state) 135 | cs = childstates(tree, ps) 136 | isempty(cs) && return nothing 137 | new_state = nextind(cs, relative_state(tree, ps, state)) 138 | iterate(cs, new_state) === nothing && return nothing 139 | update_state!(tree, ps, children(tree, ps), new_state) 140 | end 141 | 142 | function nextsibling(node, ::StoredParents, ::ImplicitSiblings, ::RegularTree) 143 | isroot(node) && return nothing 144 | p = parent(node) 145 | last_was_node = false 146 | for c in children(p) 147 | last_was_node && return c 148 | (c == node) && (last_was_node = true) 149 | end 150 | last_was_node && return nothing 151 | error("Tree inconsistency: node not a child of parent") 152 | end 153 | nextsibling(node, ::Any, ::StoredSiblings, ::Any) = error("Trees with explicit siblings must override the `nextsibling` method explicitly") 154 | nextsibling(node) = nextsibling(node, parentlinks(node), siblinglinks(node), treekind(node)) 155 | 156 | function prevsibling(node, ::StoredParents, ::ImplicitSiblings, ::RegularTree) 157 | isroot(node) && return nothing 158 | p = parent(node) 159 | last_c = nothing 160 | for c in children(p) 161 | (c == node) && return last_c 162 | last_c = c 163 | end 164 | @show p 165 | @show node 166 | error("Tree inconsistency: node not a child of parent") 167 | end 168 | prevsibling(node, ::Any, ::StoredSiblings, ::Any) = error("Trees with explicit siblings must override the `prevsibling` method explicitly") 169 | prevsibling(node) = prevsibling(node, parentlinks(node), siblinglinks(node), treekind(node)) 170 | prevsibling(tree, node) = prevsibling(node) 171 | 172 | isroot(tree, state, ::RegularTree) = tree == state 173 | isroot(tree, state, ::IndexedTree) = state == rootstate(tree) 174 | isroot(tree, state) = isroot(tree, state, treekind(tree)) 175 | 176 | struct Subtree{T,S} 177 | tree::T 178 | state::S 179 | end 180 | children(tree::Subtree) = children(tree.tree, tree.state) 181 | nodetype(tree::Subtree) = nodetype(tree.tree) 182 | idxtype(tree::Subtree) = idxtype(tree.tree) 183 | rootstate(tree::Subtree) = tree.state 184 | parentlinks(::Type{Subtree{T,S}}) where {T,S} = parentlinks(T) 185 | 186 | joinstate(tree, a, b) = b 187 | 188 | if isdefined(Base, :UnionAll) 189 | Base.@pure function get_primary(T::DataType) 190 | T.name.wrapper 191 | end 192 | else 193 | Base.@pure function get_primary(T::DataType) 194 | T.name.primary 195 | end 196 | end 197 | 198 | function stepstate(ti::TreeIterator, state) 199 | if isa(ti, PreOrderDFS) && ti.filter(getnode(ti.tree, state)) 200 | ccs = childstates(ti.tree, state) 201 | !isempty(ccs) && return first(ccs) 202 | end 203 | while !isroot(ti.tree, state) 204 | nextstate = nextsibling(ti.tree, state) 205 | if nextstate !== nothing 206 | return joinstate(ti.tree, nextstate, firststate( 207 | get_primary(typeof(ti))(Subtree(ti.tree, nextstate)))) 208 | end 209 | state = parentstate(ti.tree, state) 210 | isa(ti, PostOrderDFS) && return state 211 | end 212 | nothing 213 | end 214 | 215 | getnode(tree::AbstractShadowTree, ns::ImplicitNodeStack) = tree[ns.idx_stack.stack] 216 | getnode(tree, ns) = getnode(tree, ns, treekind(tree)) 217 | getnode(tree, ns, ::IndexedTree) = tree[ns] 218 | getnode(tree, ns, ::RegularTree) = ns 219 | getnode(tree, ::ImplicitRootState, ::RegularTree) = tree 220 | 221 | function iterate(ti::TreeIterator) 222 | state = firststate(ti) 223 | (getnode(ti.tree, state), state) 224 | end 225 | function iterate(ti::TreeIterator, state) 226 | state = stepstate(ti, state) 227 | state === nothing && return nothing 228 | (getnode(ti.tree, state), state) 229 | end 230 | 231 | """ 232 | Acends the tree, at each node choosing whether or not to continue. 233 | Note that the parent is computed before the callback is exectuted, allowing 234 | modification of the argument to the callback (as long as the overall tree 235 | structure is not altered). 236 | """ 237 | function ascend(select, node) 238 | isroot(node) && (select(node); return node) 239 | p = parent(node) 240 | while select(node) && !isroot(node) 241 | node = p 242 | p = parent(node) 243 | end 244 | node 245 | end 246 | 247 | """ 248 | Descends the tree, at each node choosing the child given by select callback 249 | or the current node if 0 is returned. 250 | """ 251 | function descend(select, tree) 252 | idx = select(tree) 253 | idx == 0 && return tree 254 | node = children(tree)[idx] 255 | while true 256 | idx = select(node) 257 | idx == 0 && return node 258 | node = children(node)[idx] 259 | end 260 | end 261 | 262 | """ 263 | Iterator to visit the nodes of a tree, all nodes of a level will be visited 264 | before their children 265 | 266 | e.g. for the tree 267 | 268 | ``` 269 | Any[1,Any[2,3]] 270 | ├─ 1 271 | └─ Any[2,3] 272 | ├─ 2 273 | └─ 3 274 | ``` 275 | 276 | we will get `[[1, [2,3]], 1, [2, 3], 2, 3]`. 277 | 278 | WARNING: This is \$O(n^2)\$, only use this if you know you need it, as opposed to 279 | a more standard statefull approach. 280 | """ 281 | struct StatelessBFS <: TreeIterator{Any} 282 | tree::Any 283 | end 284 | IteratorSize(::Type{StatelessBFS}) = SizeUnknown() 285 | 286 | function descend_left(newinds, next_node, level) 287 | # Go down until we are at the correct level or a dead end 288 | while length(newinds) != level 289 | cs = children(next_node) 290 | if isempty(cs) 291 | break 292 | end 293 | push!(newinds, 1) 294 | next_node = first(cs) 295 | end 296 | return newinds 297 | end 298 | 299 | function nextind_or_deadend(tree, ind, level) 300 | current_level = active_level = length(ind) 301 | active_inds = copy(ind) 302 | # Go up until there is a right neighbor 303 | while current_level > 0 304 | # Check for next node at the current level 305 | active_inds = ind[1:current_level-1] 306 | parent = Tree(tree)[active_inds] 307 | cur_child = ind[current_level] 308 | ni = nextind(children(parent), cur_child) 309 | current_level -= 1 310 | if iterate(children(parent), ni) !== nothing 311 | newinds = [active_inds; ni] 312 | next_node = children(parent)[ni] 313 | return descend_left(newinds, next_node, level) 314 | end 315 | end 316 | return nothing 317 | end 318 | 319 | iterate(ti::StatelessBFS) = (Tree(ti.tree)[[]], []) 320 | """ 321 | Stateless level-order bfs iteration. The algorithm is as follows: 322 | 323 | Go up. If there is a right neighbor, go right, then left until you reach the 324 | same level. If you reach the root, go left until you reach the next level. 325 | """ 326 | function iterate(ti::StatelessBFS, ind) 327 | org_level = active_level = length(ind) 328 | newinds = ind 329 | while true 330 | newinds = nextind_or_deadend(ti.tree, newinds, active_level) 331 | if newinds === nothing 332 | active_level += 1 333 | if active_level > org_level + 1 334 | return nothing 335 | end 336 | newinds = descend_left([], ti.tree, active_level) 337 | end 338 | if length(newinds) == active_level 339 | break 340 | end 341 | end 342 | Tree(ti.tree)[newinds], newinds 343 | end 344 | 345 | # Mapping over trees 346 | function treemap(f::Function, tree::PostOrderDFS) 347 | new_tree = Any[Union{}[]] 348 | current_length = 0 349 | for (ind, node) in pairs(tree) 350 | while length(new_tree) < length(ind) 351 | push!(new_tree, Union{}[]) 352 | end 353 | thechildren = Union{}[] 354 | if length(ind) < length(new_tree) 355 | thechildren = pop!(new_tree) 356 | end 357 | if ind == [] 358 | return f(ind, node, thechildren) 359 | end 360 | siblings = new_tree[end] 361 | el = f(ind, node, thechildren) 362 | S = typeof(el) 363 | T = eltype(siblings) 364 | if S === T || S <: T 365 | push!(siblings, el) 366 | else 367 | R = typejoin(T, S) 368 | new = similar(siblings, R) 369 | copy!(new,1,siblings,1,length(siblings)) 370 | push!(new,el) 371 | new_tree[end] = new 372 | end 373 | end 374 | end 375 | 376 | function treemap!(f::Function, ti::PreOrderDFS) 377 | state = firststate(ti) 378 | while state !== nothing 379 | ind = state 380 | node = getnode(ti.tree, ind) 381 | new_node = f(node) 382 | if new_node !== node 383 | if isempty(ind) 384 | return treemap!(PreOrderDFS(new_node)) do x 385 | x == new_node && return x 386 | f(x) 387 | end 388 | end 389 | Tree(ti.tree)[ind] = new_node 390 | end 391 | state = stepstate(ti, ind) 392 | end 393 | ti.tree 394 | end 395 | --------------------------------------------------------------------------------