Improve the heuristic to find pretty names for the variables in the loops

5 files changed, 64 insertions, 5 deletions

diff --git a/compiler/PrintPure.ml b/compiler/PrintPure.ml
index 532271c3..2c8d5081 100644
--- a/compiler/PrintPure.ml
+++ b/compiler/PrintPure.ml
@@ -521,7 +521,7 @@ let rec texpression_to_string (fmt : ast_formatter) (inside : bool)
       let meta_s = meta_to_string fmt meta in
       let e = texpression_to_string fmt inside indent indent_incr e in
       match meta with
-      | Assignment _ | Tag _ ->
+      | Assignment _ | SymbolicAssignment _ | Tag _ ->
           let e = meta_s ^ "\n" ^ indent ^ e in
           if inside then "(" ^ e ^ ")" else e
       | MPlace _ -> "(" ^ meta_s ^ " " ^ e ^ ")")
@@ -665,6 +665,10 @@ and meta_to_string (fmt : ast_formatter) (meta : meta) : string =
         "@assign(" ^ mplace_to_string fmt lp ^ " := "
         ^ texpression_to_string fmt false "" "" rv
         ^ rp ^ ")"
+    | SymbolicAssignment (var_id, rv) ->
+        "@symb_assign(" ^ VarId.to_string var_id ^ " := "
+        ^ texpression_to_string fmt false "" "" rv
+        ^ ")"
     | MPlace mp -> "@mplace=" ^ mplace_to_string fmt mp
     | Tag msg -> "@tag \"" ^ msg ^ "\""
   in
diff --git a/compiler/Pure.ml b/compiler/Pure.ml
index 118aec50..fe30a650 100644
--- a/compiler/Pure.ml
+++ b/compiler/Pure.ml
@@ -358,7 +358,7 @@ type qualif_id =
  *)
 type qualif = { id : qualif_id; type_args : ty list } [@@deriving show]
 
-type var_id = VarId.id [@@deriving show]
+type var_id = VarId.id [@@deriving show, ord]
 
 (** Ancestor for {!iter_expression} visitor *)
 class ['self] iter_expression_base =
@@ -513,14 +513,21 @@ and texpression = { e : expression; ty : ty }
  *)
 and mvalue = (texpression[@opaque])
 
+(** Meta-information stored in the AST *)
 and meta =
   | Assignment of mplace * mvalue * mplace option
