Make a minor improvement in printing

2 files changed, 13 insertions, 4 deletions

diff --git a/src/ExtractToFStar.ml b/src/ExtractToFStar.ml
index 0781a811..7d8e84c4 100644
--- a/src/ExtractToFStar.ml
+++ b/src/ExtractToFStar.ml
@@ -961,10 +961,18 @@ let extract_fun_def (ctx : extraction_ctx) (fmt : F.formatter)
       ctx def.inputs_lvs
   in
   (* Print the return type *)
-  F.pp_print_space fmt ();
-  F.pp_print_string fmt ":";
-  F.pp_print_space fmt ();
-  extract_ty ctx fmt false (Collections.List.to_cons_nil def.signature.outputs);
+  let _ =
+    F.pp_print_space fmt ();
+    (* Open a box for the return type *)
+    F.pp_open_hovbox fmt 0;
+    (* Print the return type *)
+    F.pp_print_string fmt ":";
+    F.pp_print_space fmt ();
+    extract_ty ctx fmt false
+      (Collections.List.to_cons_nil def.signature.outputs);
+    (* Close the box for the return type *)
+    F.pp_close_box fmt ()
+  in
   (* Print the "=" *)
   F.pp_print_space fmt ();
   F.pp_print_string fmt "=";
diff --git a/src/PureUtils.ml b/src/PureUtils.ml
index aa383d8c..9f3bd5ef 100644
--- a/src/PureUtils.ml
+++ b/src/PureUtils.ml
@@ -186,6 +186,7 @@ let functions_not_mutually_recursive (funs : fun_def list) : bool =
 
 (** We use this to check whether we need to add parentheses around expressions.
     We only look for outer monadic let-bindings.
+    This is used when printing the branches of `if ... then ... else ...`.
  *)
 let rec expression_requires_parentheses (e : texpression) : bool =
   match e.e with