(** Dijkstra's "Dutch national flag" *)

module Flag

  use int.Int
  use ref.Ref
  use array.Array
  use array.ArraySwap
  use array.ArrayPermut

  type color = Blue | White | Red

  predicate monochrome (a:array color) (i:int) (j:int) (c:color) =
    forall k:int. i <= k < j -> a[k]=c

  (* We scan the array from left to right using [i] and we maintain
     the following invariant, using indices [b] and [r]:

       0         b          i           r
      +---------+----------+-----------+-------+
      |  Blue   |  White   |     ?     |  Red  |
      +---------+----------+-----------+-------+

  *)

  let dutch_flag (a:array color) : unit
    ensures  { exists b r: int.
               monochrome a 0 b Blue /\
               monochrome a b r White /\
               monochrome a r (length a) Red }
    ensures  { permut_all (old a) a }
    =
    let b = ref 0 in
    let i = ref 0 in
    let r = ref (length a) in
    while !i < !r do
      invariant { 0 <= !b <= !i <= !r <= length a }
      invariant { monochrome a 0  !b Blue }
      invariant { monochrome a !b !i White }
      invariant { monochrome a !r (length a) Red }
      invariant { permut_all (old a) a }
      variant   { !r - !i }
      match a[!i] with
      | Blue ->
          swap a !b !i;
          b := !b + 1;
          i := !i + 1
      | White ->
          i := !i + 1
      | Red ->
          r := !r - 1;
          swap a !r !i
      end
    done

end