module LeftChildRightSiblingTrees

using AbstractTrees

# See `abstracttrees.jl` for the only dependency of this package

export Node,
    addchild,
    addsibling,
    depth,
    graftchildren!,
    isroot,
    isleaf,
    islastsibling,
    isequiv,
    lastsibling,
    prunebranch!,
    copy_subtree

mutable struct Node{T}
    data::T
    parent::Node{T}
    child::Node{T}
    sibling::Node{T}

    """
        root = Node(data)

    Construct a disconnected node, which can serve as the root of a new tree.
    `root.data` holds `data`.
    """
    function Node{T}(data) where T
        n = new{T}(data)
        n.parent = n
        n.child = n
        n.sibling = n
        n
    end

    """
        node = Node(data, parent::Node)

    Construct a `node` with `parent` as its parent node. `node.data` stores `data`.
    Node that this does *not* update any links in, e.g., `parent`'s other children.

    For a higher-level interface, see [`addchild`](@ref).
    """
    function Node{T}(data, parent::Node) where T
        n = new{T}(data, parent)
        n.child = n
        n.sibling = n
        n
    end
end
Node(data::T) where {T} = Node{T}(data)
Node(data, parent::Node{T}) where {T} = Node{T}(data, parent)

"""
    node = lastsibling(child)

Return the last sibling of `child`.
"""
function lastsibling(sib::Node)
    newsib = sib.sibling
    while !islastsibling(sib)
        sib = newsib
        newsib = sib.sibling
    end
    sib
end

"""
    node = addsibling(oldersib, data)

Append a new "youngest" sibling, storing `data` in `node`. `oldersib` must be
the previously-youngest sibling (see [`lastsibling`](@ref)).
"""
function addsibling(oldersib::Node{T}, data) where T
    if !islastsibling(oldersib)
        error("Truncation of sibling list")
    end
    youngersib = Node(data, oldersib.parent)
    oldersib.sibling = youngersib
    youngersib
end

"""
    node = addchild(parent::Node, data)

Create a new child of `parent`, storing `data` in `node.data`.
This adjusts all links to ensure the integrity of the tree.
"""
function addchild(parent::Node{T}, data) where T
    newc = Node(data, parent)
    prevc = parent.child
    if prevc === parent
        parent.child = newc
    else
        prevc = lastsibling(prevc)
        prevc.sibling = newc
    end
    newc
end

"""
    child = addchild(parent::Node{T}, data::Node{T}) where {T}

Add a node `data` as the last child of `parent`. Requires that `data` is a root node.
"""
function addchild(parent::Node{T}, data::Node{T}) where T
    if !isroot(data)
        error("Child node must be a root node")
    end
    prevc = parent.child
    if prevc === parent
        parent.child = data
    else
        prevc = lastsibling(prevc)
        prevc.sibling = data
    end
    data.parent = parent
    data
end

"""
    new_root = copy_subtree(root::Node{T}) where {T}

Get a shallow copy of the subtree rooted at `root`. Note that this does not copy the
data, and only copies the tree structure.
"""
function copy_subtree(root::Node{T}) where {T}
    new_root = Node{T}(root.data)
    if !isleaf(root)
        last_child = new_root
        for child in root
            new_child = copy_subtree(child)
            if last_child === new_root
                new_root.child = new_child
            else
                last_child.sibling = new_child
            end
            new_child.parent = new_root
            last_child = new_child
        end
    end
    return new_root
end

"""
    isroot(node)

Returns `true` if `node` is the root of a tree (meaning, it is its own parent).
"""
AbstractTrees.isroot(n::Node) = n === n.parent

"""
    islastsibling(node)

Returns `true` if `node` is the last sibling
"""
islastsibling(n::Node) = n === n.sibling

"""
    isleaf(node)

Returns `true` if `node` has no children.
"""
isleaf(n::Node) = n === n.child

makeleaf!(n::Node) = n.child = n

makelastsibling!(n::Node) = n.sibling = n

Base.show(io::IO, n::Node) = print(io, "Node(", n.data, ')')

# Iteration over children
# for c in parent
#     # do something
# end
Base.IteratorSize(::Type{<:Node}) = Base.SizeUnknown()
Base.eltype(::Type{Node{T}}) where T = Node{T}

function Base.iterate(n::Node, state::Node = n.child)
    n === state && return nothing
    return state, islastsibling(state) ? n : state.sibling
end

function showedges(io::IO, parent::Node, printfunc = identity)
    str = printfunc(parent.data)
    if str != nothing
        if isleaf(parent)
            println(io, str, " has no children")
        else
            print(io, str, " has the following children: ")
            for c in parent
                print(io, printfunc(c.data), "    ")
            end
            print(io, "\n")
            for c in parent
                showedges(io, c, printfunc)
            end
        end
    end
end
showedges(parent::Node) = showedges(stdout, parent)

depth(node::Node) = depth(node, 1)
function depth(node::Node, d)
    childd = d + 1
    for c in node
        d = max(d, depth(c, childd))
    end
    return d
end

"""
    graftchildren!(dest, src)

Move the children of `src` to become children of `dest`.
`src` becomes a leaf node.
"""
function graftchildren!(dest, src)
    # With no children to move, `dest` is left untouched and `src` is already a leaf.
    isleaf(src) && return dest
    for c in src
        c.parent = dest
    end
    if isleaf(dest)
        dest.child = src.child
    else
        lastsib = lastsibling(dest.child)
        lastsib.sibling = src.child
    end
    makeleaf!(src)
    return dest
end

"""
    prunebranch!(node)

Eliminate `node` and all its children from the tree.
"""
function prunebranch!(node)
    isroot(node) && error("cannot prune the root")
    p = node.parent
    if p.child === node
        # `node` is the first child of p
        if islastsibling(node)
            makeleaf!(p)   # p is now a leaf
        else
            p.child = node.sibling
        end
    else
        # `node` is a middle or last child of p
        child = p.child
        sib = child.sibling
        while sib !== node
            @assert sib !== child
            child = sib
            sib = child.sibling
        end
        if islastsibling(sib)
            # node is the last child of p, just truncate
            makelastsibling!(child)
        else
            # skip over node
            child.sibling = sib.sibling
        end
    end
    return p
end

"""
    isequiv(a::Node, b::Node)

Return `true` if the subtrees rooted at `a` and `b` have the same shape and
equal `data` at every corresponding node.

This is a comparison by value: it can return `true` for two distinct objects.
To test whether `a` and `b` are the same node, use `a === b`. Note that `==`
and `isequal` on `Node`s test identity, not equivalence.
"""
function isequiv(a::Node, b::Node)
    a.data == b.data || return false
    reta, retb = iterate(a), iterate(b)
    while true
        reta === retb === nothing && return true
        ((reta === nothing) || (retb === nothing)) && return false
        childa, statea = reta
        childb, stateb = retb
        isequiv(childa, childb) || return false
        reta, retb = iterate(a, statea), iterate(b, stateb)
    end
end

include("abstracttrees.jl")

end # module
