Fix a couple bugs here and there, improve Lean code-gen, still WIP

2 files changed, 47 insertions, 26 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index af510a84..44bc5e1c 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -175,8 +175,8 @@ let assumed_variants () : (assumed_ty * VariantId.id * string) list =
       ]
   | Lean ->
       [
-        (Result, result_return_id, "result.ret");
-        (Result, result_fail_id, "result.fail");
+        (Result, result_return_id, "ret");
+        (Result, result_fail_id, "fail");
         (Error, error_failure_id, "panic");
         (* No Fuel::Zero on purpose *)
         (* No Fuel::Succ on purpose *)
@@ -202,7 +202,7 @@ let assumed_llbc_functions :
     (VecIndexMut, rg0, "vec_index_mut_back");
   ]
 
-let assumed_pure_functions : (pure_assumed_fun_id * string) list =
+let assumed_pure_functions (): (pure_assumed_fun_id * string) list =
   match !backend with
   | FStar ->
       [
@@ -218,7 +218,7 @@ let assumed_pure_functions : (pure_assumed_fun_id * string) list =
   | Lean ->
       [
         (Return, "return");
-        (Fail, "fail");
+        (Fail, "fail_");
         (Assert, "massert");
         (* TODO: figure out how to deal with this *)
         (FuelDecrease, "decrease");
@@ -232,7 +232,7 @@ let names_map_init () : names_map_init =
     assumed_structs;
     assumed_variants = assumed_variants ();
     assumed_llbc_functions;
-    assumed_pure_functions;
+    assumed_pure_functions = assumed_pure_functions ();
   }
 
 let extract_unop (extract_expr : bool -> texpression -> unit)
@@ -1414,7 +1414,11 @@ let extract_adt_g_value
        *)
       let cons =
         match variant_id with
-        | Some vid -> ctx_get_variant adt_id vid ctx
+        | Some vid ->
+            if !backend = Lean then
+              ctx_get_type adt_id ctx ^ "." ^ ctx_get_variant adt_id vid ctx
+            else
+              ctx_get_variant adt_id vid ctx
         | None -> ctx_get_struct adt_id ctx
       in
       if inside && field_values <> [] then F.pp_print_string fmt "(";
@@ -1610,7 +1614,11 @@ and extract_adt_cons (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
        *)
       let cons =
         match adt_cons.variant_id with
-        | Some vid -> ctx_get_variant adt_cons.adt_id vid ctx
+        | Some vid ->
+            if !backend = Lean then
+              ctx_get_type adt_cons.adt_id ctx ^ "." ^ ctx_get_variant adt_cons.adt_id vid ctx
+            else
+              ctx_get_variant adt_cons.adt_id vid ctx
         | None -> ctx_get_struct adt_cons.adt_id ctx
       in
       let use_parentheses = inside && args <> [] in
@@ -1725,7 +1733,8 @@ and extract_lets (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
         F.pp_print_space fmt ();
         extract_texpression ctx fmt false re;
         F.pp_print_space fmt ();
-        F.pp_print_string fmt "in";
+        if !backend <> Lean then
+          F.pp_print_string fmt "in";
         ctx)
     in
     (* Close the box for the let-binding *)
