val alloc (#a: eqtype) (#b #inv: _) (r: erid)
    : ST (MDM.t r a b inv)
      (requires
        (fun h -> inv (MDM.empty_partial_dependent_map #a #b) /\ witnessed (region_contains_pred r))
      )
      (ensures
        (fun h0 x h1 ->
            inv (MDM.empty_partial_dependent_map #a #b) /\ ralloc_post r (MDM.empty #a #b) h0 x h1))
val alloc (#a: eqtype) (#b #inv: _) (r: erid)
    : ST (MDM.t r a b inv)
      (requires
        (fun h -> inv (MDM.empty_partial_dependent_map #a #b) /\ witnessed (region_contains_pred r))
      )
      (ensures
        (fun h0 x h1 ->
            inv (MDM.empty_partial_dependent_map #a #b) /\ ralloc_post r (MDM.empty #a #b) h0 x h1))
let alloc (#a:eqtype) #b #inv (r: erid):
  ST (MDM.t r a b inv)
    (requires (fun h -> 
      inv (MDM.empty_partial_dependent_map #a #b) /\ 
      witnessed (region_contains_pred r) ))
    (ensures (fun h0 x h1 ->
      inv (MDM.empty_partial_dependent_map #a #b) /\
      ralloc_post r (MDM.empty #a #b) h0 x h1))
  = MDM.alloc #a #b #inv #r ()
module MiTLS.KDF
open MiTLS

/// We define a KDF parametric in both its usage and ipkg
/// We rely on type abstraction to separate secrets with different indexes
/// For soundness, we must also prevent external calls to create derived secrets.
///
/// We idealize KDF as a random function, with lazy sampling.
/// This holds syntactically only after inlining all calls to
/// pkg.create/coerce.
///
/// This requires memoizing calls, which is a bit tricky when calling
/// stateful allocators.

open MiTLS.Mem
open MiTLS.Pkg
open MiTLS.Idx 
open MiTLS.Pkg.Tree

module MDM = FStar.Monotonic.DependentMap
module HS = FStar.HyperStack

inline_for_extraction
let sample (len:UInt32.t): ST (lbytes32 len)
    (requires fun h0 -> True)
    (ensures fun h0 r h1 -> h0 == h1)
  = assume false; Random.sample32 len

type details (ha:kdfa) =
| Log:
  loginfo: TLSInfo.logInfo ->
  hv: Hashing.Spec.anyTag {digest_info ha loginfo hv} -> details ha

type info =
| Info:
  ha:kdfa ->
  option (details ha) -> info

noextract
let rec index_info (i:id{model}) =
  let i':pre_id = i in
  match i' with
  | Preshared a _ -> Info a None
  | Derive i l (ExpandLog log hv) -> Info (ha_of_id i) (Some (Log log hv))
  | Derive i _ _ -> index_info i

let valid_info (i0:id) (v:info) =
  if model then v == index_info i0 else True

type info0 (i:id) = u:info{valid_info i u}
type iflag = b:bool{b ==> model}
type usage (ideal:iflag) = children (b2t ideal)

unfold type regid = i:(Idx?.t ii){(Idx?.registered ii) i}

let derived_key
  (#ideal:iflag)
  (u: usage ideal)
  (i: regid)
  (lbl: label {u `has_lbl` lbl})
  (ctx: context)
  : Pure Type0
  (requires wellformed_derive i lbl ctx /\
    registered (derive i lbl ctx) /\ model)
//    (model ==> Pkg?.key (child u lbl) == ktype))
  (ensures fun t -> True)
  =
  Pkg?.key (child u lbl) (Derive i lbl ctx)
//  (Pkg?.key (child u lbl)) (Derive i lbl ctx)

let kdf_tables_region : tls_rgn = new_region Mem.tls_tables_region

/// key-derivation table (memoizing create/coerce)
noextract
type domain (#ideal:iflag) (u:usage ideal) (i:regid) : eqtype =
| Domain:
  lbl: label {u `has_lbl` lbl} ->
  ctx: context {wellformed_derive i lbl ctx /\ registered (derive i lbl ctx)} ->
  domain u i

noextract
let kdf_range (#ideal:iflag) (u:usage ideal) (i:regid{model}) (d:domain u i) =
  (let Domain lbl ctx = d in derived_key u i lbl ctx)

noextract
noeq private type table (#ideal:iflag) (u:usage ideal) (i:regid) (s:squash model) =
  | KDF_table:
    r:subrgn kdf_tables_region ->
    t:MDM.t r (domain u i) (kdf_range u i) (fun _ -> True) ->
    table u i s

inline_for_extraction
let secret_len (a:info) : keylen = Hacl.Hash.Definitions.hash_len a.ha
type real_secret (i:regid) = a:info0 i & lbytes32 (secret_len a)

let safe (#ideal:iflag) (u:usage ideal) (i:regid) =
  honest i /\ b2t ideal

type secret (#ideal:iflag) (u:usage ideal) (i:regid) =
  Model.ir (safe u) (table u i) (real_secret i) i

/// Real KDF schemes, parameterized by an algorithm computed from the
/// index (existing code).

let lemma_honest_parent (i:regid) (lbl:label) (ctx:context)
  : Lemma (requires wellformed_derive i lbl ctx /\ registered (derive i lbl ctx) /\ ~(honest_idh ctx))
  (ensures honest (derive i lbl ctx) ==> honest i)
  = admit() // hard, uses the monotonic excluded middle axiom and table invariant

inline_for_extraction
let get_info (#ideal:iflag) (#u:usage ideal) (#i:regid) (k:secret u i) =
  if Model.is_safe k then index_info i
  else dfst (Model.real k)

// For KDF, we require that being fresh in the KDF table implies being fresh
// in the derived package's definition table
#set-options "--z3rlimit 30"
type local_kdf_invariant (#ideal:iflag) (#u:usage ideal) (#i:regid) (k:secret u i) (h:mem) =
  (forall (lbl:label {u `has_lbl` lbl}) (ctx':context).
    (~(honest_idh ctx') /\ wellformed_derive i lbl ctx' /\ registered (derive i lbl ctx')) ==>
    (if model then
      let ctx : c:context{wellformed_derive i lbl c /\ registered (derive i lbl c)} = ctx' in
      let i' : regid = derive i lbl ctx in
      lemma_honest_parent i lbl ctx;
      let pkg' = child u lbl in
      // Nice: we get this from the tree structure for free!
      assert(Pkg?.ideal pkg' ==> b2t ideal);
      let dt = itable pkg'.define_table in
      match Model.is_safe k with
      | true ->
        // the entries in the KDF table match those of the child's define_table
        let KDF_table r t : table u i () = Model.ideal k () in
        MDM.sel (sel h t) (Domain lbl ctx) == MDM.sel (sel h dt) i'
      | false ->
        let (| a, raw |) = Model.real k in
        assert(~(safe u i));
        assert(honest i' ==> honest i);
        assert(Pkg?.ideal pkg' ==> ~(honest i')); // to call coerceT
        // the child's define table has correctly-computed coerced entries
        (match MDM.sel (sel h dt) i' with
        | None -> True
        | Some k' -> exists (a': (Pkg?.info pkg') i').
          // we recompute the concrete key materials, and recall the
          // existence of some prior info, which (by specification
          // of coerceT) must yield exactly the same instance as the
          // one recorded earlier.
          let i'': i:regid{Pkg?.ideal pkg' ==> ~(honest i)} = i' in
	  let len' = Pkg?.len pkg' a' in
          let lb = FStar.Bytes.bytes_of_string lbl in
          let raw' = HKDF.expand_spec #a.ha raw lb len' in
          k' == Pkg?.coerceT pkg' i'' a' raw')
    else True))
#reset-options

noextract
let kdf_shared_footprint (#ideal:iflag) (u:usage ideal) : rset =
  assume false; Set.empty
  // List.Tot.fold_left (Set.empty) (fun s p -> rset_union s p.define_region) u

let kdf_footprint (#ideal:iflag) (#u:usage ideal) (#i:regid) (k:secret u i)
  : GTot (s:rset{s `Set.disjoint` kdf_shared_footprint u}) =
  assume false; // WIP(adl) Fix rset for define_region subrgn
  if Model.is_safe k then
    let KDF_table r _ = Model.ideal k () in
    Set.singleton r
  else Set.empty

let local_kdf_invariant_framing (#ideal:iflag) (#u:usage ideal) (i:regid) (k:secret u i) (h0:mem) (r:rid) (h1:mem)
  : Lemma (requires local_kdf_invariant k h0 /\ modifies_one r h0 h1
            /\ ~(r `Set.mem` kdf_footprint k) /\ ~(r `Set.mem` kdf_shared_footprint u))
          (ensures local_kdf_invariant k h1)
  = admit()

/// maybe reverse-inline sampling from low-level KeyGen?
/// otherwise we have to argue it is what Create does.
///
/// MK: what does reverse-inline of low-level KeyGen mean?

// The post-condition of creating a KDF is that its table is empty
// This is useful to re-establish the multi-KDF invariant
type kdf_post (#ideal:iflag) (#u:usage ideal) (#i:regid)
  (a:info0 i) (k:secret u i) (h:mem) =
  (safe u i ==>
    (let KDF_table r t = Model.ideal k () in
     sel h t == MDM.empty #(domain u i) #(kdf_range u i)))

// Framing for the kdf_post depends only on kdf_footprint k
let kdf_post_framing (#ideal:iflag) (#u:usage ideal) (#i:regid) (a:info0 i)
  (k:secret u i) (h0:mem) (r:rid) (h1:mem)
  : Lemma (requires (kdf_post a k h0 /\ modifies_one r h0 h1 /\ ~(r `Set.mem` kdf_footprint k)))
          (ensures (kdf_post a k h1))
  = admit()

val coerceT:
  #ideal: iflag ->
  u: usage ideal ->
  i: regid{~(safe u i)} ->
  a: info0 i ->
  repr: lbytes32 (secret_len a) ->
  GTot (secret u i)
let coerceT #ideal u i a repr =
  Model.mk_real (| a, repr |)

val coerce:
  #ideal: iflag ->
  u: usage ideal ->
  i: regid{~(safe u i)} ->
  a: info0 i ->
  repr: lbytes32 (secret_len a) ->
  ST (secret u i)
  (requires fun h0 -> True)
  (ensures fun h0 k h1 -> modifies_none h0 h1
    /\ k == coerceT u i a repr
    /\ fresh_regions (kdf_footprint k) h0 h1
    /\ kdf_post a k h1 /\ local_kdf_invariant k h1)

#reset-options "--z3rlimit 100" 
let coerce #ideal u i a repr =
  let k = Model.mk_real (| a, repr |) in
  let h1 = get() in
  // WIP stronger packaging
  (if model then assume(local_kdf_invariant k h1));
  k

/// NS:
/// MDM.alloc is a stateful function with all implicit arguments
/// F* will refuse to instantiate all those arguments, since implicit
/// instantiation in F* should never result in an effect.
///
/// Stateful functions always take at least 1 concrete argument.
///
/// I added a unit here
///
/
prims.fst.checked MiTLS.TLSInfo.fst.checked MiTLS.Random.fst.checked MiTLS.Pkg.Tree.fst.checked MiTLS.Pkg.fst.checked MiTLS.Model.fst.checked MiTLS.Mem.fst.checked MiTLS.Idx.fst.checked MiTLS.HKDF.fsti.checked MiTLS.Hashing.Spec.fst.checked Hacl.Hash.Definitions.fst.checked FStar.UInt32.fsti.checked FStar.Set.fsti.checked FStar.Pervasives.fsti.checked FStar.Monotonic.DependentMap.fsti.checked FStar.HyperStack.fst.checked FStar.Bytes.fsti.checked
MiTLS.KDF.fst
r: FStar.HyperStack.ST.erid -> FStar.HyperStack.ST.ST (FStar.Monotonic.DependentMap.t r a b inv)
Prims.eqtype FStar.DependentMap.t FStar.Monotonic.DependentMap.opt FStar.HyperStack.ST.erid FStar.Monotonic.DependentMap.alloc FStar.Monotonic.DependentMap.t FStar.Monotonic.HyperStack.mem Prims.l_and FStar.Monotonic.DependentMap.empty_partial_dependent_map FStar.HyperStack.ST.witnessed FStar.HyperStack.ST.region_contains_pred FStar.HyperStack.ST.ralloc_post FStar.Monotonic.DependentMap.imap FStar.Monotonic.DependentMap.grows FStar.Monotonic.DependentMap.empty
let alloc (#a: eqtype) #b #inv (r: erid)
    : ST (MDM.t r a b inv)
      (requires
        (fun h -> inv (MDM.empty_partial_dependent_map #a #b) /\ witnessed (region_contains_pred r))
      )
      (ensures
        (fun h0 x h1 ->
            inv (MDM.empty_partial_dependent_map #a #b) /\ ralloc_post r (MDM.empty #a #b) h0 x h1)) =
  
MDM.alloc #a #b #inv #r ()