(*
   VSTTE'10 competition http://www.macs.hw.ac.uk/vstte10/Competition.html
   Problem 4: N-queens

   Author: Jean-Christophe Filliatre (CNRS)
   Tool:   Why3 (see http://why3.lri.fr/)
*)

module NQueens

  use int.Int
  use array.Array

  predicate eq_board (b1 b2: array int) (pos: int) =
    forall q. 0 <= q < pos -> b1[q] = b2[q]

  lemma eq_board_set:
    forall b: array int, pos q i: int.
    pos <= q -> eq_board b b[q <- i] pos

  lemma eq_board_sym:
    forall b1 b2: array int, pos: int.
    eq_board b1 b2 pos -> eq_board b2 b1 pos

  lemma eq_board_trans:
    forall b1 b2 b3: array int, pos: int.
    eq_board b1 b2 pos -> eq_board b2 b3 pos -> eq_board b1 b3 pos

  lemma eq_board_extension:
    forall b1 b2: array int, pos: int.
    eq_board b1 b2 pos -> b1[pos] = b2[pos] -> eq_board b1 b2 (pos+1)

  predicate consistent_row (board: array int) (pos: int) (q: int) =
    board[q] <> board[pos] /\
    board[q] - board[pos] <> pos - q /\
    board[pos] - board[q] <> pos - q

  lemma consistent_row_eq:
    forall b1 b2: array int, pos: int.
    eq_board b1 b2 (pos+1) -> forall q: int. 0 <= q < pos ->
      consistent_row b1 pos q -> consistent_row b2 pos q

  predicate is_consistent (board: array int) (pos: int) =
    forall q. 0 <= q < pos -> consistent_row board pos q

  exception Inconsistent int

  let check_is_consistent (board: array int) (pos: int) : bool
    requires { 0 <= pos < length board }
    ensures { result <-> is_consistent board pos }
  = try
      for q = 0 to pos - 1 do
        invariant { forall j. 0 <= j < q -> consistent_row board pos j }
        let bq   = board[q]   in
        let bpos = board[pos] in
        if bq        = bpos    then raise (Inconsistent q);
        if bq - bpos = pos - q then raise (Inconsistent q);
        if bpos - bq = pos - q then raise (Inconsistent q)
      done;
      True
    with Inconsistent q ->
      assert { not (consistent_row board pos q) };
      False
    end

  predicate is_board (board: array int) (pos: int) =
    forall q. 0 <= q < pos -> 0 <= board[q] < length board

  predicate solution (board: array int) (pos: int) =
    is_board board pos /\
    forall q. 0 <= q < pos -> is_consistent board q

  lemma solution_eq_board:
    forall b1 b2: array int, pos: int. length b1 = length b2 ->
    eq_board b1 b2 pos -> solution b1 pos -> solution b2 pos

  exception Solution

  let rec bt_queens (board: array int) (n: int) (pos: int) : unit
    variant  { n - pos }
    requires { length board = n /\ 0 <= pos <= n /\ solution board pos }
    ensures  { eq_board board (old board) pos }
    ensures  { forall b:array int. length b = n -> is_board b n ->
                 eq_board board b pos -> not (solution b n) }
    raises   { Solution -> solution board n }
  = if pos = n then raise Solution;
    for i = 0 to n - 1 do
      invariant { eq_board board (old board) pos }
      invariant { forall b:array int. length b = n -> is_board b n ->
            eq_board board b pos -> 0 <= b[pos] < i -> not (solution b n) }
      board[pos] <- i;
      if check_is_consistent board pos then bt_queens board n (pos + 1)
    done

  let queens (board: array int) (n: int) : unit
    requires { length board = n }
    ensures  { forall b:array int. length b = n -> is_board b n ->
               not (solution b n) }
    raises   { Solution -> solution board n }
  = bt_queens board n 0


  exception BenchFailure

  let test8 () raises { BenchFailure -> true } =
    let a = Array.make 8 0 in
    try
      queens a 8;
      raise BenchFailure
    with Solution -> a
    end

  (** variant: counting solutions (not part of the competition)

      TODO: prove soundness i.e. we indeed count the number of solutions *)

  use ref.Refint

  let rec count_bt_queens (board: array int) (n: int) (pos: int) : (_r: int)
    variant  { n - pos }
    requires { length board = n /\ 0 <= pos <= n /\ solution board pos }
    ensures  { eq_board board (old board) pos }
  = if pos = n then
      1
    else begin
      let s = ref 0 in
      for i = 0 to n - 1 do
        invariant { eq_board board (old board) pos }
        board[pos] <- i;
        if check_is_consistent board pos then
          s += count_bt_queens board n (pos+1)
      done;
      !s
    end

  let count_queens (board: array int) (n: int) : (_r: int)
    requires { length board = n }
    ensures  { true }
  = count_bt_queens board n 0

  let test_count_8 () =
    let a = Array.make 8 0 in
    count_queens a 8

end

(** counting solutions with 63-bit machine integers *)

module NQueens63

  use int.Int
  use ref.Refint
  use seq.Seq
  use mach.array.Array63
  use mach.int.Int63

  predicate is_board (board: array int63) (pos: int) =
    forall q. 0 <= q < pos -> 0 <= board[q] < (length board)

  exception MInconsistent

  let check_is_consistent (board: array int63) (pos: int63) : bool
    requires { 0 <= pos < (length board) }
    requires { is_board board (pos + 1) }
  = try
      let q = ref 0 in
      while !q < pos do
        invariant { 0 <= !q <= pos }
        invariant { is_board board (pos + 1) }
        variant   { pos - !q }
        let bq   = board[!q]   in
        let bpos = board[pos] in
        if bq        = bpos     then raise MInconsistent;
        if bq - bpos = pos - !q then raise MInconsistent;
        if bpos - bq = pos - !q then raise MInconsistent;
        q := !q + 1
      done;
      True
    with MInconsistent ->
      False
    end

  use mach.peano.Peano as P

  let rec count_bt_queens
    (solutions: ref P.t) (board: array int63) (n: int63) (pos: int63) : unit
    requires { (length board) = n }
    requires { 0 <= pos <= n }
    requires { is_board board (pos) }
    variant  { n - pos }
    ensures  { is_board board (pos) }
  =
    if pos = n then
      solutions := P.succ !solutions
    else
      let i = ref 0 in
      while !i < n do
        invariant { 0 <= !i <= n }
        invariant { is_board board (pos) }
        variant   { n - !i }
        board[pos] <- !i;
        if check_is_consistent board pos then
          count_bt_queens solutions board n (pos + 1);
        i := !i + 1
      done

  let count_queens (n: int63) : (_p: P.t)
    requires { n >= 0 }
    ensures  { true }
  =
    let solutions = ref P.zero in
    let board = Array63.make n 0 in
    count_bt_queens solutions board n 0;
    !solutions

  let test_count_8 () =
    let n = 8 in
    count_queens n

end

