
(** Skew heaps

    Author: Jean-Christophe Filliâtre (CNRS)
*)

module Heap

  use int.Int

  type elt
  predicate le elt elt
  clone relations.TotalPreOrder with
    type t = elt, predicate rel = le, axiom .

  type t

  function size t : int

  function occ elt t : int

  predicate mem (x: elt) (t: t) = occ x t > 0

  function minimum t : elt

  predicate is_minimum (x: elt) (t: t) =
    mem x t && forall e. mem e t -> le x e

  axiom min_is_min:
    forall t: t. 0 < size t -> is_minimum (minimum t) t

  val constant empty : t
    ensures { size result = 0 }
    ensures { forall e. not (mem e result) }

  val size (t: t) : int
    ensures { result = size t }

  val is_empty (t: t) : bool
    ensures { result <-> size t = 0 }

  val get_min (t: t) : elt
    requires { 0 < size t }
    ensures  { result = minimum t }

  val merge (t1 t2: t) : t
    ensures  { forall e. occ e result = occ e t1 + occ e t2 }
    ensures  { size result = size t1 + size t2 }

  val add (x: elt) (t: t) : t
    ensures  { occ x result = occ x t + 1 }
    ensures  { forall e. e <> x -> occ e result = occ e t }
    ensures  { size result = size t + 1 }

  val remove_min (t: t) : t
    requires { 0 < size t }
    ensures  { occ (minimum t) result = occ (minimum t) t - 1 }
    ensures  { forall e. e <> minimum t -> occ e result = occ e t }
    ensures  { size result = size t - 1 }

end

module SkewHeaps

  use int.Int
  use bintree.Tree
  use export bintree.Size
  use export bintree.Occ

  type elt

  val predicate le elt elt
  clone relations.TotalPreOrder with
    type t = elt, predicate rel = le, axiom .

  (* [e] is no greater than the root of [t], if any *)
  predicate le_root (e: elt) (t: tree elt) = match t with
    | Empty      -> true
    | Node _ x _ -> le e x
  end

  (* [t] is a heap *)
  predicate heap (t: tree elt) = match t with
    | Empty      -> true
    | Node l x r -> le_root x l && heap l && le_root x r && heap r
  end

  function minimum (tree elt) : elt
  axiom minimum_def: forall l x r. minimum (Node l x r) = x

  predicate is_minimum (x: elt) (t: tree elt) =
    mem x t && forall e. mem e t -> le x e

  (* the root is the smallest element *)
  let rec lemma root_is_min (t: tree elt) : unit
     requires { heap t && 0 < size t }
     ensures  { is_minimum (minimum t) t }
     variant  { t }
  = match t with
    | Empty -> absurd
    | Node l _ r ->
        if not (is_empty l) then root_is_min l;
        if not (is_empty r) then root_is_min r
    end

  let constant empty : tree elt = Empty
    ensures { heap result }
    ensures { size result = 0 }
    ensures { forall e. not (mem e result) }

  let get_min (t: tree elt) : elt
    requires { heap t && 0 < size t }
    ensures  { result = minimum t }
  =
    match t with
      | Empty      -> absurd
      | Node _ x _ -> x
    end

  let rec merge (t1 t2: tree elt) : tree elt
    requires { heap t1 && heap t2 }
    ensures  { heap result }
    ensures  { forall e. occ e result = occ e t1 + occ e t2 }
    ensures  { size result = size t1 + size t2 }
    variant  { size t1 + size t2 }
  =
    match t1, t2 with
    | Empty, _ -> t2
    | _, Empty -> t1
    | Node l1 x1 r1, Node l2 x2 r2 ->
       if le x1 x2 then
         Node (merge r1 t2) x1 l1
       else
         Node (merge r2 t1) x2 l2
    end

  let add (x: elt) (t: tree elt) : tree elt
    requires { heap t }
    ensures  { heap result }
    ensures  { occ x result = occ x t + 1 }
    ensures  { forall e. e <> x -> occ e result = occ e t }
    ensures  { size result = size t + 1 }
  =
    merge (Node Empty x Empty) t

  let remove_min (t: tree elt) : tree elt
    requires { heap t && 0 < size t }
    ensures  { heap result }
    ensures  { occ (minimum t) result = occ (minimum t) t - 1 }
    ensures  { forall e. e <> minimum t -> occ e result = occ e t }
    ensures  { size result = size t - 1 }
  =
    match t with
      | Empty      -> absurd
      | Node l _ r -> merge l r
    end

end

(* application *)

module HeapSort

  use array.Array
  use array.ArrayPermut
  use SkewHeaps
  clone export array.Sorted with type elt = elt, predicate le = le, axiom .
  use int.Int
  use ref.Ref
  use map.Occ as M
  use bintree.Occ as H

  let heapsort (a: array elt) : unit
    ensures { sorted a }
    ensures { permut_all (old a) a }
  =
    let n = length a in
    let t = ref empty in
    for i = 0 to n - 1 do
      invariant { heap !t && size !t = i }
      invariant { forall e.
                  M.occ e a.elts i n + H.occ e !t = M.occ e a.elts 0 n }
      t := add a[i] !t;
      assert { M.occ a[i] a.elts i n = 1 + M.occ a[i] a.elts (i+1) n }
    done;
    label I in
    for i = 0 to n - 1 do
      invariant { sorted_sub a 0 i }
      invariant { heap !t && size !t = n - i }
      invariant { forall j. 0 <= j < i -> forall e. mem e !t -> le a[j] e }
      invariant { forall e.
                  M.occ e a.elts 0 i + H.occ e !t = M.occ e (a.elts at I) 0 n }
      a[i] <- get_min !t;
      t := remove_min !t
    done

end
