Improve the formatting of the generated code

7 files changed, 107 insertions, 69 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index f9c4d10a..9aaf753e 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -1203,8 +1203,7 @@ let rec extract_texpression (ctx : extraction_ctx) (fmt : F.formatter)
   | Qualif _ ->
       (* We use the app case *)
       extract_App ctx fmt inside e []
-  | Let (monadic, lv, re, next_e) ->
-      extract_Let ctx fmt inside monadic lv re next_e
+  | Let (_, _, _, _) -> extract_lets ctx fmt inside e
   | Switch (scrut, body) -> extract_Switch ctx fmt inside scrut body
   | Meta (_, e) -> extract_texpression ctx fmt inside e
 
@@ -1394,45 +1393,60 @@ and extract_Abs (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
   (* Close the box for the abs expression *)
   F.pp_close_box fmt ()
 
-and extract_Let (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
-    (monadic : bool) (lv : typed_pattern) (re : texpression)
-    (next_e : texpression) : unit =
+and extract_lets (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
+    (e : texpression) : unit =
+  let lets, next_e = destruct_lets e in
   (* Open a box for the whole expression *)
   F.pp_open_hvbox fmt 0;
   (* Open parentheses *)
   if inside then F.pp_print_string fmt "(";
-  (* Open a box for the let-binding *)
-  F.pp_open_hovbox fmt ctx.indent_incr;
+  (* Extract the let-bindings *)
+  let extract_let (ctx : extraction_ctx) (monadic : bool) (lv : typed_pattern)
+      (re : texpression) : extraction_ctx =
+    (* Open a box for the let-binding *)
+    F.pp_open_hovbox fmt ctx.indent_incr;
+    let ctx =
+      if monadic then (
+        (* Note that in F*, the left value of a monadic let-binding can only be
+         * a variable *)
+        let ctx = extract_typed_pattern ctx fmt true lv in
+        F.pp_print_space fmt ();
+        let arrow = match !backend with FStar -> "<--" | Coq -> "<-" in
+        F.pp_print_string fmt arrow;
+        F.pp_print_space fmt ();
+        extract_texpression ctx fmt false re;
+        F.pp_print_string fmt ";";
+        ctx)
+      else (
+        F.pp_print_string fmt "let";
+        F.pp_print_space fmt ();
+        let ctx = extract_typed_pattern ctx fmt true lv in
+        F.pp_print_space fmt ();
+        let eq = match !backend with FStar -> "=" | Coq -> ":=" in
+        F.pp_print_string fmt eq;
+        F.pp_print_space fmt ();
+        extract_texpression ctx fmt false re;
+        F.pp_print_space fmt ();
+        F.pp_print_string fmt "in";
+        ctx)
+    in
+    (* Close the box for the let-binding *)
+    F.pp_close_box fmt ();
+    F.pp_print_space fmt ();
+    (* Return *)
+    ctx
+  in
   let ctx =
-    if monadic then (
-      (* Note that in F*, the left value of a monadic let-binding can only be
-       * a variable *)
-      let ctx = extract_typed_pattern ctx fmt true lv in
-      F.pp_print_space fmt ();
-      let arrow = match !backend with FStar -> "<--" | Coq -> "<-" in
-      F.pp_print_string fmt arrow;
-      F.pp_print_space fmt ();
-      extract_texpression ctx fmt false re;
-      F.pp_print_string fmt ";";
-      ctx)
-    else (
-      F.pp_print_string fmt "let";
-      F.pp_print_space fmt ();
-      let ctx = extract_typed_pattern ctx fmt true lv in
-      F.pp_print_space fmt ();
-      let eq = match !backend with FStar -> "=" | Coq -> ":=" in
-      F.pp_print_string fmt eq;
-      F.pp_print_space fmt ();
-      extract_texpression ctx fmt false re;
-      F.pp_print_space fmt ();
-      F.pp_print_string fmt "in";
-      ctx)
+    List.fold_left
+      (fun ctx (monadic, lv, re) -> extract_let ctx monadic lv re)
+      ctx lets
   in
-  (* Close the box for the let-binding *)
-  F.pp_close_box fmt ();
+  (* Open a box for the next expression *)
+  F.pp_open_hovbox fmt ctx.indent_incr;
   (* Print the next expression *)
-  F.pp_print_space fmt ();
   extract_texpression ctx fmt false next_e;
+  (* Close the box for the next expression *)
+  F.pp_close_box fmt ();
   (* Close parentheses *)
   if inside then F.pp_print_string fmt ")";
   (* Close the box for the whole expression *)
diff --git a/compiler/PureUtils.ml b/compiler/PureUtils.ml
index 5a024d9e..728a4fe6 100644
--- a/compiler/PureUtils.ml
+++ b/compiler/PureUtils.ml
@@ -206,6 +206,15 @@ let mk_arrows (inputs : ty list) (output : ty) =
   in
   aux inputs
 
+(** Destruct an expression into a list of nested lets *)
+let rec destruct_lets (e : texpression) :
+    (bool * typed_pattern * texpression) list * texpression =
+  match e.e with
+  | Let (monadic, lv, re, next_e) ->
+      let lets, last_e = destruct_lets next_e in
+      ((monadic, lv, re) :: lets, last_e)
+  | _ -> ([], e)
+
 (** Destruct an [App] expression into an expression and a list of arguments.
 
     We simply destruct the expression as long as it is of the form [App (f, x)].
diff --git a/tests/coq/misc/External__Funs.v b/tests/coq/misc/External__Funs.v
index 021acd6e..e8ecfab8 100644
--- a/tests/coq/misc/External__Funs.v
+++ b/tests/coq/misc/External__Funs.v
@@ -18,7 +18,8 @@ Definition swap_fwd
   p0 <- core_mem_swap_back0 T x y st st0;
   let (st1, _) := p0 in
   p1 <- core_mem_swap_back1 T x y st st1;
-  let (st2, _) := p1 in Return (st2, tt)
+  let (st2, _) := p1 in
+  Return (st2, tt)
   .
 
 (** [external::swap] *)
@@ -31,7 +32,8 @@ Definition swap_back
   p0 <- core_mem_swap_back0 T x y st st1;
   let (st2, x0) := p0 in
   p1 <- core_mem_swap_back1 T x y st st2;
-  let (_, y0) := p1 in Return (st0, (x0, y0))
+  let (_, y0) := p1 in
+  Return (st0, (x0, y0))
   .
 
 (** [external::test_new_non_zero_u32] *)
@@ -40,13 +42,16 @@ Definition test_new_non_zero_u32_fwd
   p <- core_num_nonzero_non_zero_u32_new_fwd x st;
   let (st0, opt) := p in
   p0 <- core_option_option_unwrap_fwd Core_num_nonzero_non_zero_u32_t opt st0;
-  let (st1, nzu) := p0 in Return (st1, nzu)
+  let (st1, nzu) := p0 in
+  Return (st1, nzu)
   .
 
 (** [external::test_vec] *)
 Definition test_vec_fwd : result unit :=
   let v := vec_new u32 in
-  v0 <- vec_push_back u32 v (0 %u32); let _ := v0 in Return tt
+  v0 <- vec_push_back u32 v (0 %u32);
+  let _ := v0 in
+  Return tt
   .
 
 (** [external::custom_swap] *)
@@ -57,7 +62,8 @@ Definition custom_swap_fwd
   p0 <- core_mem_swap_back0 T x y st st0;
   let (st1, x0) := p0 in
   p1 <- core_mem_swap_back1 T x y st st1;
-  let (st2, _) := p1 in Return (st2, x0)
+  let (st2, _) := p1 in
+  Return (st2, x0)
   .
 
 (** [external::custom_swap] *)
@@ -70,7 +76,8 @@ Definition custom_swap_back
   p0 <- core_mem_swap_back0 T x y st st1;
   let (st2, _) := p0 in
   p1 <- core_mem_swap_back1 T x y st st2;
-  let (_, y0) := p1 in Return (st0, (ret, y0))
+  let (_, y0) := p1 in
+  Return (st0, (ret, y0))
   .
 
 (** [external::test_custom_swap] *)
@@ -84,7 +91,9 @@ Definition test_custom_swap_back
   result (state * (u32 * u32))
   :=
   p <- custom_swap_back u32 x y st (1 %u32) st0;
-  let (st1, p0) := p in let (x0, y0) := p0 in Return (st1, (x0, y0))
+  let (st1, p0) := p in
+  let (x0, y0) := p0 in
+  Return (st1, (x0, y0))
   .
 
 (** [external::test_swap_non_zero] *)
@@ -94,7 +103,8 @@ Definition test_swap_non_zero_fwd
   let (st0, _) := p in
   p0 <- swap_back u32 x (0 %u32) st st0;
   let (st1, p1) := p0 in
-  let (x0, _) := p1 in if x0 s= 0 %u32 then Fail_ else Return (st1, x0)
+  let (x0, _) := p1 in
+  if x0 s= 0 %u32 then Fail_ else Return (st1, x0)
   .
 
 End External__Funs .
diff --git a/tests/coq/misc/NoNestedBorrows.v b/tests/coq/misc/NoNestedBorrows.v
index 6dc41204..9075d01d 100644
--- a/tests/coq/misc/NoNestedBorrows.v
+++ b/tests/coq/misc/NoNestedBorrows.v
@@ -94,7 +94,8 @@ Definition get_max_fwd (x : u32) (y : u32) : result u32 :=
 Definition test3_fwd : result unit :=
   x <- get_max_fwd (4 %u32) (3 %u32);
   y <- get_max_fwd (10 %u32) (11 %u32);
-  z <- u32_add x y; if negb (z s= 15 %u32) then Fail_ else Return tt
+  z <- u32_add x y;
+  if negb (z s= 15 %u32) then Fail_ else Return tt
   .
 
 (** Unit test for [no_nested_borrows::test3] *)
@@ -139,7 +140,8 @@ Check (test_list1_fwd )%return.
 (** [no_nested_borrows::test_box1] *)
 Definition test_box1_fwd : result unit :=
   let b := 1 %i32 in
-  let x := b in if negb (x s= 1 %i32) then Fail_ else Return tt
+  let x := b in
+  if negb (x s= 1 %i32) then Fail_ else Return tt
   .
 
 (** Unit test for [no_nested_borrows::test_box1] *)
@@ -186,7 +188,8 @@ Definition split_list_fwd
 Definition test_split_list_fwd : result unit :=
   let l := ListNil in
   p <- split_list_fwd i32 (ListCons (0 %i32) l);
-  let (hd, _) := p in if negb (hd s= 0 %i32) then Fail_ else Return tt
+  let (hd, _) := p in
+  if negb (hd s= 0 %i32) then Fail_ else Return tt
   .
 
 (** Unit test for [no_nested_borrows::test_split_list] *)
@@ -281,7 +284,8 @@ Fixpoint list_nth_mut_back
     else
       (
         i0 <- u32_sub i 1 %u32;
-        tl0 <- list_nth_mut_back T tl i0 ret; Return (ListCons x tl0) )
+        tl0 <- list_nth_mut_back T tl i0 ret;
+        Return (ListCons x tl0) )
   | ListNil => Fail_
   end
   .
@@ -298,7 +302,8 @@ Fixpoint list_rev_aux_fwd
 (** [no_nested_borrows::list_rev] *)
 Definition list_rev_fwd_back (T : Type) (l : List_t T) : result (List_t T) :=
   let li := mem_replace_fwd (List_t T) l ListNil in
-  l0 <- list_rev_aux_fwd T li ListNil; Return l0
+  l0 <- list_rev_aux_fwd T li ListNil;
+  Return l0
   .
 
 (** [no_nested_borrows::test_list_functions] *)
diff --git a/tests/coq/misc/Paper.v b/tests/coq/misc/Paper.v
index 5d9598eb..e15e0dc1 100644
--- a/tests/coq/misc/Paper.v
+++ b/tests/coq/misc/Paper.v
@@ -76,7 +76,8 @@ Fixpoint list_nth_mut_back
     else
       (
         i0 <- u32_sub i 1 %u32;
-        tl0 <- list_nth_mut_back T tl i0 ret; Return (ListCons x tl0) )
+        tl0 <- list_nth_mut_back T tl i0 ret;
+        Return (ListCons x tl0) )
   | ListNil => Fail_
   end
   .
@@ -97,7 +98,8 @@ Definition test_nth_fwd : result unit :=
   x <- list_nth_mut_fwd i32 (ListCons (1 %i32) l1) (2 %u32);
   x0 <- i32_add x 1 %i32;
   l2 <- list_nth_mut_back i32 (ListCons (1 %i32) l1) (2 %u32) x0;
-  i <- sum_fwd l2; if negb (i s= 7 %i32) then Fail_ else Return tt
+  i <- sum_fwd l2;
+  if negb (i s= 7 %i32) then Fail_ else Return tt
   .
 
 (** Unit test for [paper::test_nth] *)
@@ -108,7 +110,9 @@ Definition call_choose_fwd (p : (u32 * u32)) : result u32 :=
   let (px, py) := p in
   pz <- choose_fwd u32 true px py;
   pz0 <- u32_add pz 1 %u32;
-  p0 <- choose_back u32 true px py pz0; let (px0, _) := p0 in Return px0
+  p0 <- choose_back u32 true px py pz0;
+  let (px0, _) := p0 in
+  Return px0
   .
 
 End Paper .
diff --git a/tests/fstar/betree/BetreeMain.Funs.fst b/tests/fstar/betree/BetreeMain.Funs.fst
index 9ba5d3e7..8cb5eb41 100644
--- a/tests/fstar/betree/BetreeMain.Funs.fst
+++ b/tests/fstar/betree/BetreeMain.Funs.fst
@@ -151,7 +151,8 @@ let rec betree_list_split_at_fwd
         | Fail -> Fail
         | Return p ->
           let (ls0, ls1) = p in
-          let l = ls0 in Return (BetreeListCons hd l, ls1)
+          let l = ls0 in
+          Return (BetreeListCons hd l, ls1)
         end
       end
     | BetreeListNil -> Fail
@@ -161,7 +162,8 @@ let rec betree_list_split_at_fwd
 let betree_list_push_front_fwd_back
   (t : Type0) (self : betree_list_t t) (x : t) : result (betree_list_t t) =
   let tl = mem_replace_fwd (betree_list_t t) self BetreeListNil in
-  let l = tl in Return (BetreeListCons x l)
+  let l = tl in
+  Return (BetreeListCons x l)
 
 (** [betree_main::betree::List::{1}::pop_front] *)
 let betree_list_pop_front_fwd (t : Type0) (self : betree_list_t t) : result t =
@@ -211,7 +213,8 @@ let rec betree_list_partition_at_pivot_fwd
       | Fail -> Fail
       | Return p ->
         let (ls0, ls1) = p in
-        let l = ls0 in Return (BetreeListCons (i, x) l, ls1)
+        let l = ls0 in
+        Return (BetreeListCons (i, x) l, ls1)
       end
   | BetreeListNil -> Return (BetreeListNil, BetreeListNil)
   end
@@ -252,13 +255,8 @@ let betree_leaf_split_fwd
                   params.betree_params_split_size) in
                 let n0 = BetreeNodeLeaf (Mkbetree_leaf_t id1
                   params.betree_params_split_size) in
-                Return
-                (st1,
-                Mkbetree_internal_t
-                self.betree_leaf_id
-                pivot
-                n
-                n0)
+                Return (st1, Mkbetree_internal_t self.betree_leaf_id pivot n
+                  n0)
               end
             end
           end
diff --git a/tests/fstar/betree_back_stateful/BetreeMain.Funs.fst b/tests/fstar/betree_back_stateful/BetreeMain.Funs.fst
index ea8344fa..eebed6e6 100644
--- a/tests/fstar/betree_back_stateful/BetreeMain.Funs.fst
+++ b/tests/fstar/betree_back_stateful/BetreeMain.Funs.fst
@@ -151,7 +151,8 @@ let rec betree_list_split_at_fwd
         | Fail -> Fail
         | Return p ->
           let (ls0, ls1) = p in
-          let l = ls0 in Return (BetreeListCons hd l, ls1)
+          let l = ls0 in
+          Return (BetreeListCons hd l, ls1)
         end
       end
     | BetreeListNil -> Fail
@@ -161,7 +162,8 @@ let rec betree_list_split_at_fwd
 let betree_list_push_front_fwd_back
   (t : Type0) (self : betree_list_t t) (x : t) : result (betree_list_t t) =
   let tl = mem_replace_fwd (betree_list_t t) self BetreeListNil in
-  let l = tl in Return (BetreeListCons x l)
+  let l = tl in
+  Return (BetreeListCons x l)
 
 (** [betree_main::betree::List::{1}::pop_front] *)
 let betree_list_pop_front_fwd (t : Type0) (self : betree_list_t t) : result t =
@@ -211,7 +213,8 @@ let rec betree_list_partition_at_pivot_fwd
       | Fail -> Fail
       | Return p ->
         let (ls0, ls1) = p in
-        let l = ls0 in Return (BetreeListCons (i, x) l, ls1)
+        let l = ls0 in
+        Return (BetreeListCons (i, x) l, ls1)
       end
   | BetreeListNil -> Return (BetreeListNil, BetreeListNil)
   end
@@ -252,13 +255,8 @@ let betree_leaf_split_fwd
                   params.betree_params_split_size) in
                 let n0 = BetreeNodeLeaf (Mkbetree_leaf_t id1
                   params.betree_params_split_size) in
-                Return
-                (st1,
-                Mkbetree_internal_t
-                self.betree_leaf_id
-                pivot
-                n
-                n0)
+                Return (st1, Mkbetree_internal_t self.betree_leaf_id pivot n
+                  n0)
               end
             end
           end