PriorityQueue.fs

[<AutoOpen>]
module PriorityQueue =

    [<AutoOpen>]
    module private Heap =

        type 'a Node =
            | Empty
            | Node of 'a NodeData
        and 'a NodeData =
            { Value: 'a
              Left: 'a Node
              Right: 'a Node
              Height: uint32
              Priority: uint16 }

        let singleton (priority, value) =
            Node { Height = 1u; Priority = priority; Value = value; Left = Empty; Right = Empty }

        let height = function
            | Empty -> 0u
            | Node {Height = h} -> h

        let private makeNode (priority, value, left, right) =
            if height left >= height right then
                Node { Height = height right + 1u; Priority = priority; Value = value; Left = left; Right = right }
            else
                Node { Height = height left + 1u; Priority = priority; Value = value; Left = right; Right = left }

        let rec merge = function
            | Empty, h -> h
            | h, Empty -> h
            | Node left, Node right ->
                if left.Priority <= right.Priority then
                    makeNode (left.Priority, left.Value, left.Left, merge (left.Right, Node right))
                else
                    makeNode (right.Priority, right.Value, right.Left, merge (right.Right, Node left))

        let insert heap (priority, value) =
            merge (heap, singleton (priority, value))

        let tryPeak = function
            | Empty -> None
            | Node node -> (node.Priority, node.Value) |> Some

        let tryPop = function
            | Empty -> None
            | Node node -> ((node.Priority, node.Value), merge (node.Left, node.Right)) |> Some

    // wrapper for convinient API
    type PriorityQueue<'a>
        private (node: 'a Node, count: uint32) =

        static member singleton (priority, value) = PriorityQueue((singleton (priority, value)), 1u)
        member __.Count = count
        member __.Insert(priority, value) =
            let newData = insert node (priority, value)
            PriorityQueue(newData, count + 1u)
        member __.TryPop() =
            match tryPop node with
            | None -> None
            | Some (pop, tail) -> Some (pop, PriorityQueue(tail, count - 1u))

        member __.TryPeak() = tryPeak node

    [<RequireQualifiedAccess>]
    [<CompilationRepresentation(CompilationRepresentationFlags.ModuleSuffix)>]
    module PriorityQueue =
        let singleton (priority, value) = PriorityQueue<_>.singleton (priority, value)
        let insert (queue: PriorityQueue<_>) (priority, value) = queue.Insert(priority, value)
        let tryPop (queue: PriorityQueue<_>) = queue.TryPop()
        let tryPeak (queue: PriorityQueue<_>) = queue.TryPeak()
        let count (queue: PriorityQueue<_>) = queue.Count