
(* Verification of the following 2-lines C program solving the N-queens puzzle:

   t(a,b,c){int d=0,e=a&~b&~c,f=1;if(a)for(f=0;d=(e-=d)&-e;f+=t(a-d,(b+d)*2,(
   c+d)/2));return f;}main(q){scanf("%d",&q);printf("%d\n",t(~(~0<<q),0,0));}
*)

theory S

  use int.Int
  use export set.SetAppInt

  val succ (s: set): set
    ensures { forall i: int. mem i result <-> i >= 1 /\ mem (i - 1) s }

  val pred (s: set): set
    ensures { forall i: int. mem i result <-> i >= 0 /\ mem (i + 1) s }

end

module NQueensSetsTermination

  use S
  use ref.Refint

  let rec t (a b c : set) variant { cardinal a } =
    if not (is_empty a) then begin
      let e = ref (diff (diff a b) c) in
      let f = ref 0 in
      while not (is_empty !e) do
        invariant { subset !e (diff (diff a b) c) }
        variant { cardinal !e }
        let d = min_elt !e in
        f += t (remove d a) (succ (add d b)) (pred (add d c));
        e := remove d !e
      done;
      !f
    end else
      1

end

theory Solution

  use int.Int
  use export map.Map

  (* the number of queens *)
  constant n : int

  type solution = map int int

  (* solutions t and u have the same prefix [0..i[ *)
  predicate eq_prefix (t u: map int 'a) (i: int) =
    forall k: int. 0 <= k < i -> t[k] = u[k]

  predicate eq_sol (t u: solution) = eq_prefix t u n

  (* s stores a partial solution, for the rows 0..k-1 *)
  predicate partial_solution (k: int) (s: solution) =
    forall i: int. 0 <= i < k ->
      0 <= s[i] < n /\
      (forall j: int. 0 <= j < i -> s[i]      <> s[j] /\
                                    s[i]-s[j] <> i-j  /\
                                    s[i]-s[j] <> j-i)

  predicate solution (s: solution) = partial_solution n s

  lemma partial_solution_eq_prefix:
    forall u t: solution, k: int.
    partial_solution k t -> eq_prefix t u k -> partial_solution k u

  predicate lt_sol (s1 s2: solution) =
    exists i: int. 0 <= i < n /\ eq_prefix s1 s2 i /\ s1[i] < s2[i]

  type solutions = map int solution

  (* s[a..b[ is sorted for lt_sol *)
  predicate sorted (s: solutions) (a b: int) =
    forall i j: int. a <= i < j < b -> lt_sol s[i] s[j]

  (* a sorted array of solutions contains no duplicate *)
  lemma no_duplicate:
    forall s: solutions, a b: int. sorted s a b ->
    forall i j: int. a <= i < j < b -> not (eq_sol s[i] s[j])

end

(* 1. Abstract version of the code using sets (not ints) *******************)

module NQueensSets

  use int.Int
  use S
  use ref.Refint
  use Solution

  val ghost col: ref solution  (* solution under construction *)
  val ghost k  : ref int       (* current row in the current solution *)

  val ghost sol: ref solutions (* all solutions *)
  val ghost s  : ref int       (* next slot for a solution
                                  = current number of solutions *)

  let rec t3 (a b c : set) variant { cardinal a }
    requires { 0 <= !k /\ !k + cardinal a = n /\ !s >= 0 /\
      (forall i: int. mem i a <->
        (0<=i<n /\ forall j: int. 0 <= j < !k ->  !col[j] <> i)) /\
      (forall i: int. i>=0 -> (not (mem i b) <->
        (forall j: int. 0 <= j < !k -> !col[j] <> i + j - !k))) /\
      (forall i: int. i>=0 -> (not (mem i c) <->
        (forall j: int. 0 <= j < !k -> !col[j] <> i + !k - j))) /\
      partial_solution !k !col }
    ensures { result = !s - old !s >= 0 /\ !k = old !k /\
      sorted !sol (old !s) !s /\
      (forall t: solution.
         ((solution t /\ eq_prefix !col t !k) <->
          (exists i: int. old !s <= i < !s /\ eq_sol t !sol[i]))) /\
      (* assigns *)
      eq_prefix (old !col) !col !k /\
      eq_prefix (old !sol) !sol (old !s) }
  = if not (is_empty a) then begin
      let e = ref (diff (diff a b) c) in
      (* first, you show that if u is a solution with the same k-prefix as col, then
          u[k] (the position of the queen on the row k) must belong to e *)
      assert { forall u:solution. solution u /\ eq_prefix !col u !k ->
               S.mem u[!k] !e };
      let f = ref 0 in
      label L in
      while not (is_empty !e) do
        invariant {
          !f = !s - (!s at L) >= 0 /\ !k = !k at L /\
          subset !e (diff (diff a b) c) /\
          partial_solution !k !col /\
          sorted !sol (!s at L) !s /\
          (forall i j: int. mem i (diff (!e at L) !e) -> mem j !e -> i < j) }
        invariant { forall i:int. (!s at L) <= i < !s -> 
           solution !sol[i] /\ eq_prefix !col !sol[i] !k /\ 
           mem !sol[i][!k] (diff (!e at L) !e) }
        invariant {
          (forall t: solution.
           (solution t /\ eq_prefix !col t !k /\ mem t[!k] (diff (!e at L) !e))
           ->
           mem t[!k] (!e at L) && not (mem t[!k] !e) &&
           (exists i: int. (!s at L) <= i < !s /\ eq_sol t !sol[i])) }
        (* assigns *)
        invariant { eq_prefix (!col at L) !col !k }
        invariant { eq_prefix (!sol at L) !sol (!s at L) }
        variant { cardinal !e }
        let d = min_elt !e in
        ghost col := !col[!k <- d];
        ghost incr k;
        f += t3 (remove d a) (succ (add d b)) (pred (add d c));
        ghost decr k;
        e := remove d !e
      done;
      !f
    end else begin
      ghost sol := !sol[!s <- !col];
      ghost incr s;
      1
    end

  let queens3 (q: int)
    requires { 0 <= q = n /\ !s = 0 /\ !k = 0 }
    ensures { result = !s /\ sorted !sol 0 !s /\
      forall t: solution.
        solution t <-> (exists i: int. 0 <= i < result /\ eq_sol t !sol[i]) }
  = t3 (interval 0 q) (empty ()) (empty ())

end
