Tips for extraction with polymorphism and erasure

module Extraction_PolymorphismAndErasure
open FStar.Ghost

(* This module exercises the extraction of various polymorphic types,
   including assumed polymorphic values.

   In all cases, assumed values are extracted to OCaml with bodies
   `failwith "Not yet implemented"`

   Meanwhile, in Kremlin, those assumed F* functions that extract to
   polymorphic OCaml functions are discarded altogether.
 *)

//f1:
//produces in OCaml `let f1 : 'Aa . 'Aa -> 'Aa = ...
//so, discarded in Kremlin
assume
val f1 (a:Type) (x:a) : Tot a

//f2:
//produces in OCaml `let f2 : unit -> int -> int = ...`
//so, retained in Kremlin
assume
val f2 (x:erased int) (y:int) : Tot int

//f3:
//produces in OCaml `let f3 : 'Apre . Prims.int -> Prims.int ...`
//The type function is extracted as a type in OCaml, since it can
//in some cases improve typability in OCaml (see examples 6 and 8 below)
//But, this is discarded in Kremlin
assume
val f3 (pre:(int -> Type)) (y:int) : Tot int

//f4:
//produces in OCaml `let f4 : unit -> Prims.int -> Prims.int ...`
//same as before, retained in Kremlin
assume
val f4 (pre:(int -> erased bool)) (y:int) : Tot int

//f5:
//This is a variant of f3, but no longer using prenex quantification
//So, F* cannot extract pre to a type in OCaml, instead erasing it to unit
//We get ``let (f5 : Prims.int -> unit -> Prims.int) =```
//And this is retained in Kremlin
assume
val f5 (y:int) (pre:(int -> Type)) : Tot int

//f6: This is similar to f3, except it uses the dependency in the result
//This extracts to `let f6 : 'Apre . Prims.int -> 'Apre =`
//Hinting at the usefulness of extracting the first argument as a type
//since it also is used to type the result
//Discarded in Kremlin
assume
val f6 (pre:int -> Type) (y:int) : Tot (pre y)

// As in f5, in this case, we no longer have prenex quantificaion,
// so, we get: let (f7 : Prims.int -> unit -> Obj.t)
// which is not polymorphic but not sure how Kremlin extracts Obj.t
assume
val f7 (y:int) (pre:int -> Type) : Tot (pre y)

// Now to illustrate why retaining a type function is useful

// The next three lines are extracted to
// ```
// let f9 : 'Apre . Prims.int -> 'Apre -> 'Apre = fun y  -> fun z  -> z
// type 'Ax t9 = Prims.int
// let (test9 : unit t9) = f9 Prims.int_zero Prims.int_zero
// ```
// Note, we avoid the use of any magics
let f9 (pre:int -> Type) (y:int) (z:pre y) : Tot (pre y) = z
let t9 (x:int) = y:int{y = x}
let test9 = f9 t9 0 0


// when the type function is erased (as in f10 and f5), we
// get magics in extraction
// ```
// let (f10 : unit -> unit -> Prims.int -> Obj.t -> Obj.t) =
//  fun uu____1556  -> fun pre  -> fun y  -> fun z  -> z
// let (test10 : unit t9) =
//   Obj.magic (f10 () () Prims.int_zero (Obj.magic Prims.int_zero))
// ```
let f10 (_:unit) (pre:int -> Type) (y:int) (z:pre y) : Tot (pre y) = z
let test10 = f10 () t9 0 0

// Note, this only really helps for type functions dependent on values
// whose application results in an well-formed Ocaml type
// For instnace, this kind of higher-rank polymorphism still produces
// magics in OCaml
// ```
// let f11 : 'Apre 'Aa . 'Apre -> 'Apre = fun z  -> z
// let (test11 : Prims.int Prims.list) = Obj.magic (f11 (Obj.magic []))
// ```
let f11 (pre:Type -> Type) (a:Type) (z:pre a) : Tot (pre a) = z
let test11 = f11 list int []

(*
 In summary, this suggests a style where if you want a polymorphic assume
 to persist when extracted to Kremlin, it may be effective to not use
 prenex quantification. For example, this function from FStar.Bytes


`
val of_buffer (#p #q:_) (l:UInt32.t) (buf:B.mbuffer UInt8.t p q{B.length buf == U32.v l})
  : Stack (b:bytes{length b = UInt32.v l})
`

 will not survive extraction to Kremlin, since the two preorder
 arguments will be extracted as types.

 Instead, at least in this case, a trick is to write its type as

`
val of_buffer (l:UInt32.t) (#p #q:_) (buf:B.mbuffer UInt8.t p q{B.length buf == U32.v l})
  : Stack (b:bytes{length b = UInt32.v l})
`
 which will result in p and q showing up as unit arguments in both
 OCaml and Kremlin.

*)