(* Priority queue, implemented as a binary heap within a fixed capacity array

   Authors: Aymeric Walch (ENS Lyon)
            Jean-Christophe Filliâtre (CNRS)
*)

use int.Int
use mach.int.Int
use seq.Seq
use array.Array
use bag.Bag
use option.Option

type elt = int

predicate heap_order (a: array elt) (len: int) =
  forall i. 0 <= i < len ->
  (0 < 2*i + 1 < len -> a[i] <= a[2*i + 1]) /\
  (0 < 2*i + 2 < len -> a[i] <= a[2*i + 2])

type t = {
          cap: int;
  mutable len: int;
            a: array elt;
}
  invariant { 0 <= len <= length a = cap }
  invariant { heap_order a len }
  by { cap = 0; len = 0; a = Array.make 0 0; }

let create (cap: int) : t
  requires { 0 <= cap }
= { cap = cap; len = 0; a = Array.make cap 0 }

(* proving that h.a[0] is the minimum requires induction *)
let rec lemma first_is_min (h: t) (i: int)
  requires { 0 <= i < h.len }
  ensures  { h.a[0] <= h.a[i] }
  variant  { i }
= if i > 0 then first_is_min h (div (i-1) 2)

let rec move_down (a: array elt) (i: int) (x: elt) (len: int) : unit
  requires { heap_order a len }
  requires { 0 <= i < len <= length a }
  requires { i > 0 -> a[div (i - 1) 2] <= x }
  ensures  { heap_order a len }
  variant  { len - i }
= if 2*i + 1 < len then begin
    let j =
      let j = 2 * i + 1 in
      if j + 1 < len && a[j+1] < a[j] then j+1 else j in
    if a[j] < x then begin
      a[i] <- a[j];
      move_down a j x len
    end else begin
      a[i] <- x;
    end
  end else begin (* 2i+1 outside *)
    a[i] <- x;
  end

exception Empty

let extract_min_exn (h : t) : elt
  raises  { Empty -> h.len = old h.len = 0 }
  ensures { result = old h.a[0] }
  ensures { h.len = old h.len - 1 }
= if h.len = 0 then raise Empty;
  h.len <- h.len - 1;
  let x = h.a[h.len] in
  if h.len <> 0 then begin
    let min = h.a[0] in
    move_down h.a 0 x h.len;
    min
  end else begin
    x
  end

let rec move_up (a: array elt) (i: int) (x: elt) (len: int) : unit
  requires { 0 <= i < len <= length a }
  requires { heap_order a len }
  requires { 0 < 2*i + 1 < len -> x <= a[2*i + 1] }
  requires { 0 < 2*i + 2 < len -> x <= a[2*i + 2] }
  variant  { i }
  ensures  { heap_order a len }
= if i = 0 then begin
    a[i] <- x;
  end
  else begin
    let j = (i-1) / 2 in
    let y = a[j] in
    if y > x then begin
      a[i] <- y;
      move_up a j x len
    end
    else begin
      a[i] <- x;
    end
  end

exception Full

let insert (x: elt) (h: t) : unit
  raises  { Full -> h.len = old h.len = h.cap }
  ensures { h.len = old h.len + 1 }
= if h.len = h.cap then raise Full;
  let n = h.len in
  h.len <- h.len + 1;
  if n = 0 then
    h.a[0] <- x
  else
    let j = (n - 1) / 2 in
    let y = h.a[j] in
    if y > x then begin
      h.a[n] <- y;
      move_up h.a j x h.len
    end else
      h.a[n] <- x