@@ -1734,6 +1743,13 @@ and extract_lets (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
     (* Return *)
     ctx
   in
+  (* If Lean, we rely on monadic blocks, so we insert a do and open a new box
+     immediately *)
+  if !backend = Lean then begin
+    F.pp_open_vbox fmt ctx.indent_incr;
+    F.pp_print_string fmt "do";
+    F.pp_print_space fmt ();
+  end;
   let ctx =
     List.fold_left
       (fun ctx (monadic, lv, re) -> extract_let ctx monadic lv re)
@@ -1745,6 +1761,9 @@ and extract_lets (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
   extract_texpression ctx fmt false next_e;
   (* Close the box for the next expression *)
   F.pp_close_box fmt ();
+  (* do-box (Lean only) *)
+  if !backend = Lean then
+    F.pp_close_box fmt ();
   (* Close parentheses *)
   if inside then F.pp_print_string fmt ")";
   (* Close the box for the whole expression *)
@@ -1852,7 +1871,9 @@ and extract_Switch (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
 
       (* End the match *)
       F.pp_print_space fmt ();
-      F.pp_print_string fmt "end");
+      (* Relying on indentation in Lean *)
+      if !backend <> Lean then
+        F.pp_print_string fmt "end");
   (* Close parentheses *)
   if inside then F.pp_print_string fmt ")";
   (* Close the box for the whole expression *)
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index 3ad55d37..152dfc99 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -107,10 +107,10 @@ type type_decl_kind = Enum | Struct
     For instance, provided some information about a function (its basename,
     information about the region group, etc.) it should come up with an
     appropriate name for the forward/backward function.
-    
+
     It can of course apply many transformations, like changing to camel case/
     snake case, adding prefixes/suffixes, etc.
-    
+
     2. Format some specific terms, like constants.
  *)
 type formatter = {
@@ -120,12 +120,12 @@ type formatter = {
   str_name : string;
   type_decl_kind_to_qualif : decl_kind -> type_decl_kind option -> string;
       (** Compute the qualified for a type definition/declaration.
-    
+
           For instance: "type", "and", etc.
        *)
   fun_decl_kind_to_qualif : decl_kind -> string;
       (** Compute the qualified for a function definition/declaration.
-    
+
           For instance: "let", "let rec", "and", etc.
        *)
   field_name : name -> FieldId.id -> string option -> string;
@@ -133,7 +133,7 @@ type formatter = {
           - type name
           - field id
           - field name
-          
+
           Note that fields don't always have names, but we still need to
           generate some names if we want to extract the structures to records...
           We might want to extract such structures to tuples, later, but field
@@ -147,13 +147,13 @@ type formatter = {
        *)
   struct_constructor : name -> string;
       (** Structure constructors are used when constructing structure values.
-      
+
           For instance, in F*:
           {[
             type pair = { x : nat; y : nat }
             let p : pair = Mkpair 0 1
           ]}
-          
+
           Inputs:
           - type name
        *)
@@ -194,7 +194,7 @@ type formatter = {
       (** Generates the name of the definition used to prove/reason about
           termination. The generated code uses this clause where needed,
           but its body must be defined by the user.
-      
+
           Inputs:
           - function id: this is especially useful to identify whether the
             function is an assumed function or a local function
@@ -205,7 +205,7 @@ type formatter = {
        *)
   var_basename : StringSet.t -> string option -> ty -> string;
       (** Generates a variable basename.
-      
+
           Inputs:
           - the set of names used in the context so far
           - the basename we got from the symbolic execution, if we have one
@@ -227,7 +227,7 @@ type formatter = {
        *)
   extract_primitive_value : F.formatter -> bool -> primitive_value -> unit;
       (** Format a constant value.
-      
+
           Inputs:
           - formatter
           - [inside]: if [true], the value should be wrapped in parentheses
@@ -242,7 +242,7 @@ type formatter = {
     texpression ->
     unit;
       (** Format a unary operation
-      
+
           Inputs:
           - a formatter for expressions (called on the argument of the unop)
           - extraction context (see below)
@@ -262,7 +262,7 @@ type formatter = {
     texpression ->
     unit;
       (** Format a binary operation
-      
+
           Inputs:
           - a formatter for expressions (called on the arguments of the binop)
           - extraction context (see below)
@@ -289,7 +289,7 @@ type id =
   | StructId of type_id
       (** We use this when we manipulate the names of the structure
           constructors.
-          
+
           For instance, in F*:
           {[
             type pair = { x: nat; y : nat }
@@ -328,7 +328,7 @@ module IdMap = Collections.MakeMap (IdOrderedType)
 (** The names map stores the mappings from names to identifiers and vice-versa.
 
     We use it for lookups (during the translation) and to check for name clashes.
-    
+
     [id_to_string] is for debugging.
   *)
 type names_map = {
@@ -385,7 +385,7 @@ let names_map_add_function (id_to_string : id -> string) (fid : fun_id)
 
     We do this in an inefficient manner (by testing all indices starting from
     0) but it shouldn't be a bottleneck.
-    
+
     Also note that at some point, we thought about trying to reuse names of
     variables which are not used anymore, like here:
     {[
@@ -394,7 +394,7 @@ let names_map_add_function (id_to_string : id -> string) (fid : fun_id)
       let x0 = ... in // We could use the name "x" if [x] is not used below
       ...
     ]}
-    
+
     However it is a good idea to keep things as they are for F*: as F* is
     designed for extrinsic proofs, a proof about a function follows this
     function's structure. The consequence is that we often end up
@@ -402,7 +402,7 @@ let names_map_add_function (id_to_string : id -> string) (fid : fun_id)
     when calling lemmas) we often need to talk about the "past" (i.e.,
     previous values), it is very useful to generate code where all variable
     names are assigned at most once.
-    
+
     [append]: function to append an index to a string
  *)
 let basename_to_unique (names_set : StringSet.t)