Improve the pretty names generation for loops

4 files changed, 59 insertions, 16 deletions

diff --git a/compiler/Pure.ml b/compiler/Pure.ml
index 777d4308..912e05fb 100644
--- a/compiler/Pure.ml
+++ b/compiler/Pure.ml
@@ -530,6 +530,7 @@ and meta =
   | SymbolicAssignment of (var_id[@opaque]) * mvalue
       (** Informationg linking a variable (from the pure AST) to an
           expression.
+
           We use this to guide the heuristics which derive pretty names.
         *)
   | MPlace of mplace  (** Meta-information about the origin of a value *)
diff --git a/compiler/PureMicroPasses.ml b/compiler/PureMicroPasses.ml
index e670570b..3614487e 100644
--- a/compiler/PureMicroPasses.ml
+++ b/compiler/PureMicroPasses.ml
@@ -1723,6 +1723,8 @@ let apply_passes_to_def (ctx : trans_ctx) (def : fun_decl) :
       ^ Print.option_to_string T.RegionGroupId.to_string def.back_id
       ^ ")"));
 
+  log#ldebug (lazy ("original decl:\n\n" ^ fun_decl_to_string ctx def ^ "\n"));
+
   (* First, find names for the variables which are unnamed *)
   let def = compute_pretty_names def in
   log#ldebug
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index 531a13e9..c6ef4297 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -914,6 +914,13 @@ let rec unbox_typed_value (v : V.typed_value) : V.typed_value =
       | _ -> raise (Failure "Unreachable"))
   | _ -> v
 
+(** Translate a symbolic value *)
+let symbolic_value_to_texpression (ctx : bs_ctx) (sv : V.symbolic_value) :
+    texpression =
+  (* Translate the type *)
+  let var = lookup_var_for_symbolic_value sv ctx in
+  mk_texpression_from_var var
+
 (** Translate a typed value.
 
     It is used, for instance, on values used as inputs for function calls.
@@ -988,9 +995,7 @@ let rec typed_value_to_texpression (ctx : bs_ctx) (ectx : C.eval_ctx)
         | V.MutBorrow (_, v) ->
             (* Borrows are the identity in the extraction *)
             translate v)
-    | Symbolic sv ->
-        let var = lookup_var_for_symbolic_value sv ctx in
-        mk_texpression_from_var var
+    | Symbolic sv -> symbolic_value_to_texpression ctx sv
   in
   (* Debugging *)
   log#ldebug
@@ -1288,6 +1293,14 @@ let get_abs_ancestors (ctx : bs_ctx) (abs : V.abs) (call_id : V.FunCallId.id) :
   let abs_ancestors = list_ancestor_abstractions ctx abs call_id in
   (call_info.forward, abs_ancestors)
 
+(** Add meta-information to an expression *)
+let mk_meta_symbolic_assignments (vars : var list) (values : texpression list)
+    (e : texpression) : texpression =
+  let var_values = List.combine vars values in
+  List.fold_right
+    (fun (var, arg) e -> mk_meta (SymbolicAssignment (var_get_id var, arg)) e)
+    var_values e
+
 let rec translate_expression (e : S.expression) (ctx : bs_ctx) : texpression =
   match e with
   | S.Return (ectx, opt_v) -> translate_return ectx opt_v ctx
@@ -1891,8 +1904,10 @@ and translate_end_abstraction_loop (ectx : C.eval_ctx) (abs : V.abs)
          need to *transmit* to the loop backward function): they are not the
          values consumed upon ending the abstraction (i.e., we don't use
          [abs_to_consumed]) *)
-      let back_inputs = T.RegionGroupId.Map.find rg_id ctx.backward_inputs in
-      let back_inputs = List.map mk_texpression_from_var back_inputs in
+      let back_inputs_vars =
+        T.RegionGroupId.Map.find rg_id ctx.backward_inputs
+      in
+      let back_inputs = List.map mk_texpression_from_var back_inputs_vars in
       (* If the function is stateful:
        * - add the state input argument
        * - generate a fresh state variable for the returned state
@@ -1953,7 +1968,39 @@ and translate_end_abstraction_loop (ectx : C.eval_ctx) (abs : V.abs)
          * a value containing mutable borrows, which can't be the case... *)
         assert (List.length inputs = List.length fwd_inputs);
         next_e)
-      else mk_let effect_info.can_fail output call next_e
+      else
+        (* Add meta-information - this is slightly hacky: we look at the
+           values consumed by the abstraction (note that those come from
+           *before* we applied the fixed-point context) and use them to
+           guide the naming of the output vars.
+
+           Also, we need to convert the backward outputs from patterns to
+           variables.
+
+           Finally, in practice, this works well only for loop bodies:
+           we do this only in this case.
+           TODO: improve the heuristics, to give weight to the hints for
+           instance.
+        *)
+        let next_e =
+          if ctx.inside_loop then
+            let consumed_values = abs_to_consumed ctx ectx abs in
+            let var_values = List.combine outputs consumed_values in
+            let var_values =
+              List.filter_map
+                (fun (var, v) ->
+                  match var.Pure.value with
+                  | PatVar (var, _) -> Some (var, v)
+                  | _ -> None)
+                var_values
+            in
+            let vars, values = List.split var_values in
+            mk_meta_symbolic_assignments vars values next_e
+          else next_e
+        in
+
+        (* Create the let-binding *)
+        mk_let effect_info.can_fail output call next_e
 
 and translate_global_eval (gid : A.GlobalDeclId.id) (sval : V.symbolic_value)
     (e : S.expression) (ctx : bs_ctx) : texpression =
@@ -2335,16 +2382,7 @@ and translate_forward_end (ectx : C.eval_ctx)
          We then remove all the meta information from the body *before* calling
          {!PureMicroPasses.decompose_loops}.
       *)
-      let e =
-        let var_values = List.combine loop_info.input_vars org_args in
-        List.fold_right
-          (fun (var, arg) e ->
-            mk_meta (SymbolicAssignment (var_get_id var, arg)) e)
-          var_values e
-      in
-
-      (* Return *)
-      e
+      mk_meta_symbolic_assignments loop_info.input_vars org_args e
 
 and translate_loop (loop : S.loop) (ctx : bs_ctx) : texpression =
   let loop_id = V.LoopId.Map.find loop.loop_id ctx.loop_ids_map in
diff --git a/compiler/SynthesizeSymbolic.ml b/compiler/SynthesizeSymbolic.ml
index 976b781d..6668c043 100644
--- a/compiler/SynthesizeSymbolic.ml
+++ b/compiler/SynthesizeSymbolic.ml
@@ -20,6 +20,8 @@ let mk_opt_place_from_op (op : E.operand) (ctx : Contexts.eval_ctx) :
   | E.Copy p | E.Move p -> Some (mk_mplace p ctx)
   | E.Constant _ -> None
 
+let mk_meta (m : meta) (e : expression) : expression = Meta (m, e)
+
 let synthesize_symbolic_expansion (sv : V.symbolic_value)
     (place : mplace option) (seel : V.symbolic_expansion option list)
     (el : expression list option) : expression option =