Lean: preserve more types from sail

1. Preserve the let type annotations from the sail source
2. Preserve the parameters of types

src/sail_lean_backend/pretty_print_lean.ml

@@ -57,15 +57,16 @@ let is_enum env id = match Env.lookup_id id env with Enum _ -> true | _ -> false
 
 let pat_is_plain_binder env (P_aux (p, _)) =
   match p with
-  | (P_id id | P_typ (_, P_aux (P_id id, _))) when not (is_enum env id) -> Some (Some id)
-  | P_wild | P_typ (_, P_aux (P_wild, _)) -> Some None
-  | P_var (_, _) -> Some (Some (Id_aux (Id "var", Unknown)))
-  | P_app (_, _) -> Some (Some (Id_aux (Id "app", Unknown)))
-  | P_vector _ -> Some (Some (Id_aux (Id "vect", Unknown)))
-  | P_tuple _ -> Some (Some (Id_aux (Id "tuple", Unknown)))
-  | P_list _ -> Some (Some (Id_aux (Id "list", Unknown)))
-  | P_cons (_, _) -> Some (Some (Id_aux (Id "cons", Unknown)))
-  | P_lit _ -> Some (Some (Id_aux (Id "lit", Unknown)))
+  | P_id id when not (is_enum env id) -> Some (Some id, None)
+  | P_typ (typ, P_aux (P_id id, _)) when not (is_enum env id) -> Some (Some id, Some typ)
+  | P_wild | P_typ (_, P_aux (P_wild, _)) -> Some (None, None)
+  | P_var (_, _) -> Some (Some (Id_aux (Id "var", Unknown)), None)
+  | P_app (_, _) -> Some (Some (Id_aux (Id "app", Unknown)), None)
+  | P_vector _ -> Some (Some (Id_aux (Id "vect", Unknown)), None)
+  | P_tuple _ -> Some (Some (Id_aux (Id "tuple", Unknown)), None)
+  | P_list _ -> Some (Some (Id_aux (Id "list", Unknown)), None)
+  | P_cons (_, _) -> Some (Some (Id_aux (Id "cons", Unknown)), None)
+  | P_lit _ -> Some (Some (Id_aux (Id "lit", Unknown)), None)
   | _ -> None
 
 (* Copied from the Coq PP *)
@@ -171,9 +172,9 @@ and doc_nconstraint ctx (NC_aux (nc, _)) =
   match nc with
   | NC_and (n1, n2) -> flow (break 1) [doc_nconstraint ctx n1; string "∧"; doc_nconstraint ctx n2]
   | NC_or (n1, n2) -> flow (break 1) [doc_nconstraint ctx n1; string "∨"; doc_nconstraint ctx n2]
