Pigeon Principle for setoids

PigeonPrinciple.fst

module PigeonPrinciple
 
open FStar.IntegerIntervals 

open FStar.Seq
  
type binary_relation (a: Type) = a -> a -> bool
 
[@@"opaque_to_smt"]
let is_reflexive  (#a:Type) (r: binary_relation a) = forall (x:a).     x `r` x
[@@"opaque_to_smt"]
let is_symmetric  (#a:Type) (r: binary_relation a) = forall (x y:a).   x `r` y == y `r` x
[@@"opaque_to_smt"]
let is_transitive (#a:Type) (r: binary_relation a) = forall (x y z:a). x `r` y /\ y `r` z ==> x `r` z 
 
type equivalence_relation (a: Type) 
  = r:binary_relation a{is_reflexive r /\ is_symmetric r /\ is_transitive r}
 
let trans_lemma (#a:Type) (eq: equivalence_relation a) (x y z:a)
  : Lemma (requires (eq x y \/ eq y x) /\ (eq y z \/ eq z y))  
          (ensures (x `eq` z) && (z `eq` x))   
  = reveal_opaque (`%is_transitive) (is_transitive eq);
    reveal_opaque (`%is_symmetric) (is_symmetric eq)

let refl_lemma #a (eq:equivalence_relation a) (x:a) 
  : Lemma (eq x x) = reveal_opaque (`%is_reflexive) (is_reflexive #a) 

let symm_lemma (#a:Type) (eq:equivalence_relation a) (x y:a) 
  : Lemma (eq x y == eq y x) = reveal_opaque (`%is_symmetric) (is_symmetric eq) 
  
let contains_eq #a (eq: equivalence_relation a) (s: seq a) (x:a)
  = exists (i:under (length s)). eq x (index s i)
  
let rec find (#a: Type) (s: seq a) (p: (a -> bool)) (i: under (length s))
    : Pure (option (under (length s)))
      (requires True)
      (ensures
        (function
          | None -> (forall (k: nat{i <= k /\ k < length s}). p (index s k) == false)
          | Some j -> i <= j /\ p (index s j)))
      (decreases (length s - i)) =
  if p (index s i) then Some i else if i + 1 < length s then find #a s p (i + 1) else None

let items_of #a (eq: equivalence_relation a) (s: seq a) = x:a{contains_eq eq s x} 


let rec count_eq #a (eq: equivalence_relation a) (s: seq a) (x:a)
  : Tot(nat) (decreases length s) =
  if length s = 0 then 0
  else if x `eq`  head s then 1 + count_eq eq (tail s) x
  else count_eq eq (tail s) x

let tail_contains_eq #a (eq: equivalence_relation a) (s:seq a) (x:a{contains_eq eq s x /\ ~(eq x (head s))})
  : Lemma (contains_eq eq (tail s) x) = 
  assert (length s > 0);
  let t = tail s in 
  assert (forall (i: under (length t)). index s (i+1) == index t i);
  eliminate exists (i: under (length s)). eq x (index s i)
  returns exists (k: under (length t)). eq x (index t k) with _.
  begin
    assert (i>0);
    assert (eq x (index s i));
    assert (index s i == index t (i-1))
  end 

let rec index_eq #a (eq:equivalence_relation a) (s: seq a) (x:a{contains_eq eq s x})
  : Tot(i:nat{(i<length s) /\ (x `eq` index s i) /\ (forall (j: under i). not (x `eq` index s j))}) (decreases length s) =
  assert(length s > 0);
  if length s = 1 then 0
  else begin
    if x `eq` index s 0 then 0 else begin
      tail_contains_eq eq s x;
      assert (contains_eq eq (tail s) x);
      let ieq = index_eq eq (tail s) x in
      let res = 1 + ieq in
      let aux (i: under res) : Lemma (not (x `eq` index s i)) =
        if i > 0 then assert (index (tail s) (i-1) == index s i) 
        in Classical.forall_intro aux;
      res
    end 
  end

#set-options "--query_stats"
let rec pigeonhole #a (eq: equivalence_relation a) 
                      (n: seq a{length n > 0}) 
                      (s: seq (items_of eq n))
  : Pure (under (length s) * under (length s))
         (requires length s > length n)
         (ensures (fun (i1,i2) -> i1<i2 /\ (index s i1 `eq` index s i2)))
         (decreases length n) = 
  if length n = 1 then begin 
    assert (length s > 1);
    trans_lemma eq (index s 0) (index n 0) (index s 1); 
    (0,1)
  end
  else begin
    let k0 = index s 0 in
    match find s (fun k -> eq k k0) 1 with
    | Some i -> symm_lemma eq (index s 0) (index s i);
               0,i
    | None ->
      let index_of_k0 = index_eq eq n k0 in
      let n_no_k0 = append (slice n 0 (index_of_k0)) (slice n (index_of_k0+1) (length n)) in
  //    assert (eq k0 ( index n index_of_k0 ));
  //    if index_of_k0 + 1 < length n then
  //      assert (index n (index_of_k0+1) == index n_no_k0 (index_of_k0));
  //    assert (length n_no_k0 = length n - 1);
      let aux (x:items_of eq n{not (eq x k0)}) 
        : Lemma (contains_eq eq n_no_k0 x) = 
          let ieq = index_eq eq n x in
           
          if ieq < index_of_k0 then begin
            assert (index n ieq == index n_no_k0 ieq)
          end else begin   
            reveal_opaque (`%is_transitive) (is_transitive eq);
            reveal_opaque (`%is_symmetric) (is_symmetric eq);
    //        assert (ieq > index_of_k0);
            assert (index n ieq == index n_no_k0 (ieq-1)) 
          end; 
        () in Classical.forall_intro aux; 
      let reduced_t = init (length s - 1) 
                      (fun i -> let k = index s (i+1) in
                            // assert (not(k `eq` k0));
                             (index s (i+1)) <: (items_of eq n_no_k0))
      in
      let i1, i2 = pigeonhole eq n_no_k0 reduced_t in
      (i1+1, i2+1)
  end