(* Companion source code to                                       *)
(*     Chris Okasaki                                              *)
(*     "The Role of Lazy Evaluation in Amortized Data Structures" *)
(*     ICFP'96                                                    *)

(* worst-case binomial queues *)

(* insert takes O(1) worst-case time                       *)
(* findMin, deleteMin, merge take O(log n) worst-case time *)


(* Streams *)

open System.Unsafe.Susp
    (* import susp, delay, force *)

datatype 'a StreamNode = Nil | Cons of 'a * 'a Stream
withtype 'a Stream = 'a StreamNode susp

val emptyStream = delay (fn () => Nil)
fun isEmptyStream s = case force s of Nil => true | Cons (x, s) => false
fun cons (x, s) = delay (fn () => Cons (x, s))
fun normalize s = case force s of Nil => () | Cons (x, s) => normalize s

(* Binomial Queues *)

type Elem = int  (* elements may be any ordered type *)

datatype Tree = Node of Elem * Tree list (* children in decreasing order of size *)
datatype Digit = Zero | One of Tree
type Schedule = Digit Stream list  (* list of delayed calls to addUnique *)
type BinQueue = Digit Stream * Schedule  (* digits/trees in increasing order of size *)

exception Empty

local  (* auxiliary functions *)
  fun root (Node (x,c)) = x

  fun link (Node (x,c)) (Node (y,d)) =
        if x <= y then Node (x, Node (y,d) :: c)
                  else Node (y, Node (x,c) :: d)

  fun addUnique t q = (* add one to low-order digit, link/carry if already a one *)
        delay (fn () =>
            case force q of
                Nil => Cons (One t, emptyStream)
              | Cons (Zero, q) => Cons (One t, q)
              | Cons (One t', q) => Cons (Zero, addUnique (link t t') q))

  fun smerge q1 q2 = (* add digit streams, link/carry when two ones are in the same position *)
        case (force q1, force q2) of
            (Nil, _) => q2
          | (_, Nil) => q1
          | (Cons (Zero, q1), Cons (digit, q2)) => cons (digit, smerge q1 q2)
          | (Cons (digit, q1), Cons (Zero, q2)) => cons (digit, smerge q1 q2)
          | (Cons (One t1, q1), Cons (One t2, q2)) =>
                 cons (Zero, addUnique (link t1 t2) (smerge q1 q2))

  fun getMin q = (* find and remove the tree with the minimum root *)
        case force q of
            Nil => raise Empty
          | Cons (Zero, q) =>  (* zero is never the last digit *)
                let val (t, q) = getMin q
                in (t, cons (Zero, q)) end
          | Cons (One t, q) =>
                if isEmptyStream q then (t, emptyStream)
                else let val (t', q') = getMin q
                     in 
                         if root t <= root t' then (t, q)
                                              else (t', cons (One t, q'))
                     end

  fun execute [] = []
    | execute (job :: schedule) =  (* execute first job in schedule *)
        case force job of
            Cons (One t, _) => schedule            (* addUnique terminates *)
          | Cons (Zero, job') => job' :: schedule  (* addUnique continues *)

  val execute2 = execute o execute  (* execute two jobs *)

in

  val empty = (emptyStream, [])

  fun isEmpty (q, schedule) = isEmptyStream q

  fun insert x (q, schedule) =
        let val q' = addUnique (Node (x,[])) q
        in (q', execute2 (q' :: schedule)) end

  fun merge (q1,schedule1) (q2,schedule2) =
        let val q = smerge q1 q2
        in
            normalize q;  (* force and memoize entire stream *)
            (q, [])
        end

  fun findMin (q, schedule) =
        let val (t, _) = getMin q
        in root t end

  fun deleteMin (q, schedule) =
        let val (Node (x,c), q') = getMin q
            fun ones [] = emptyStream
              | ones (t :: ts) = cons (One t, ones ts)
            val c' = ones (rev c)  (* convert children into a queue *)
            val q'' = smerge c' q'
        in
            normalize q'';  (* force and memoize entire stream *)
            (q'', [])
        end
end