braun-tree.timl

(* Braun trees *)

(* Braun trees are a kind of balanced binary tree where the element stored in each node is no larger than elements in its children, and the size of each left child is equal to or larger by one than its right sibling. They are used to implement functional priority queues (or binary heaps), where the dequeue operation pops and returns the minimal element of the queue. 
 *)

(* The main functions are [insert], [peek] and [pop]. *)

functor BraunTree (K : sig
                               type t
                               idx m : Time
                               val le : t * t -- m --> bool
              end) =
struct

open Basic
open K
           
datatype braun_tree : {Nat} =
         Leaf of braun_tree {0}
       | Node {size : Nat} of t * braun_tree {ceil ($size/2)} * braun_tree {floor ($size/2)} --> braun_tree {size+1}

fun min_max {m : Time} (le : t * t -- m --> bool) (x, y) =
    if le (x, y) then
      (x, y)
    else
      (y, x)

absidx T_insert : BigO (fn n => log2 $n) with
fun insert {size : Nat} (tr : braun_tree {size}) (x : t) return braun_tree {size + 1} using T_insert size =
    case tr of
        (* implicit-index-argument inference failed here, so we have to supply index arguments explicitly *)
        Leaf => @Node {0} (x, Leaf, Leaf)
      | @Node {size'} (y, l, r) =>
        let
          val (smaller, bigger) = min_max le (x, y)
        in
          @Node {size' + 1} (smaller, insert r bigger, l)
        end
end

(* retrieve without pop *)  
absidx T_peek: Time with
fun peek {size : Nat} (tr : braun_tree {size + 1}) return t using T_peek =
    case tr of
        Node (x, _, _) => x
      | Leaf => never
end

absidx T : Time with
fun peek_option {size : Nat} (tr : braun_tree {size}) return option t using T =
    case tr of
        Leaf => NONE
      | Node (x, _, _) => SOME x
end

absidx T_sift : BigO (fn n => log2 $n) with
fun sift {size : Nat} (x, l : braun_tree {ceil ($size/2)}, r : braun_tree {floor ($size/2)}) return braun_tree {size+1} using T_sift size =
    case (l, r) of
        (Leaf, Leaf) =>
        @Node {0} (x, Leaf, Leaf)
      | (Node (y, _, _), Leaf) =>
        let
          val (smaller, bigger) = min_max le (x, y)
        in
          @Node {1} (smaller, @Node {0} (bigger, Leaf, Leaf), Leaf)
        end
      | (@Node {nl'} (lx, l1, l2), @Node {nr'} (rx, r1, r2)) =>
        (case (b2B <| le (x, lx), b2B <| le (x, rx), b2B <| le (lx, rx)) of
             (True, True, _) =>
             @Node {size} (x, l, r)
           | (_, _, True) =>
             @Node {size} (lx, @sift {nl'} (x, l1, l2), r)
           | (_, _, _) =>
             @Node {size} (rx, l, @sift {nr'} (x, r1, r2))
        )
      | _ => never
end

(* pop out the root element (minimal element) *)  
absidx T_pop : BigO (fn n => log2 $n * log2 $n) with
fun pop {size : Nat} (tr : braun_tree {size + 1}) return t * braun_tree {size} using T_pop size =
    case tr of
        Node (x, l, r) =>
        (case (l, r) of
             (Leaf, Leaf) => (x, Leaf)
           | (Node (y, _, _), Leaf) =>
             (x, @Node {0} (y, Leaf, Leaf))
           | (@Node {nl'} (lx, _, _), @Node {nr'} (rx, r1, r2)) =>
             let
               val (_, l) = @pop {nl'} l
             in
               if le (lx, rx) then
                 (x, @Node {nl'+1+nr'} (lx, r, l))
               else
                 (x, @Node {nl'+1+nr'} (rx, @sift {nr'} (lx, r1, r2), l))
             end
           | _ => never
        )
      | Leaf => never
end

end