fmap sandbox

Sandbox.fst

module Sandbox

open FStar.List
open FStar.Set

type nat_from (l: set nat) = x:nat{mem x l}

type nonempty_set a = l:set a { l =!= empty }

let list_len = FStar.List.Tot.Base.length 

type nonempty_list a = l:list a{ list_len l > 0 }

/// Tot version of nth. Fail-safe, total, pure and all. 
let rec at #a (l: list a{list_len l > 0}) (i: nat{i < list_len l}) : Tot a = 
  match l with 
  | h::t -> if i=0 then h else at t (i-1)

noeq
type fmap a = 
  | Empty : fmap a
  | FMap : dom : nonempty_set nat -> (f : (nat_from dom -> a)) -> fmap a

let nonempty_fmap a = (m:fmap a{m =!= Empty})

let dom_of #a (m: fmap a) = match m with
  | Empty -> empty
  | FMap dom _ -> dom

type keys_of #a (m: fmap a) = nat_from (dom_of m)

let ( @> ) #a (m: nonempty_fmap a) (i: keys_of m) = 
  match m with 
  | FMap dom f -> f i

let is_atomic_update_of #a (original updated: nonempty_fmap a) (i: keys_of original) (new_value: a) =
  dom_of original == dom_of updated /\
  (forall (k: keys_of original). 
     (updated @> k) == (if i=k then new_value else (original @> k)))

type atomic_update_of #a (original: nonempty_fmap a) (i: keys_of original) (new_value: a)
  = updated: nonempty_fmap a { is_atomic_update_of original updated i new_value }

let fmap_update #a (m: nonempty_fmap a) (i: keys_of m) (v:a)
  : Tot (atomic_update_of m i v) = 
  match m with 
  | Empty -> m
  | FMap dom f -> FMap dom (fun (k: nat_from dom) -> (if i=k then v else f k))

let is_batch_update_of #a (original updated: nonempty_fmap a) (new_keys: set (keys_of original)) (new_value: a) =
  dom_of updated == dom_of original /\ 
  (forall (i: keys_of original). updated @> i == (if (mem i new_keys) then new_value else original @> i))

type batch_update_of #a (original: nonempty_fmap a) (new_keys: set (keys_of original)) (new_value: a)
 = updated: nonempty_fmap a { is_batch_update_of original updated new_keys new_value }

let fmap_update_batch #a (m: nonempty_fmap a) (l: set (keys_of m)) (v: a)
  : Tot (batch_update_of m l v)  =  
  FMap (dom_of m) (fun i -> if mem i l then v else m @> i) 

let fmap_init #a (indices: nonempty_set nat) (value: a) :
  (result: fmap a{ dom_of result == indices /\ 
                  (forall (i: keys_of result). result @> i == value)}) = 
  FMap indices (fun (x:nat_from indices) -> value) 

let singleton_is_not_empty (#a:eqtype) (x:a) : Lemma (singleton x =!= empty) = 
  mem_empty x  

let fmap_add_or_update #a (m: fmap a) (i: nat) (v:a) : z:nonempty_fmap a { 
  (dom_of z == union (dom_of m) (singleton i)) /\
  (forall (k: keys_of z). z @> k == (if k=i then v else m @> k) )
} = 
  let new_indices = add i (dom_of m) in 
  mem_union i new_indices (dom_of m) ;
  FMap new_indices (fun (k: nat_from new_indices) -> if k=i then v else m @> k)

let at_cons_lemma #a (h:a) (t: list a) (i:nat{i<=length t}) 
  : Lemma (at (h::t) i == (if i=0 then h else at t (i-1))) = ()

let rec keys_of_pairs #a (mappings: nonempty_list (nat * a)) : z:list nat{
  length mappings == length z /\
  (forall (i: nat{i<length mappings}). fst (at mappings i) == at z i) 
} = match mappings with
  | [] -> []
  | [h] -> [fst h]
  | h::t -> Classical.forall_intro (at_cons_lemma h t);
          Cons (fst h) (keys_of_pairs t) 

let fmap_initializer #a (mappings: list (nat*a)) 
  : m:fmap a  { list_len mappings > 0 ==> m =!= Empty }
  =
  if list_len mappings = 0 then Empty else
  let rec aux (l: list (nat*a)) (map: fmap a)
  : (m: fmap a{ list_len l > 0 ==> m =!= Empty /\
       dom_of m == as_set (keys_of_pairs l) `union` dom_of map
  })
  = 
    match l with 
    | [] -> map
    | [(i,v)] -> singleton_is_not_empty i;
                assert (keys_of_pairs l == [i]);
                assert (as_set (keys_of_pairs l) == union (singleton i) empty);
                let next_dom = union (singleton i) (dom_of map) in
                mem_union i (singleton i) (dom_of map);
                mem_empty i;
                
                assert (as_set (keys_of_pairs l) == singleton i);
                admit();
                FMap (singleton i `union` dom_of map) (fun k -> if k=i then v else map @> i)
    | (i,v)::t -> 
                mem_union i (singleton i) (dom_of (aux t (fmap_add_or_update map i v)));
//                lemma_equal_intro (dom_of (aux t (fmap_add_or_update map i v)))
  //                                (singleton i `union` 
                aux t (fmap_add_or_update map i v) in
  aux mappings Empty

let test = fmap_initializer [ (0, -1); (3, 1); (4, 5) ]

let _ = assert (test @> 0 == -1)

let rec fmap_update1 #a (m: fmap a) (l: list nat) (v:a)
  : Tot (result:fmap a{ result._dom == m._dom }) = 
 match l with
 | []     -> m
 | h :: t  -> let m' = { m with _f = (fun y -> if h = y then v else m._f y)} in 
             fmap_update1 m t v

let rec fmap_update_does_nothing #a (m: fmap a) (l: list nat) (v: a)
  : Lemma (fmap_update1 m l v == m) = 
  match l with 
  | [] -> ()
  | h::t -> fmap_update_does_nothing m t v

let rec fmap_update_lemma #a (m: fmap a) (l: list nat) (v:a)
   : Lemma (ensures (forall i. mem i l ==>  (fmap_update1 m l v)._f i ==  v)) =
 match l with
 | []     -> ()
 | h :: t  -> let m' = { m with _f = (fun y -> if h = y then v else m._f y)} in 
             fmap_update_lemma m t v