-  | NC_equal (a1, a2) -> flow (break 1) [doc_typ_arg ctx a1; string "="; doc_typ_arg ctx a2]
-  | NC_not_equal (a1, a2) -> flow (break 1) [doc_typ_arg ctx a1; string "≠"; doc_typ_arg ctx a2]
-  | NC_app (f, args) -> doc_id_ctor f ^^ parens (separate_map comma_sp (doc_typ_arg ctx) args)
+  | NC_equal (a1, a2) -> flow (break 1) [doc_typ_arg ctx `All a1; string "="; doc_typ_arg ctx `All a2]
+  | NC_not_equal (a1, a2) -> flow (break 1) [doc_typ_arg ctx `All a1; string "≠"; doc_typ_arg ctx `All a2]
+  | NC_app (f, args) -> doc_id_ctor f ^^ parens (separate_map comma_sp (doc_typ_arg ctx `All) args)
   | NC_false -> string "false"
   | NC_true -> string "true"
   | NC_ge (n1, n2) -> flow (break 1) [doc_nexp ctx n1; string "≥"; doc_nexp ctx n2]
@@ -190,8 +191,13 @@ and doc_nconstraint ctx (NC_aux (nc, _)) =
         ]
   | NC_var ki -> doc_kid ctx ki
 
-and doc_typ_arg ctx (A_aux (t, _)) =
-  match t with A_typ t -> doc_typ ctx t | A_nexp n -> doc_nexp ctx n | A_bool nc -> parens (doc_nconstraint ctx nc)
+and doc_typ_arg ctx rel (A_aux (t, _)) =
+  match t with
+  | A_typ t -> doc_typ ctx t
+  | A_nexp n -> doc_nexp ctx n
+  | A_bool nc -> (
+      match rel with `Only_relevant -> empty | `All -> parens (doc_nconstraint ctx nc)
+    )
 
 and provably_nneg ctx x = Type_check.prove __POS__ ctx.env (nc_gteq x (nint 0))
 
@@ -219,7 +225,7 @@ and doc_typ ctx (Typ_aux (t, _) as typ) =
   | Typ_app (Id_aux (Id "range", _), [A_aux (A_nexp low, _); A_aux (A_nexp high, _)]) ->
       if provably_nneg ctx low then string "Nat" else string "Int"
   | Typ_var kid -> doc_kid ctx kid
-  | Typ_app (id, _) -> doc_id_ctor id
+  | Typ_app (id, args) -> doc_id_ctor id ^^ space ^^ separate_map space (doc_typ_arg ctx `Only_relevant) args
   | Typ_exist (_, _, typ) -> doc_typ ctx typ
   | _ -> failwith ("Type " ^ string_of_typ_con typ ^ " " ^ string_of_typ typ ^ " not translatable yet.")
 
@@ -452,17 +458,18 @@ and doc_exp (as_monadic : bool) ctx (E_aux (e, (l, annot)) as full_exp) =
       else wrap_with_pure as_monadic (parens (separate space [doc_exp false ctx e; colon; doc_typ ctx typ]))
   | E_tuple es -> wrap_with_pure as_monadic (parens (separate_map (comma ^^ space) d_of_arg es))
   | E_let (LB_aux (LB_val (lpat, lexp), _), e) ->
-      let id =
+      let id_typ =
         match pat_is_plain_binder env lpat with
-        | Some (Some (Id_aux (Id id, _))) -> id
-        | Some None -> "x" (* TODO fresh name or wildcard instead of x *)
+        | Some (Some (Id_aux (Id id, _)), Some typ) -> string id ^^ space ^^ colon ^^ space ^^ doc_typ ctx typ
+        | Some (Some (Id_aux (Id id, _)), None) -> string id
+        | Some (None, _) -> string "x" (* TODO fresh name or wildcard instead of x *)
         | _ -> failwith "Let pattern not translatable yet."
       in
       let decl_val =
         if effectful (effect_of lexp) then [string "←"; string "do"; doc_exp true ctx lexp]
         else [coloneq; doc_exp false ctx lexp]
       in
-      nest 2 (flow (break 1) ([string "let"; string id] @ decl_val)) ^^ hardline ^^ doc_exp as_monadic ctx e
+      nest 2 (flow (break 1) ([string "let"; id_typ] @ decl_val)) ^^ hardline ^^ doc_exp as_monadic ctx e
   | E_internal_return e -> doc_exp false ctx e (* ??? *)
   | E_struct fexps ->
       let args = List.map d_of_field fexps in
@@ -523,8 +530,8 @@ let doc_funcl_init global (FCL_aux (FCL_funcl (id, pexp), annot)) =
     pats
     |> List.map (fun (pat, typ) ->
            match pat_is_plain_binder env pat with
-           | Some (Some id) -> (id, typ)
-           | Some None -> (Id_aux (Id "x", l), typ) (* TODO fresh name or wildcard instead of x *)
+           | Some (Some id, _) -> (id, typ)
+           | Some (None, _) -> (Id_aux (Id "x", l), typ) (* TODO fresh name or wildcard instead of x *)
            | _ -> failwith "Argument pattern not translatable yet."
        )
   in
@@ -612,11 +619,15 @@ let doc_typdef ctx (TD_aux (td, tannot) as full_typdef) =
       nest 2 (flow (break 1) [string "abbrev"; string id; coloneq; doc_typ ctx t])
   | TD_abbrev (Id_aux (Id id, _), tq, A_aux (A_nexp ne, _)) ->
       nest 2 (flow (break 1) [string "abbrev"; string id; colon; string "Int"; coloneq; doc_nexp ctx ne])
-  | TD_variant (Id_aux (Id id, _), typq, ar, _) ->
+  | TD_variant (Id_aux (Id id, _), tq, ar, _) ->
       let pp_tus = concat (List.map (fun tu -> hardline ^^ doc_type_union ctx tu) ar) in
-      nest 2 (nest 2 (flow space [string "inductive"; string id; string "where"]) ^^ pp_tus)
-      ^^ hardline ^^ hardline
-      ^^ flow space [string "open"; string id]
+      let rectyp = doc_typ_quant_relevant ctx tq in
+      let rectyp = List.map (fun d -> parens d) rectyp in
+      let decl_start =
+        if List.is_empty rectyp then [string "inductive"; string id; string "where"]
+        else [string "inductive"; string id; separate space rectyp; string "where"]
+      in
+      nest 2 (nest 2 (flow space decl_start) ^^ pp_tus) ^^ hardline ^^ hardline ^^ flow space [string "open"; string id]
   | _ -> failwith ("Type definition " ^ string_of_type_def_con full_typdef ^ " not translatable yet.")
 
 (* Copied from the Coq PP *)

test/lean/let.expected.lean

@@ -8,6 +8,10 @@ def foo (lit : Unit) : (BitVec 16) :=
   let z := (HOr.hOr (0xFFFF : (BitVec 16)) (0xABCD : (BitVec 16)))
   (HAnd.hAnd (0x0000 : (BitVec 16)) z)
 
+def bar (lit : Unit) : (BitVec 16) :=
+  let z : (BitVec 16) := (HOr.hOr (0xFFFF : (BitVec 16)) (0xABCD : (BitVec 16)))
+  (HAnd.hAnd (0x0000 : (BitVec 16)) z)
+
 def initialize_registers (lit : Unit) : Unit :=
   ()
 

test/lean/let.sail

@@ -7,3 +7,7 @@ function foo() -> bits(16) = {
   0x0000 & z
 }
 
+function bar() -> bits(16) = {
+  let z : bits(16) = 0xFFFF | 0xABCD in
+  0x0000 & z
+}

test/lean/union.expected.lean

@@ -15,7 +15,7 @@ inductive shape where
 
 open shape
 
-inductive my_option where
+inductive my_option (k_a : Type) where
   | MySome (_ : k_a)
   | MyNone (_ : Unit)
 
@@ -31,11 +31,15 @@ def undefined_circle (lit : Unit) : SailM circle := do
   (pure { radius := sorry })
 
 /-- Type quantifiers: k_a : Type -/
-def is_none (opt : my_option) : Bool :=
+def is_none (opt : my_option k_a) : Bool :=
   match opt with
   | MySome _ => false
   | MyNone () => true
 
+/-- Type quantifiers: k_a : Type -/
+def use_is_none (opt : my_option k_a) : Bool :=
+  (is_none opt)
+
 def initialize_registers (lit : Unit) : Unit :=
   ()
 

test/lean/union.sail

@@ -22,3 +22,8 @@ function is_none opt = match opt {
   MyNone()  => true
 }
 
+val use_is_none : forall ('a : Type). my_option('a) -> bool
+
+function use_is_none opt = {
+  is_none(opt)
+}