Simple key-based search in a seq of key-value pairs

Sandbox.fst

module Sandbox

open FStar.Seq

let under (x: nat) = z:nat{z<x}

let index_for #t (x:seq t) = under (length x)

let rec select_by_key (#key_type: eqtype) 
                      (#value_type: Type) 
                      (key: key_type) 
                      (set: seq (key_type * value_type))
  : Pure (seq value_type) 
         (requires True)
         (ensures fun l -> ( // if a key is found in set, 
                          // then the value should appear in result
                       forall (i: index_for set). 
                           ((fst (index set i) = key) 
                           ==> (exists (j: index_for l). index l j == snd (index set i)))))
         (decreases length set)
  = if length set = 0 
    then empty
    else match (head set) with 
         | (x, v) -> // we isolate the recursion in one place
                    let subquery = select_by_key key (tail set) in 
                    // we have the (ensures) clause in form of lemma 
                    let aux (i: index_for (tail set))
                      : Lemma (requires (fst (index (tail set) i) = key)) 
                              (ensures exists (j: index_for subquery).
                                         index subquery j == snd (index (tail set) i)) 
                      = () in // which F* automatically proves of course
                    if key = x 
                    then begin // here's the tricky part: we link tail indexing 
                               // expressions with current seq indexing ones
                           let joined = cons v subquery in // return value to reason about
                           let aux_full (i: index_for set) // (post-condition for return value)
                             : Lemma (requires (fst (index set i) = key))
                                     (ensures exists (j: index_for joined).
                                                index joined j == snd (index set i)) 
                             // Even though I dislike assertions in finalized proofs,
                             // this assertion actually helps F* to prove the post-condition.
                             = if i = 0 then assert (index joined 0 == snd (index set 0))
                               else begin
                                 aux (i-1); // we obtain the exists statement from recursion
                                 eliminate exists (j: index_for subquery). // and use the new F* sugar
                                             index subquery j == snd (index (tail set) (i-1))
                                 returns exists (k: index_for joined).
                                           index joined k == snd (index set i) 
                                 with _. index_cons_r v subquery (j+1)
                                 // Note how we can't get away with just 
                                 // invoking forall_intro (index_cons_r v subquery) instead...
                               end 
                             in Classical.forall_intro (Classical.move_requires aux_full);
                             joined // ...but to ensure post-condition, we invoke it on aux_full.
                         end
                    else begin // if the first pair's key does not match our search,
                               // things get simplier and we can just get away with 
                               // just aux (i-1) since we know for sure that (i > 0).
                           let aux_full (i: index_for set) 
                             : Lemma (requires (fst (index set i) = key))
                                     (ensures exists (j: index_for subquery).
                                                index subquery j == snd (index set i)) 
                             = if i > 0 then aux (i-1) 
                             in Classical.forall_intro (Classical.move_requires aux_full);
                           subquery
                         end