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.

*)

This suggests maybe adding an attribute to some type functions so that they will be unconditionally erased in extraction to unit, rather than surviving as a type variable. In this case, FStar.Preorder.preorder could be tagged with the attribute. This may solve those extraction issues for those functions where there are not enough arguments to be permuted with those type functions (infamously, LowParse.Slice.slice)