Improve simplify_aggregates to introduce structure updates

2 files changed, 48 insertions, 3 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index 8dd5910f..f2b5f00f 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -2054,7 +2054,13 @@ and extract_StructUpdate (ctx : extraction_ctx) (fmt : F.formatter)
     let var_name = ctx_get_var (Option.get supd.init) ctx in
     F.pp_print_string fmt var_name;
     F.pp_print_space fmt ();
-    F.pp_print_string fmt "where";
+    let where_keyword =
+      match !backend with
+      | FStar -> "where"
+      | Lean -> "with"
+      | Coq -> raise (Failure "Unreachable")
+    in
+    F.pp_print_string fmt where_keyword;
     F.pp_print_space fmt ());
   F.pp_open_hvbox fmt 0;
   let delimiter = match !backend with Lean -> "," | Coq | FStar -> ";" in
diff --git a/compiler/PureMicroPasses.ml b/compiler/PureMicroPasses.ml
index 7d01a622..74f3c576 100644
--- a/compiler/PureMicroPasses.ml
+++ b/compiler/PureMicroPasses.ml
@@ -1178,8 +1178,47 @@ let simplify_aggregates (ctx : trans_ctx) (def : fun_decl) : fun_decl =
                 List.length filt_var_projs = List.length updates
                 && List.for_all (fun y -> y = x) filt_var_projs
               then { e with e = Var x }
-              else (* TODO *)
-                super#visit_texpression env e
+              else if
+                (* Attempt to create an "update" expression (i.e., of
+                   the shape [{ x with f := v }]).
+
+                   This is not supported by Coq *)
+                !Config.backend <> Coq
+              then (
+                (* Compute the number of occurrences of each variable *)
+                let occurs = ref VarId.Map.empty in
+                List.iter
+                  (fun x ->
+                    let num =
+                      match VarId.Map.find_opt x !occurs with
+                      | None -> 1
+                      | Some n -> n + 1
+                    in
+                    occurs := VarId.Map.add x num !occurs)
+                  filt_var_projs;
+                (* Find the max - note that we can initialize the max at 0,
+                   because there is at least one variable *)
+                let max = ref 0 in
+                let x = ref x in
+                List.iter
+                  (fun (y, n) ->
+                    if n > !max then (
+                      max := n;
+                      x := y))
+                  (VarId.Map.bindings !occurs);
+                (* Create the update expression *)
+                let updates =
+                  List.filter_map
+                    (fun ((fid, fe), y_opt) ->
+                      if y_opt = Some !x then None else Some (fid, fe))
+                    (List.combine updates var_projs)
+                in
+                let supd =
+                  StructUpdate { struct_id; init = Some !x; updates }
+                in
+                let e = { e with e = supd } in
+                super#visit_texpression env e)
+              else super#visit_texpression env e
         | _ -> super#visit_texpression env e
     end
   in