From 8e4bda876dbdafdbc3e4de1d9dad495750102be6 Mon Sep 17 00:00:00 2001
From: Jonathan Protzenko <protz@microsoft.com>
Date: Wed, 30 Oct 2024 16:20:48 -0700
Subject: [PATCH] Peephole optimzations to get correct mutability inference for
 splits -- all thanks to @r1km

---
 lib/OptimizeMiniRust.ml | 21 ++++++++++++++++++++-
 src/Karamel.ml          |  1 +
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/lib/OptimizeMiniRust.ml b/lib/OptimizeMiniRust.ml
index 812e0fd4..0774db3e 100644
--- a/lib/OptimizeMiniRust.ml
+++ b/lib/OptimizeMiniRust.ml
@@ -612,7 +612,7 @@ let rec infer_expr (env: env) valuation (expected: typ) (known: known) (e: expr)
 
   (* Special case for array slices. This occurs, e.g., when calling a function with
      a struct field *)
-  | Field (Open { atom; name }, "0", None) | Field (Open { atom; name }, "1", None) ->
+  | Field (Open { atom; _ }, "0", None) | Field (Open { atom; _ }, "1", None) ->
       if is_mut_borrow expected then
         add_mut_borrow atom known, e
       else
@@ -1174,6 +1174,21 @@ let remove_deref_addrof = object
     | _ -> Deref e
 end
 
+let cleanup_splits = object(self)
+  inherit [_] map_expr as super
+  method! visit_Match _ e_scrut t branches =
+    match t with
+    | Tuple [ _; _ ] ->
+        let bs, p, e = KList.one branches in
+        let b1, b2 = KList.two bs in
+        assert (match p with TupleP _ -> true | _ -> false);
+        Let (b1, Field (e_scrut, "0", None),
+        Let (b2, Field (lift 1 e_scrut, "1", None),
+        self#visit_expr () e))
+    | _ ->
+        super#visit_Match () e_scrut t branches
+end
+
 let map_funs f_map files =
   let files =
     List.fold_left (fun files (filename, decls) ->
@@ -1261,6 +1276,10 @@ let add_derives valuation files =
     ) decls
   ) files
 
+let cleanup_minirust files =
+  let files = map_funs cleanup_splits#visit_expr files in
+  files
+
 let simplify_minirust files =
   let files = map_funs unroll_loops#visit_expr files in
   let files = map_funs remove_auto_deref#visit_expr files in
diff --git a/src/Karamel.ml b/src/Karamel.ml
index 7276f458..62db3050 100644
--- a/src/Karamel.ml
+++ b/src/Karamel.ml
@@ -758,6 +758,7 @@ Supported options:|}
     if Options.debug "rs" then
       print PrintAst.print_files files;
     let files = AstToMiniRust.translate_files files in
+    let files = OptimizeMiniRust.cleanup_minirust files in
     let files = OptimizeMiniRust.infer_mut_borrows files in
     let files = OptimizeMiniRust.simplify_minirust files in
     OutputRust.write_all files