(** {1 Priority Queues}

  Below is a traditional implementation of priority queues (module
  `Pqueue`) as a binary heap stored inside a resizeable array (module
  `Vector`).

  The code below was developed and verified as part of the VOCaL project
  (A Verified OCaml Library, project ANR-15-CE25-0008 funded by the French
  Agence Nationale de la Recherche, 2015--2021).

  This file can be translated to OCaml with the following command:
{h<pre>
    why3 extract -D ocaml64 pqueue.mlw  --modular -o /dest/dir
</pre>}

  The resulting file `pqueue__Pqueue.ml` provides an implementation of
  priority queues in a functor `Make`, with the following interface:
{h<pre>
    module Make(X : sig
      type t
      val dummy : t
      val compare : t -> t -> int
    end) : sig
      type elt = X.t
      type heap
      val create : unit -> heap
      val is_empty : heap -> bool
      val size : heap -> int
      exception Empty
      val find_min_exn : heap -> X.t
      val find_min : heap -> X.t option
      val extract_min_exn : heap -> X.t
      val delete_min_exn : heap -> unit
      val insert : X.t -> heap -> unit
    end
</pre>}

  Authors: Aymeric Walch (ENS Lyon)
           Mário Pereira (NOVA LINCS)
           Jean-Christophe Filliâtre (CNRS)
*)

