FStar Seq Partition Theory

FStar.Seq.Partitions.fst

module FStar.Seq.Partitions
 
open FStar.Seq
open FStar.IntegerIntervals

(**
  This module contains theory for no-duplicates seqence partitions.
  We only work with eqtypes here.

  Here we call 
  - a (finite) [set] a sequence with no repeating elements,
  - a [part] of a given set [s] a subset sequence entirely contained in [s], 
  and 
  - a [partition] of a given set [s] a set of non-intersecting subsets
    such that its union contains the same set of elements as [s].
    More precisely, a (p: seq (seq t)) is a partition of (s: seq t)
    exactly when 
      * [s] is a set (i.e. no repeating elements)
      * all sequences in [p] are subsets of [s]
      * there are no intersections between subsets in [p]
      * any element [x] in [s] is contained in some subset among [p]

  We prove that total length of all subsets in partition [p]
  should always add up to the length of [s].
*)

(* a seq with no intersections can be thought of as a set, if
   we discard the order and only count membership *)
let seq_is_set (#t: eqtype) (s: seq t) = forall (x:t). count x s <= 1

(* in this module, we restrict most of our definitions to sets,
   even though some would have worked with general seqs *)
let is_subset_of (#t: eqtype) (subset: seq t) (superset: seq t) 
  : Pure prop (requires seq_is_set subset /\ seq_is_set superset)
              (ensures fun _ -> True)
  = forall (x:t). (mem x subset ==> mem x superset)

(* set equality only accounts for mem, disregarding element order *)
let equal_as_sets (#t: eqtype) (a b: seq t)
  : Pure prop (requires seq_is_set a /\ seq_is_set b)
              (ensures fun _ -> True)
  = forall (x:t). (mem x a <==> mem x b)

(* True if two sets share no elements *)
let no_intersections (#t: eqtype) (a b: seq t)
  : Pure prop (requires seq_is_set a /\ seq_is_set b)
              (ensures fun _ -> True)
  = forall (x:t). not (mem x a && mem x b)
(* True if any pair of sets from [p] has empty intersection *)
let rec no_intersections_among (#t: eqtype) (p: seq (seq t)) 
  : Pure prop (requires (forall (i: under (length p)). seq_is_set (index p i)))
              (ensures fun _ -> True)         
              (decreases length p)
   = (length p = 0) \/
   ( length p > 0 /\
   ((forall (i: under (length (tail p))). no_intersections (head p) (index (tail p) i)) /\
     no_intersections_among (tail p)))

(* True if each set in the collection [parts] is a subset of [s] *)
let are_parts_of (#t: eqtype) (parts: seq (seq t)) (s: seq t)
  = seq_is_set s /\
    (forall (i: under (length parts)). seq_is_set (index parts i) /\ is_subset_of (index parts i) s)

(* true if given element is found in one of the sets in [parts] collection *)
let rec is_found_in (#t:eqtype) (x:t) (parts: seq (seq t))
  : Tot bool (decreases length parts)
  = if length parts = 0 then false
    else mem x (head parts) || is_found_in x (tail parts)

(* auxiliary lemmas *)
let count_head_lemma (#t:eqtype) (s: seq t{length s > 0}) (x:t)
  : Lemma (requires head s = x) (ensures count x (tail s) == count x s - 1) = ()
  
let count_non_head_lemma (#t:eqtype) (s: seq t{length s > 0}) (x:t)
  : Lemma (requires head s <> x) (ensures count x (tail s) == count x s) = ()

(* element removal with all relevant properties assured *)
let rec remove_element (#t: eqtype) (s: seq t) (x: t)
  : Pure (seq t) (requires seq_is_set s /\ mem x s)
                 (ensures fun r -> seq_is_set r 
                              /\ not (mem x r)
                              /\ is_subset_of r s
                              /\ (forall (p:t). mem p r == (mem p s && (p <> x)))
                              /\ (length r == length s - 1))
                 (decreases length s) = 
  assert (length s > 0);
  if head s = x then tail s
  else begin 
    let tl = (remove_element (tail s) x) in
    let r = cons (head s) tl in
    lemma_eq_elim (tail r) tl; 
    let aux_set_property (p:t) : Lemma (count p r <= 1) = 
      if p = head s then count_head_lemma r p 
      else count_non_head_lemma r p
    in Classical.forall_intro aux_set_property;
    r
  end                             

(* subset subtraction with all relevant properties assured *)
let rec subtract_subset (#t: eqtype) (a b: seq t)
  : Pure (seq t)
         (requires seq_is_set a 
                 /\ seq_is_set b 
                 /\ is_subset_of b a)
         (ensures fun diff -> seq_is_set diff
                         /\ is_subset_of diff a
                         /\ (length diff == length a - length b)
                         /\ (forall (x:t). mem x diff == (mem x a && not (mem x b))))
         (decreases length b) =
    if length b = 0 then a 
    else subtract_subset (remove_element a (head b)) (tail b) 

(* sum of lengths of all sequences in collection [s] *)
let rec total_length (#t: eqtype) (s: seq (seq t)) : Tot nat (decreases length s)
  = if length s = 0 then 0
    else length (head s) + total_length (tail s)

(* true if the collection [parts] is actually a partition of set [s] *)
let is_partition (#t: eqtype) (parts: seq (seq t)) (s: seq t)
  : Pure prop (requires are_parts_of parts s)
              (ensures fun _ -> True) = no_intersections_among parts /\
                                  (forall (x:t). mem x s ==> x `is_found_in` parts)
                                  //length s == total_length parts

(* auxiliary to establish provable equality of sequences of length zero *)
let empty_lemma (#t:eqtype) (s: seq t) 
  : Lemma (requires seq_is_set s /\ equal_as_sets s empty)
          (ensures s == empty) = 
    if length s > 0 then assert (mem (head s) s)
    else lemma_eq_elim s empty

let rec equal_sets_have_equal_lengths (#t: eqtype) (a b: seq t)
  : Lemma (requires seq_is_set a /\ seq_is_set b /\ equal_as_sets a b)
          (ensures length a == length b)
          (decreases length a) = 
  if length a = 0 then empty_lemma b
  else equal_sets_have_equal_lengths (remove_element a (head a)) (remove_element b (head a)) 

let rec equal_length_subset_equals_as_set (#t: eqtype) (superset subset: seq t)
  : Lemma (requires seq_is_set superset /\ 
                    seq_is_set subset /\ 
                    is_subset_of subset superset /\
                    (length superset == length subset))
          (ensures equal_as_sets superset subset)
          (decreases length subset) = 
  if (length superset > 0) 
  then equal_length_subset_equals_as_set (remove_element superset (head subset))
                                         (remove_element subset (head subset))

(* True if [sub_parts] is a partition of [sub], 
   and [sub] is a subset of [s], 
   given [parts] is a partition of [s]  *)
let is_subpartition (#t: eqtype) (s: seq t) (parts: seq (seq t)) (sub: seq t) (sub_parts: seq (seq t))
  : Pure prop (requires seq_is_set s 
                      /\ seq_is_set sub
                      /\ parts `are_parts_of` s
                      /\ sub_parts `are_parts_of` sub
                      /\ s `is_subset_of` s)
              (ensures fun _ -> True) = 
    is_partition sub_parts sub

(* Couldn't rewrite this in terms of dependent tuples *)
type partition (t: eqtype) = 
  | Partition : (s: seq t) -> 
                (parts: seq (seq t) { seq_is_set s /\ parts `are_parts_of` s /\ is_partition parts s }) -> 
                partition t

(* This function extracts a subpartition of [s `subtract_subset` head pt.parts],
   ensuring total length to be decreased precisely by (length (head pt.parts)),
   thus making the final step to the proof of partition length sum breakdown. *)
#push-options "--split_queries"
let partition_split (#t: eqtype) (pt: partition t)
  : Pure (partition t) (requires length pt.parts > 0)
                       (ensures fun p -> is_subpartition pt.s pt.parts p.s p.parts
                                    /\ (length p.parts == length pt.parts - 1)
                                    /\ (length p.s == length pt.s - length (head pt.parts))) =   
  let ps = subtract_subset pt.s (head pt.parts) in
  let pp = (tail pt.parts) in   
  let aux (i: under (length pp)) (x:t)
    : Lemma (mem x (index pp i) ==> mem x ps) = 
      if mem x (index pp i) then assert (mem x pt.s) 
    in Classical.forall_intro_2 aux; 
  assert (is_subpartition pt.s pt.parts ps pp); 
  assert (length ps == length pt.s - length (head pt.parts));
  Partition ps pp
#pop-options

(* Final result, any partition will have total length 
   of [parts] equal to the length of base set [s] *)
let rec total_length_lemma (#t: eqtype) (parts: seq (seq t)) (s: seq t)
  : Lemma (requires parts `are_parts_of` s 
                  /\ parts `is_partition` s)
          (ensures length s == total_length parts) 
          (decreases length parts) = 
  if length parts > 0 then
    let p = partition_split (Partition s parts) in
    total_length_lemma p.parts p.s;
    assert (total_length parts == length (head parts) + total_length (tail parts));
    assert (length p.s == length s - length (head parts))