-      (** Meta-information stored in the AST.
+      (** Information about an assignment which occured in LLBC.
+          We use this to guide the heuristics which derive pretty names.
 
           The first mplace stores the destination.
           The mvalue stores the value which is put in the destination
           The second (optional) mplace stores the origin.
         *)
+  | 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 *)
   | Tag of string  (** A tag - typically used for debugging *)
 [@@deriving
diff --git a/compiler/PureMicroPasses.ml b/compiler/PureMicroPasses.ml
index b9441397..09cc2533 100644
--- a/compiler/PureMicroPasses.ml
+++ b/compiler/PureMicroPasses.ml
@@ -290,6 +290,17 @@ let compute_pretty_names (def : fun_decl) : fun_decl =
     (* Add the constraint *)
     match (unmeta rv).e with Var vid -> add_constraint mp vid ctx | _ -> ctx
   in
+  let add_pure_var_value_constraint (var_id : VarId.id) (rv : texpression)
+      (ctx : pn_ctx) : pn_ctx =
+    (* Add the constraint *)
+    match (unmeta rv).e with
+    | Var vid -> (
+        (* Try to find a name for the vid *)
+        match VarId.Map.find_opt vid ctx.pure_vars with
+        | None -> ctx
+        | Some name -> add_pure_var_constraint var_id name ctx)
+    | _ -> ctx
+  in
   (* Specific case of constraint on left values *)
   let add_left_constraint (lv : typed_pattern) (ctx : pn_ctx) : pn_ctx =
     let obj =
@@ -484,6 +495,8 @@ let compute_pretty_names (def : fun_decl) : fun_decl =
             | _ -> ctx
           in
           ctx
+      | SymbolicAssignment (var_id, rvalue) ->
+          add_pure_var_value_constraint var_id rvalue ctx
       | MPlace mp -> add_right_constraint mp e ctx
       | Tag _ -> ctx
     in
diff --git a/compiler/PureUtils.ml b/compiler/PureUtils.ml
index a60bcd78..e86d2a78 100644
--- a/compiler/PureUtils.ml
+++ b/compiler/PureUtils.ml
@@ -57,6 +57,8 @@ let compute_primitive_value_ty (cv : primitive_value) : ty =
   | Char _ -> Char
   | String _ -> Str
 
+let var_get_id (v : var) : VarId.id = v.id
+
 let mk_typed_pattern_from_primitive_value (cv : primitive_value) : typed_pattern
     =
   let ty = compute_primitive_value_ty cv in
@@ -73,6 +75,10 @@ let mk_tag (msg : string) (next_e : texpression) : texpression =
   let ty = next_e.ty in
   { e; ty }
 
+let mk_mplace (var_id : E.VarId.id) (name : string option)
+    (projection : mprojection) : mplace =
+  { var_id; name; projection }
+
 (** Type substitution *)
 let ty_substitute (tsubst : TypeVarId.id -> ty) (ty : ty) : ty =
   let obj =
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index 49b696b2..ab9e40df 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -103,6 +103,7 @@ type call_info = {
  *)
 type loop_info = {
   loop_id : LoopId.id;
+  input_vars : var list;
   input_svl : V.symbolic_value list;
   type_args : ty list;
   forward_inputs : texpression list option;
@@ -149,7 +150,7 @@ type bs_ctx = {
   fuel0 : VarId.id;
       (** The original fuel taken as input by the function (if we use fuel) *)
   fuel : VarId.id;
-  (* The fuel to use for the recursive calls (if we use fuel) *)
+      (** The fuel to use for the recursive calls (if we use fuel) *)
   forward_inputs : var list;
       (** The input parameters for the forward function corresponding to the
           translated Rust inputs (no fuel, no state).
@@ -169,6 +170,8 @@ type bs_ctx = {
           [None] if we are not inside a loop, [Some] otherwise (and whatever
           the kind of function we are translating: it will be [Some] even
           though we are synthesizing a forward function).
+
+          TODO: move to {!loop_info}
        *)
   calls : call_info V.FunCallId.Map.t;
       (** The function calls we encountered so far *)
@@ -2254,6 +2257,7 @@ and translate_forward_end (ectx : C.eval_ctx)
       let args =
         List.map (typed_value_to_texpression ctx ectx) loop_input_values
       in
+      let org_args = args in
 
       (* Lookup the effect info for the loop function *)
       let fid = A.Regular ctx.fun_decl.A.def_id in
@@ -2316,7 +2320,31 @@ and translate_forward_end (ectx : C.eval_ctx)
       in
 
       (* Create the let expression with the loop call *)
-      mk_let effect_info.can_fail out_pat loop_call next_e
+      let e = mk_let effect_info.can_fail out_pat loop_call next_e in
+
+      (* Add meta-information linking the loop input parameters and the
+         loop input values - we use this to derive proper names.
+
+         There is something important here: as we group the end of the function
+         and the loop body in a {!Loop} node, when exploring the function
+         and applying micro passes, we introduce the variables specific to
+         the loop body before exploring both the loop body and the end of
+         the function. It means it is ok to reference some variables which might
+         actually be defined, in the end, in a different branch.
+
+         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
 
 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
@@ -2416,6 +2444,7 @@ and translate_loop (loop : S.loop) (ctx : bs_ctx) : texpression =
     let loop_info =
       {
         loop_id;
+        input_vars = inputs;
         input_svl = loop.input_svalues;
         type_args;
         forward_inputs = None;