From 1f6f6d7e50fb36c29dced9281070c3ea34227c18 Mon Sep 17 00:00:00 2001
From: Jonathan Protzenko <protz@microsoft.com>
Date: Mon, 16 Sep 2024 15:39:35 -0700
Subject: [PATCH] WIP: handling match nodes properly

---
 lib/DataTypes.ml        | 28 +++++++++++++++++++++++++++-
 lib/OptimizeMiniRust.ml |  5 +++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/lib/DataTypes.ml b/lib/DataTypes.ml
index fdf3acdc..d29aaeba 100644
--- a/lib/DataTypes.ml
+++ b/lib/DataTypes.ml
@@ -1314,11 +1314,37 @@ let optimize files =
   let files = remove_unit_fields#visit_files () files in
   files
 
+(* For Rust, we leave `match` nodes in the AST -- Rust *does* have
+   pattern-matching, after all. However, F* code contains a lot of one-branch
+   matches, of the form `match e with _ -> ...` or `match e with { f = x; } ->
+   ...`. Therefore, for code quality, we replace the latter with `...` and the
+   former with `let x = e.f in ...`. *)
+let remove_one_branch_matches = object (self)
+
+  inherit [_] map as _super
+
+  method! visit_EMatch (_, _t as env) c scrut branches =
+    let scrut = self#visit_expr env scrut in
+    match branches with
+    | [ [], _, e ] ->
+        (* A match made of a single branch; e is always evaluated. *)
+        (self#visit_expr env e).node
+    | [ [ b ], { node = PBound _; _ }, e ] ->
+        (* match scrut with \ b. x -> e  ~~>  let x = scrut in e *)
+        ELet (b, scrut, self#visit_expr env e)
+    | [ [ b ], { node = PRecord ps; _ }, e ] ->
+        (* match scrut with \ b. { f = x } -> e  ~~>  let x = scrut.f in e *)
+        let f = Option.get (List.find_map (fun (f, x) -> match x.node with PBound _ -> Some f | _ -> None) ps) in
+        ELet (b, with_type b.typ (EField (scrut, f)), self#visit_expr env e)
+    | _ ->
+        EMatch (c, scrut, self#visit_branches env branches)
+end
+
 (* General compilation scheme *)
 let everything files =
   let map = build_scheme_map files in
   let files = (compile_simple_matches map)#visit_files () files in
-  let files = if Options.rust () then files else (compile_all_matches map)#visit_files () files in
+  let files = if Options.rust () then remove_one_branch_matches#visit_files () files else (compile_all_matches map)#visit_files () files in
   let files = remove_non_scalar_casts#visit_files () files in
   let files = remove_empty_structs files in
   map, files
diff --git a/lib/OptimizeMiniRust.ml b/lib/OptimizeMiniRust.ml
index 4bb744a7..aa84a708 100644
--- a/lib/OptimizeMiniRust.ml
+++ b/lib/OptimizeMiniRust.ml
@@ -221,7 +221,7 @@ let rec infer (env: env) (expected: typ) (known: known) (e: expr): known * expr
       known, Assign (Index (e1, e2), e3, t)
 
   (* x.f[e2] = e3 *)
-  | Assign (Index (Field (_, f, st) as e1, e2), e3, t) ->
+  | Assign (Index (Field (_, f, st (* optional type *)) as e1, e2), e3, t) ->
       let known = add_mut_field st f known in
       let known, e2 = infer env usize known e2 in
       let known, e3 = infer env t known e3 in
@@ -361,7 +361,8 @@ let rec infer (env: env) (expected: typ) (known: known) (e: expr): known * expr
          we would need to add the expected type of the scrutinee to the Match node,
          similar to what is done for Assign and Field, in order to recurse on
          the scrutinee *)
-      (* FIXME per comment above *)
+      (* FIXME per comment above + need to open binders in match branches + need
+         to mark some of those binders as mutable *)
       let known, arms = List.fold_left_map (fun known (bs, pat, e) ->
           let known, e = infer env expected known e in
           known, (bs, pat, e)