module Vector

  use import mach.array.Array63 as A
  use int.Int
  use int.ComputerDivision
  use int.MinMax as MinMax
  use seq.Seq
  use mach.int.MinMax63
  use mach.int.Int63
  use mach.int.Refint63
  use option.Option
  use import map.Map as M
  use ref.Ref
  use seq.Seq
  use ocaml.Sys
  use ocaml.Pervasives

  type t 'a = {
           dummy: 'a;
    mutable size: int63;
    mutable data: A.array 'a;
    ghost mutable view: seq 'a;
  } invariant { length view = size }
    invariant { forall i. 0 <= i < size -> view[i] = data[i] }
    invariant { 0 <= size <= length data <= max_array_length }
    invariant { forall i. size <= i < length data -> data[i] = dummy }
    by {
      dummy = any 'a; size = zero; data = A.make zero (any 'a); view = empty;
    }

  let create (capacity [@ocaml:optional]: option int63)
             (dummy [@ocaml:named]: 'a) : t 'a
    requires { let capacity =
                 match capacity with None -> zero | Some c -> c end in
               0 <= capacity <= max_array_length }
    ensures  { result.size = 0 }
  = let capacity = match capacity with None -> zero | Some c -> c end in
    { dummy = dummy; size = zero; data = A.make capacity dummy; view = empty }

  let make (dummy [@ocaml:optional]: option 'a) (n: int63) (x: 'a) : t 'a
    requires { 0 <= n <= max_array_length }
    returns  { a -> n = a.size = A.length a.data }
    returns  { a -> forall i. 0 <= i < n -> Seq.([]) a.view i = x }
    returns  { a -> a.dummy = match dummy with None -> x | Some d -> d end }
  = let dummy = match dummy with None -> x | Some d -> d end in
    { dummy = dummy; size = n; data = A.make n x;
      view = Seq.create (to_int n) (fun _ -> x) }

  let init (dummy [@ocaml:named]: 'a) (n: int63) (f: int63 -> 'a) : t 'a
    requires { 0 <= n <= max_array_length }
    returns  { a -> n = a.size }
    returns  { a -> forall i: int63. 0 <= i < n -> Seq.([]) a.view i = f i }
  = let a = make None n dummy in
    for i = 0 to n - 1 do
      invariant { forall j:int63. 0 <= j < i -> Seq.([]) a.view j = f j }
      invariant { forall j. 0 <= j < i -> Seq.([]) a.view j = a.data[j] }
      invariant { Seq.length a.view = a.size }
      invariant { forall j. i <= j < A.length a.data ->
                            Seq.([]) a.view j = a.dummy = a.data[j] }
      A.([]<-) a.data i (f i);
      a.view <- Seq.set a.view (to_int i) (f i);
    done;
    a

  let length (a: t 'a) : int63
    returns  { n -> n = a.size }
  = a.size

  let get (a: t 'a) (i: int63) : 'a
    requires { 0 <= i < a.size }
    returns  { x -> x = Seq.([]) a.view i }
  = A.([]) a.data i

  let set (a: t 'a) (n: int63) (x: 'a) : unit
    requires { 0 <= n < a.size }
    ensures  { a.data.A.elts = (old a).data.A.elts[n <- x] }
  = A.([]<-) a.data n x; a.view <- Seq.set a.view (to_int n) x

  val ghost seq_fill (s: seq 'a) (ofs: int) (len: int) (v: 'a) : seq 'a
    requires { 0 <= ofs /\ 0 <= len /\ ofs + len <= Seq.length s }
    ensures  { Seq.length s = Seq.length result }
    ensures  { forall i. (0 <= i < ofs \/ ofs + len <= i < Seq.length s) ->
                         Seq.get s i = Seq.get result i }
    ensures  { forall i. ofs <= i < ofs + len -> Seq.get result i = v }

  let unsafe_resize (a: t 'a) (n: int63) : unit
    requires { 0 <= n <= max_array_length }
    ensures  { n = a.size }
    ensures  { forall i. 0 <= i < MinMax.min ((old a).size) n ->
                 Seq.([]) a.view i = Seq.([]) (old a).view i }
  = let n_old = A.length a.data in
    if n <= a.size then
      (* shrink *)
      if n < n_old / 4 then begin (* reallocate into a smaller array *)
        a.data <- A.sub a.data zero n;
        a.view <- a.view[.. to_int n] end
      else begin
        A.fill a.data n (a.size - n) a.dummy;
        a.view <- a.view[.. to_int n] end
    else begin
      (* grow *)
      if n > n_old then begin (* reallocate into a larger array *)
        let n_div2 = n / 2 in
        let n' =
          if n_div2 >= n_old then
            if max_array_length / 2 >= n_div2 then n
            else max_array_length
          else if max_array_length / 2 >= n_old then 2 * n_old
          else max_array_length in
        let a' = A.make n' a.dummy in
        A.blit a.data zero a' zero a.size;
        a.data <- a';
      end;
      let ghost dummy = a.dummy in
      let a_view' = Seq.create (to_int (n - a.size)) (fun _ -> dummy) in
      a.view <- Seq.(++) a.view a_view';
    end;
    a.size <- n

  let resize (a: t 'a) (n: int63) : unit
    ensures  { n = a.size }
    ensures  { forall i. 0 <= i < MinMax.min ((old a).size) n ->
                 Seq.([]) a.view i = Seq.([]) (old a).view i }
    raises   { Invalid_argument _ -> not (0 <= n <= max_array_length) }
  = if not (zero <= n <= max_array_length) then raise Invalid_argument "resize";
    unsafe_resize a n

  (** Array interface *)

  let clear (a: t 'a) : unit
    ensures  { a.size = 0 }
  = unsafe_resize a zero

  let is_empty (a: t 'a) : bool
    returns  { r -> r <-> a.size = 0 }
  = length a = zero

  let sub (a: t 'a) (ofs n: int63) : t 'a
    requires { 0 <= ofs /\ 0 <= n /\ ofs + n <= a.size }
    returns  { r -> n = r.size }
    returns  { r -> forall i. 0 <= i < n ->
                      Seq.([]) r.view i = Seq.([]) a.view (ofs + i) }
  = { dummy = a.dummy; size = n; data = A.sub a.data ofs n;
      view = a.view[to_int ofs .. to_int (ofs + n)] }

  let fill (a: t 'a) (ofs n: int63) (x: 'a) : unit
    requires { 0 <= ofs /\ 0 <= n /\ ofs + n <= a.size }
    writes   { a.data.elts, a.view }
    ensures  { forall i. (0 <= i < ofs \/ ofs + n <= i < a.size) ->
                 Seq.([]) a.view i = Seq.([]) (old a).view i }
    ensures  { forall i. ofs <= i < ofs + n -> Seq.([]) a.view i = x }
  = A.fill a.data ofs n x;
    a.view <- seq_fill a.view (to_int ofs) (to_int n) x

  val ghost seq_blit (s1: seq 'a) (ofs1: int) (s2: seq 'a) (ofs2: int)
                     (len: int) : seq 'a
    requires { 0 <= ofs1 /\ 0 <= len /\ ofs1 + len <= Seq.length s1 }
    requires { 0 <= ofs2 /\             ofs2 + len <= Seq.length s2 }
    ensures  { result = Seq.create (Seq.length s2)
                 (fun i -> if 0 <= i < ofs2 ||
                              ofs2 + len <= i < Seq.length s2 then
                                Seq.([]) s2 i
                           else Seq.([]) s1 (ofs1 + i - ofs2)) }

  let blit (a1: t 'a) (ofs1: int63) (a2: t 'a) (ofs2: int63) (n: int63)
    requires { 0 <= n }
    requires { 0 <= ofs1 /\ ofs1 + n <= a1.size }
    requires { 0 <= ofs2 /\ ofs2 + n <= a2.size }
    writes   { a2.data.elts, a2.view }
    ensures  { forall i. (0 <= i < ofs2 \/ ofs2 + n <= i < a2.size) ->
                 Seq.([]) a2.view i = Seq.([]) (old a2).view i }
    ensures  { forall i. ofs2 <= i < ofs2 + n ->
                 Seq.([]) a2.view i = Seq.([]) (old a1).view (ofs1 + i - ofs2) }
  = A.blit a1.data ofs1 a2.data ofs2 n;
    a2.view <- seq_blit a1.view (to_int ofs1) a2.view (to_int ofs2) (to_int n)

  let append (a1: t 'a) (a2: t 'a)
    requires { a1.size + a2.size <= max_array_length }
    returns  { a3 -> a3.size = a1.size + a2.size }
    returns  { a3 -> forall i. 0 <= i < a1.size ->
                       Seq.([]) a3.view i = Seq.([]) a1.view i }
    returns  { a3 -> forall i. 0 <= i < a2.size ->
                       Seq.([]) a3.view (a1.size + i) = Seq.([]) a2.view i }
  = let n1 = length a1 in
    let n2 = length a2 in
    let a = make None (n1 + n2) a1.dummy in
    blit a1 zero a zero n1;
    blit a2 zero a n1   n2;
    a

  (** Stack interface *)

  let push (a: t 'a) (x: 'a) : unit
    requires { a.size < max_array_length }
    ensures  { a.size = (old a).size + 1 }
    ensures  { Seq.([]) a.view (a.size - 1) = x }
    ensures  { forall i. 0 <= i < (old a).size ->
                 Seq.([]) a.view i = Seq.([]) (old a).view i }
  = let n = a.size in
    unsafe_resize a (n+one);
    a.view <- Seq.set a.view (to_int n) x;
    A.([]<-) a.data n x

  exception Empty

  let pop (a: t 'a) : 'a
    raises   { Empty -> a.size = (old a).size = 0 }
    returns  { _ -> a.size = (old a).size - 1 }
    returns  { x -> x = Seq.([]) (old a).view (a.size) }
    returns  { _ -> forall i. 0 <= i < a.size ->
                      Seq.([]) a.view i = Seq.([]) (old a).view i }
  = let n = length a - one in
    if n < zero then raise Empty;
    let r = A.([]) a.data n in
    unsafe_resize a n;
    r

  let pop_opt (a: t 'a) : option 'a
    returns  { r -> match r with
                    | None   -> a.size = (old a).size = 0
                    | Some x -> a.size = (old a).size - 1 /\
                                x = Seq.([]) (old a).view (a.size) /\
                                forall i. 0 <= i < a.size ->
                                  Seq.([]) a.view i
                                = Seq.([]) (old a).view i end }
  = let n = length a - one in
    if n < zero then None else begin
      let r = A.([]) a.data n in
      unsafe_resize a n;
      Some r
    end

end

module Pqueue
  use int.Int
  use int.ComputerDivision
  use int.Div2
  use int.NumOf
  use bag.Bag
  use mach.int.Int63
  use mach.int.Refint63
  use import Vector as V
  use seq.Seq
  use option.Option
  use ocaml.Pervasives
  use ocaml.Sys

  predicate is_pre_order (cmp: 'a -> 'a -> int63) =
    (forall x. cmp x x = 0) /\
    (forall x y. cmp x y = 0 <-> cmp y x = 0) /\
    (forall x y. cmp x y < 0 <-> cmp y x > 0) /\
    (forall x y z.
      (cmp x y = 0 -> cmp y z = 0 -> cmp x z = 0) /\
      (cmp x y = 0 -> cmp y z < 0 -> cmp x z < 0) /\
      (cmp x y < 0 -> cmp y z = 0 -> cmp x z < 0) /\
      (cmp x y < 0 -> cmp y z < 0 -> cmp x z < 0))

  scope Make
    scope X
      type t
      val dummy : t
      function cmp : t -> t -> int63
      axiom is_pre_order: is_pre_order cmp
      val compare (x : t) (y : t) : int63
        ensures { result = cmp x y }
    end

    type elt = X.t

    predicate le (x y: elt) = X.cmp x y <= 0

    function numocc (a: V.t elt) (c: elt) (l u: int) : int =
      numof (fun i -> a.view[i] = c) l u

    function numocc' (a : V.t elt) (c: elt) : int =
      numocc a c 0 (length a.view)

    predicate heap_order (a : V.t elt) = forall i.
      (0 < 2*i + 1 < length a.view -> le a.view[i] a.view[2*i + 1]) /\
      (0 < 2*i + 2 < length a.view -> le a.view[i] a.view[2*i + 2])

    type heap = {
                      a : V.t elt;
      ghost mutable bag : bag elt;
    }
      invariant { length a.view = card bag }
      invariant { forall x. numocc' a x = nb_occ x bag }
      invariant { heap_order a } (* Heap order is preserved *)
      by { a = V.create (Some (0 : int63)) X.dummy; bag = empty_bag }

    function minimum (h: heap) : elt

    axiom min_def: forall h. card h.bag <> 0 -> minimum h = h.a.view[0]

    predicate is_minimum (x: elt) (h: heap) =
      mem x h.bag && forall e. mem e h.bag -> X.cmp x e <= 0

    lemma num_exist:
      forall h x. numocc' h.a x > 0 ->
      exists i. 0 <= i < length h.a.view /\ h.a.view[i] = x

    let lemma min_coherent (h: heap)
      requires { 0 < card h.bag }
      ensures  { is_minimum (minimum h) h }
    = let s = h.a.view in
      let n = length s in
      let rec lemma ismin (i: int)
        requires { 0 <= i < n } variant { i } ensures  { le s[0] s[i] }
      = if i > 0 then ismin (div (i-1) 2)
     in ()

    (* Useful lemmas for proving invariants, most of the work will be here *)

    predicate substitution (a1 a2: V.t elt) (i : int63) =
      length a1.view = length a2.view /\
      0 <= i < length a1.view /\
      (forall k. 0 <= k < length a1.view -> k <> i -> a1.view[k] = a2.view[k])

    lemma sub_occ_1: forall a1 a2 : V.t elt, i : int63, x : elt.
      substitution a1 a2 i ->
      (x <> a1.view[i] /\ x <> a2.view[i]) ->
      numocc' a1 x = numocc' a2 x

    lemma sub_occ_2: forall a1 a2 i.
      substitution a1 a2 i -> a1.view[i] <> a2.view[i] ->
      numocc' a1 a1.view[i]
    = numocc' a2 a1.view[i] + 1
      by let x = a1.view[i] in
         numocc' a1 x = numocc a1 x 0 i + 1 + numocc a1 x (i+1) (length a1.view)
      && numocc' a2 x = numocc a2 x 0 i + 0 + numocc a2 x (i+1) (length a2.view)

    lemma sub_occ_3: forall a1 a2 i.
      substitution a1 a2 i -> a1.view[i] = a2.view[i] ->
      numocc' a1 a1.view[i]
    = numocc' a2 a1.view[i]

    let lemma sub_order (a1 a2 : V.t elt) (i : int63) =
      requires {substitution a1 a2 i}
      requires { heap_order a1 }
      requires { i > 0 -> X.cmp a1.view[div (i - 1) 2] a2.view[i] <= 0 }
      requires { 2*i + 1 < length a1.view -> X.cmp a1.view[2*i + 1] a2.view[i] >= 0 }
      requires { 2*i + 2 < length a1.view -> X.cmp a1.view[2*i + 2] a2.view[i] >= 0 }
      ensures { heap_order a2 }
      ()

    predicate pop (a1 a2: V.t elt) =
      length a2.view = length a1.view - 1 /\
      forall i. 0 <= i < length a2.view -> a1.view[i] = a2.view[i]

    lemma pop_occ_1: forall a1 a2 x.
      pop a1 a2 -> x <> a1.view[length a1.view - 1] ->
      numocc' a1 x = numocc' a2 x

    lemma pop_occ_2: forall a1 a2.
      pop a1 a2 ->
      numocc' a1 a1.view[length a1.view - 1] - 1
    = numocc' a2 a1.view[length a1.view - 1]

    lemma pop_order: forall a1 a2. pop a1 a2 -> heap_order a1 -> heap_order a2

    predicate push (a1 a2 : V.t elt ) =
      length a2.view = length a1.view + 1 /\
      forall i. 0 <= i < length a1.view -> a1.view[i] = a2.view[i]

    let ghost push_occ (a1 a2 : V.t elt) =
      requires { push a1 a2 }
      ensures { forall x. x <> a2.view[length a2.view - 1] -> numocc' a1 x = numocc' a2 x }
      ensures { numocc' a1 a2.view[length a2.view - 1] + 1 = numocc' a2 a2.view[length a2.view - 1] }
      ()

    let ghost push_order (a1 a2 : V.t elt) =
      requires { length a1.view > 0 }
      requires { push a1 a2 }
      requires { le a2.view[div (length a1.view - 1) 2] a2.view[length a1.view] }
      requires { heap_order a1 }
      ensures { heap_order a2 }
      ()

    lemma same_occ: forall b: bag elt, x a.
      x <> a -> nb_occ x (Bag.diff b (Bag.singleton a)) = nb_occ x b

    let ghost ancestor (a : V.t elt) (i : int) =
      requires { heap_order a }
      requires { 0 <= i < length a.view }
      ensures { i > 0 -> le a.view[div (i-1) 2] a.view[i] }
      ()

    let ghost children (a : V.t elt) (i : int) =
      requires { heap_order a }
      requires { 0 <= i < length a.view }
      ensures { 2*i+1 < length a.view -> le a.view[i] a.view[2*i+1] }
      ensures { 2*i+2 < length a.view -> le a.view[i] a.view[2*i+2] }
      ()

    let ghost trans (a : V.t elt) (i : int) =
       requires { 0 <= i < length a.view }
       requires { heap_order a }
       ensures {i > 0 /\ 2*i+1 < length a.view -> le a.view[div (i-1) 2] a.view[2*i+1] }
       ensures {i > 0 /\ 2*i+2 < length a.view -> le a.view[div (i-1) 2] a.view[2*i+2] }
       ancestor a i;
       children a i;
       ()

    let create () : heap
      ensures { result.bag = empty_bag }
    = { a = V.create (Some (0 : int63)) X.dummy; bag = empty_bag }

    let is_empty (h: heap) : bool
      ensures { result <-> h.bag = empty_bag }
    = V.is_empty h.a

    let size (h: heap) : int63
      ensures { result = card h.bag }
    = V.length h.a

    exception Empty

    let find_min_exn (h: heap) : elt
      raises  { Empty -> card h.bag = 0 }
      ensures { card h.bag > 0 /\ result = minimum h }
    = if V.is_empty h.a then raise Empty;
      V.get h.a 0

    let find_min (h: heap) : option elt
      ensures { match result with
      | None   -> card h.bag = 0
      | Some x -> card h.bag > 0 && x = minimum h end }
    = if V.is_empty h.a then None else Some (V.get h.a 0)

   let rec move_down (a : V.t elt) (i : int63) (x : elt) : unit
      requires { heap_order a }
      requires { 0 <= i < length a.view }
      requires { i > 0 -> le a.view[div (i - 1) 2] x }
      ensures  { heap_order a }
      ensures  { forall e. e <> (old a).view[i] -> e <> x ->
                   numocc' a e = numocc' (old a) e }
      ensures  { let o = (old a).view[i] in
                 if x = o then numocc' a x = numocc' (old a) x
                 else
                   numocc' a x = numocc' (old a) x + 1 /\
                   numocc' a o = numocc' (old a) o - 1 }
      ensures  { length a.view = length (old a).view }
      (* Needed for delete_min *)
      variant  { length a.view - i }
    = let n = V.length a in
      let q = if n=1 then -1 else (n-2)/2 in
      if i <= q then begin (* 2i+1 exists *)
        let j =
          let j = 2 * i + 1 in
          if j + 1 < n && X.compare (V.get a (j+1)) (V.get a j) < 0 then
            j+1
          else j in
        ancestor a (to_int i);
        children a (to_int i);
        let olda = pure {a} in
        if X.compare (V.get a j) x < 0 then begin
          V.set a i (V.get a j);
          assert { substitution (old a) a i };
          sub_order olda a i;
          assert { heap_order a };
          move_down a j x
        end else begin
          V.set a i x;
          assert { substitution (old a) a i };
          sub_order olda a i
        end
      end else begin (* 2i+1 outside *)
        V.set a i x;
        assert { substitution (old a) a i };
      end

    let extract_min_exn (h: heap) : elt
      raises  { Empty -> card h.bag = 0 /\ h.bag = old h.bag }
      ensures { result = minimum (old h) }
      ensures { (old h).bag = add result h.bag }
      (* Property about order is guaranteed by the invariant *)
    = try
        let x = V.pop h.a in
        let n = V.length h.a in
        if n <> 0 then begin
          let min = V.get h.a 0 in
          assert { min = minimum (old h) };
          assert { card (Bag.diff h.bag (Bag.singleton min)) = card h.bag - 1 };
          move_down h.a 0 x;
          h.bag <- Bag.diff h.bag (Bag.singleton min);
          min
        end else begin
          h.bag <- empty_bag;
          x
        end
      with V.Empty -> raise Empty end

    let delete_min_exn (h: heap) : unit
      raises  { Empty -> card h.bag = 0 /\ h.bag = old h.bag }
      ensures { (old h).bag = add (minimum (old h)) h.bag }
    = let _ = extract_min_exn h in
      ()

    let rec move_up (a : V.t elt) (i : int63) (x : elt)
      requires { heap_order a }
      requires { 0 <= i < length a.view }
      requires { 0 <= 2*i + 1 < length a.view -> le x a.view[2*i + 1] }
      requires { 0 <= 2*i + 2 < length a.view -> le x a.view[2*i + 2] }
      variant { to_int i }
      ensures { heap_order a }
      ensures { forall e. e <> (old a).view[i] -> e <> x -> numocc' a e = numocc' (old a) e }
      ensures  { let o = (old a).view[i] in
                 if x = o then
                   numocc' a x = numocc' (old a) x
                 else
                   numocc' a x = numocc' (old a) x + 1 /\
                   numocc' a o = numocc' (old a) o - 1 }
      ensures { length a.view = length (old a).view } (* Needed for add *)
    = let ghost olda = pure { a } in
      if i = 0 then begin
        V.set a i x;
        assert { substitution olda a i };
        sub_order olda a i
      end
      else begin
        let j = (i-1)/2 in
        let y = V.get a j in
        children a (to_int j);
        children a (to_int i);
        if X.compare y x > 0 then begin
          V.set a i y;
          assert { substitution olda a i };
          trans olda (to_int i);
          sub_order olda a i;
          move_up a j x
        end
        else begin
          V.set a i x;
          assert { substitution olda a i };
          sub_order olda a i
        end
      end

    let insert (x : elt) (h: heap) : unit
      raises  { Invalid_argument _ -> card h.bag >= Sys.max_array_length }
      ensures { h.bag = add x (old h).bag }
    = if size h = Sys.max_array_length then raise Invalid_argument "insert";
      let n = V.length h.a in
      let ghost olda = pure { h.a } in
      if n = 0 then
        V.push h.a x
      else begin
        let j = (n-1)/2 in
        let y = V.get h.a j in
        children h.a (to_int j);
        if X.compare y x > 0 then begin
          V.push h.a y;
          push_order olda h.a;
          push_occ olda h.a;
          assert { numocc' h.a x = numocc' olda x };
          move_up h.a j x
        end
        else begin
          V.push h.a x;
          push_occ olda h.a;
          push_order olda h.a
        end
      end;
      assert { numocc' h.a x = numocc' olda x + 1 };
      h.bag <- add x h.bag;
    end

end