Monkeying Around : Fun With Trees

Tree.fs

module Tree =
    type NodeId = | Id of string
    with
        member this.unapply = match this with | Id x -> x

    [<AutoOpen>]
    module internal Node =
        [<AbstractClass>]
        type NodeBase<'a> () = class
            member val internal Children : NodeBase<'a> list = [] with get, set 
            abstract member Level : int
        end

        type Root<'a> () = class
            inherit NodeBase<'a> ()
            override this.Level = 0
        end

        type Node<'a> (id: NodeId, value : 'a option, parent : NodeBase<'a>) = class
            inherit NodeBase<'a> ()
            member val Value   = value with get, set
            member this.Id     = id
            member this.Parent = parent
            override this.Level  = parent.Level + 1
        end

    type internal ConstructOperation<'a> =
        | PushChild   of NodeId * 'a option
        | AddSibling  of NodeId * 'a option
        | ModifyValue of ('a option -> 'a option)
        | Pop         of int option

    type VisitOperation<'a> =
        | VisitRoot
        | VisitChild of NodeId
        | ReadValue  of NodeId * 'a option
        | Pop

    type Tree<'a> () = class
        member val private ops : ConstructOperation<'a> list = [] with get, set

        member this.PushChild    x = this.ops <- PushChild              x :: this.ops; this
        member this.AddSibling   x = this.ops <- AddSibling             x :: this.ops; this
        member this.ModifyValue  x = this.ops <- ModifyValue            x :: this.ops; this
        member this.Pop         ?x = this.ops <- ConstructOperation.Pop x :: this.ops; this

        member this.Build () = TreeCursor<'a> (this.ops)
    end

    and TreeCursor<'a> internal (ops) = class
        let rec applyOp op (node : NodeBase<'a>)  : NodeBase<'a> = 
            match op with
            // Push a child on to the given node and return it        
            | PushChild (x, v) -> 
                let child = Node (x, v, node)
                node.Children <- upcast child :: node.Children
                upcast child

            // Add a sibling to the given node and return it
            | AddSibling (x, v) ->
                match node with
                | :? Node<'a> as n ->
                    let sibling = Node(x, v, n.Parent)
                    n.Parent.Children <- upcast sibling :: n.Parent.Children
                    upcast sibling
                | _ -> failwith "Cannot add sibling to root"

            // Modify the tag of the given node and return it
            | ModifyValue f -> 
                match node with
                | :? Node<'a> as n ->
                    n.Value <- f n.Value
                    upcast n
                | _ -> failwith "Cannot modify value of root"

            // Pop (recursively) to an ancestor of this node and return it
            | ConstructOperation.Pop l ->
                match node with
                | :? Node<'a> as n -> 
                    let level = l |> Option.defaultValue (n.Level - 1)
                    if (n.Parent.Level = level) then
                        n.Parent
                    elif (n.Level > level) then
                        applyOp (ConstructOperation.Pop (Some level)) (n.Parent)
                    else
                        failwith "How did we get here?"
                | _ -> failwith "Cannot pop root"

        let current = List.foldBack applyOp ops (upcast (Root()))

        let visitRoot start = 
            let rec visit (node : NodeBase<'a>) =
                match node with
                | :? Root<'a> as r -> seq { yield node }
                | :? Node<'a> as n -> seq { yield node; yield! visit n.Parent }
                | _ -> Seq.empty
            visit start

        let root = visitRoot current |> Seq.last :?> Root<'a>

        member this.PathToRoot = 
            let readValue (nb : NodeBase<'a>) = 
                match nb with
                | :? Node<'a> as n -> ReadValue (n.Id, n.Value)
                | _ -> VisitRoot 
            visitRoot current |> Seq.map readValue
        
        member this.PreOrderPath =
            let rec visit (node : NodeBase<'a>) =
                seq {
                    match node with
                    | :? Node<'a> as n -> yield ReadValue (n.Id, n.Value)
                    | _ -> yield! Seq.empty

                    for child in node.Children |> List.rev do
                        match child with
                        | :? Node<'a> as c -> 
                            yield! seq {
                                yield VisitChild c.Id
                                yield! visit child
                                yield VisitOperation.Pop
                            }
                        | _ -> yield! Seq.empty
                }
            visit root

        member this.VisitRoot<'o>     (processor : 'o -> VisitOperation<'a> -> 'o) (seed : 'o) = Seq.fold processor seed this.PathToRoot 
        member this.VisitPreOrder<'o> (processor : 'o -> VisitOperation<'a> -> 'o) (seed : 'o) = Seq.fold processor seed this.PreOrderPath
    end

open Tree

[<EntryPoint>]
let main argv = 
    let printNode res curr = 
        let c = 
            match curr with
            | ReadValue (id, vo) -> sprintf "%s%s" id.unapply (vo |> Option.map (sprintf " (%A)") |> Option.defaultValue "")
            | VisitRoot          -> "|"
            | VisitChild id      -> "↓"
            | VisitOperation.Pop -> "↑"
        sprintf "%s %s" res c

    let t = Tree<string> ()
    let t = t.PushChild  (Id "a",  None)
    let t = t.PushChild  (Id "b",  None)
    let t = t.PushChild  (Id "b1", None)
    let t = t.AddSibling (Id "b2", None)
    let t = t.Pop ()
    let t = t.AddSibling (Id "c",  None)
    let t = t.PushChild  (Id "c1", None)
    let t = t.AddSibling (Id "c2", None)
    let t = t.Pop ()
    let t = t.Pop ()
    let t = t.PushChild  (Id "d",  None)

    let tc = t.Build ()

    printfn "Path to root  : %s" <| tc.VisitRoot     printNode ""
    printfn "Pre-Order walk: %s" <| tc.VisitPreOrder printNode ""

    0 // return an integer exit code