Companion fix to FStarLang/karamel#483

bin/main.ml

@@ -250,6 +250,7 @@ Supported options:|}
   let files = Krml.Simplify.hoist#visit_files [] files in
   Eurydice.Logging.log "Phase2.75" "%a" pfiles files;
   let files = Krml.Simplify.fixup_hoist#visit_files () files in
+  (*good*)
   Eurydice.Logging.log "Phase2.8" "%a" pfiles files;
   let files = Krml.Simplify.misc_cosmetic#visit_files () files in
   let files = Krml.Simplify.let_to_sequence#visit_files () files in

lib/AstOfLlbc.ml

@@ -1485,12 +1485,12 @@ let rec expression_of_raw_statement (env : env) (ret_var : C.var_id) (s : C.raw_
   | Continue _ -> K.(with_type TAny EContinue)
   | Nop -> Krml.Helpers.eunit
   | Sequence (s1, s2) ->
-      let e1 = expression_of_raw_statement env ret_var s1.content in
-      let e2 = expression_of_raw_statement env ret_var s2.content in
+      let e1 = expression_of_statement env ret_var s1 in
+      let e2 = expression_of_statement env ret_var s2 in
       K.(with_type e2.typ (ESequence [ e1; e2 ]))
   | Switch (If (op, s1, s2)) ->
-      let e1 = expression_of_raw_statement env ret_var s1.content in
-      let e2 = expression_of_raw_statement env ret_var s2.content in
+      let e1 = expression_of_statement env ret_var s1 in
+      let e2 = expression_of_statement env ret_var s2 in
       let t = lesser e1.typ e2.typ in
       K.(with_type t (EIfThenElse (expression_of_operand env op, e1, e2)))
   | Switch (SwitchInt (scrutinee, _int_ty, branches, default)) ->
@@ -1501,10 +1501,10 @@ let rec expression_of_raw_statement (env : env) (ret_var : C.var_id) (s : C.raw_
             List.map
               (fun sv ->
                 ( K.SConstant (constant_of_scalar_value sv),
-                  expression_of_raw_statement env ret_var stmt.C.content ))
+                  expression_of_statement env ret_var stmt ))
               svs)
           branches
-        @ [ K.SWild, expression_of_raw_statement env ret_var default.C.content ]
+        @ [ K.SWild, expression_of_statement env ret_var default ]
       in
       let t = Krml.KList.reduce lesser (List.map (fun (_, e) -> e.K.typ) branches) in
       K.(with_type t (ESwitch (scrutinee, branches)))
@@ -1533,7 +1533,7 @@ let rec expression_of_raw_statement (env : env) (ret_var : C.var_id) (s : C.raw_
                 let variant_name, n_fields = variant_name_of_variant_id variant_id in
                 let dummies = List.init n_fields (fun _ -> K.(with_type TAny PWild)) in
                 let pat = K.with_type p.typ (K.PCons (variant_name, dummies)) in
-                [], pat, expression_of_raw_statement env ret_var e.C.content)
+                [], pat, expression_of_statement env ret_var e)
               variant_ids)
           branches
       in
@@ -1545,7 +1545,7 @@ let rec expression_of_raw_statement (env : env) (ret_var : C.var_id) (s : C.raw_
             [
               ( [],
                 K.with_type p.typ K.PWild,
-                expression_of_raw_statement env ret_var default.C.content );
+                expression_of_statement env ret_var default);
             ]
         | None -> []
       in
@@ -1554,9 +1554,13 @@ let rec expression_of_raw_statement (env : env) (ret_var : C.var_id) (s : C.raw_
   | Loop s ->
       K.(
         with_type TUnit
-          (EWhile (Krml.Helpers.etrue, expression_of_raw_statement env ret_var s.content)))
+          (EWhile (Krml.Helpers.etrue, expression_of_statement env ret_var s)))
   | Error _ -> failwith "TODO: error"
 
+and expression_of_statement (env: env) (ret_var: C.var_id) (s: C.statement): K.expr =
+  { (expression_of_raw_statement env ret_var s.content) with
+    meta = List.map (fun x -> K.CommentBefore x) s.comments_before }
+
 (** Top-level declarations: orchestration *)
 
 let of_declaration_group (dg : 'id C.g_declaration_group) (f : 'id -> 'a) : 'a list =
@@ -1749,7 +1753,7 @@ let decl_of_id (env : env) (id : C.any_decl_id) : K.decl option =
                 let body =
                   try
                     with_locals env return_type (return_var :: locals) (fun env ->
-                        expression_of_raw_statement env return_var.index body.content)
+                        expression_of_statement env return_var.index body)
                   with e ->
                     let msg =
                       Krml.KPrint.bsprintf "Eurydice error: %s\n%s" (Printexc.to_string e)
@@ -1798,7 +1802,7 @@ let decl_of_id (env : env) (id : C.any_decl_id) : K.decl option =
             let ret_var = List.hd body.locals in
             let body =
               with_locals env ty body.locals (fun env ->
-                  expression_of_raw_statement env ret_var.index body.body.content)
+                  expression_of_statement env ret_var.index body.body)
             in
             Some (K.DGlobal ([ Krml.Common.Const "" ], lid_of_name env name, 0, ty, body))
         | None -> Some (K.DExternal (None, [], 0, 0, lid_of_name env name, ty, []))

lib/Cleanup1.ml

@@ -14,48 +14,58 @@ let count_atoms =
     method! visit_EOpen _ _ a = AtomSet.singleton a
   end
 
+type remove_env = (string * typ * node_meta list) AtomMap.t
+
+let pmeta buf ({ meta; _ }: 'a with_type) =
+  List.iter (function
+    | CommentBefore s | CommentAfter s -> Buffer.add_string  buf s;
+    Buffer.add_char buf '\n') meta
+
+let mk typ meta node = { node; typ; meta }
+
 let remove_assignments =
   object (self)
     inherit [_] map
 
-    method private peel_lets to_close e =
+    method private peel_lets (to_close: remove_env) e =
       match e.node with
       | ELet (b, e1, e2) ->
           (if not (e1.node = EAny || e1.node = EUnit) then
              Krml.(Warn.fatal_error "Initializer of let-binding is %a" PrintAst.Ops.pexpr e1));
           (* Krml.(KPrint.bprintf "peeling %s\n" b.node.name); *)
           let b, e2 = open_binder b e2 in
-          let to_close = AtomMap.add b.node.atom (b.node.name, b.typ) to_close in
+          (* Krml.KPrint.bprintf "peel: let-binding meta %a\n" pmeta b; *)
+          let to_close = AtomMap.add b.node.atom (b.node.name, b.typ, b.meta) to_close in
           self#peel_lets to_close e2
       | _ ->
           let e = Krml.Simplify.sequence_to_let#visit_expr_w () e in
           (* Krml.(KPrint.bprintf "after peeling:\n%a" PrintAst.Ops.ppexpr e); *)
           self#visit_expr_w to_close e
 
-    method! visit_DFunction to_close cc flags n_cgs n t name bs e =
+    method! visit_DFunction (to_close: remove_env) cc flags n_cgs n t name bs e =
       (* Krml.(KPrint.bprintf "visiting %a\n" PrintAst.Ops.plid name); *)
       assert (AtomMap.is_empty to_close);
       DFunction (cc, flags, n_cgs, n, t, name, bs, self#peel_lets to_close e)
 
-    method! visit_DGlobal to_close flags n t name e =
+    method! visit_DGlobal (to_close: remove_env) flags n t name e =
       assert (AtomMap.is_empty to_close);
       DGlobal (flags, n, t, name, self#peel_lets to_close e)
 
-    method! visit_ELet (not_yet_closed, t) b e1 e2 =
+    method! visit_ELet ((not_yet_closed: remove_env), t) b e1 e2 =
       (* If [not_yet_closed] represents the set of bindings that have yet to be
          closed (i.e. for which we have yet to insert a let-binding, as close as
          possible to the first use-site), and [candidates] represents the atoms
          that we know for sure must be closed right now, then [close_now_over]
          inserts suitable let-bindings for the candidates that have not yet been
          closed, then calls the continuation with the remaining subset of
          not_yet_closed. *)
-      let close_now_over not_yet_closed candidates mk_node =
+      let close_now_over (not_yet_closed: remove_env) candidates mk_node =
         let to_close_now = AtomSet.inter candidates (set_of_map_keys not_yet_closed) in
         let bs = List.of_seq (AtomSet.to_seq to_close_now) in
         let bs =
           List.map
             (fun atom ->
-              let name, typ = AtomMap.find atom not_yet_closed in
+              let name, typ, meta = AtomMap.find atom not_yet_closed in
               ( {
                   node =
                     {
@@ -67,6 +77,7 @@ let remove_assignments =
                       attempt_inline = false;
                     };
                   typ;
+                  meta;
                 },
                 if typ = TUnit then
                   Krml.Helpers.eunit
@@ -80,6 +91,7 @@ let remove_assignments =
           AtomMap.filter (fun a _ -> not (AtomSet.mem a to_close_now)) not_yet_closed
         in
         let node = mk_node not_yet_closed in
+        (* XXX why is the with_type necessary here?? *)
         (Krml.Helpers.nest bs t (with_type t node)).node
       in
 
@@ -91,14 +103,14 @@ let remove_assignments =
          the general case if it's a let-node; special treatment if it's
          control-flow (match, if, while); otherwise, just close everything now and
          move on (wildcard case). *)
-      let rec recurse_or_close not_yet_closed e =
-        let t = e.typ in
-        match e.node with
+      let rec recurse_or_close (not_yet_closed: remove_env) e0 =
+        let t = e0.typ in
+        match e0.node with
         | ELet _ ->
             (* let node: restart logic and jump back to match below *)
-            self#visit_expr_w not_yet_closed e
+            self#visit_expr_w not_yet_closed e0
         | EIfThenElse (e, e', e'') ->
-            with_type t
+            mk t e0.meta
             @@ close_now_over not_yet_closed
                  ((* We must now bind: *)
                   count e
@@ -114,11 +126,11 @@ let remove_assignments =
                        recurse_or_close not_yet_closed e',
                        recurse_or_close not_yet_closed e'' ))
         | EWhile (e, e') ->
-            with_type t
+            mk t e0.meta
             @@ close_now_over not_yet_closed (count e) (fun not_yet_closed ->
                    EWhile (self#visit_expr_w not_yet_closed e, recurse_or_close not_yet_closed e'))
         | ESwitch (e, branches) ->
-            with_type t
+            mk t e0.meta
             @@ close_now_over not_yet_closed
                  ((* We must now bind: *)
                   count e
@@ -130,7 +142,7 @@ let remove_assignments =
                      ( self#visit_expr_w not_yet_closed e,
                        List.map (fun (p, e) -> p, recurse_or_close not_yet_closed e) branches ))
         | EMatch (c, e, branches) ->
-            with_type t
+            mk t e0.meta
             @@ close_now_over not_yet_closed
                  ((* We must now bind: *)
                   count e
@@ -149,19 +161,19 @@ let remove_assignments =
                an assignment in terminal position, *and* the variable has yet to
                be closed, it means that the assignment is useless since no one
                else will be using the variable after that. *)
-            with_type e.typ
-              (close_now_over not_yet_closed (count e) (fun not_yet_closed ->
+            mk t e0.meta
+            @@ close_now_over not_yet_closed (count e0) (fun not_yet_closed ->
                    (* not_yet_closed should be empty at this stage *)
-                   (self#visit_expr_w not_yet_closed e).node))
+                (self#visit_expr_w not_yet_closed e0).node)
       in
 
       match e1.node with
-      | EAssign ({ node = EOpen (_, atom); _ }, e1) when AtomMap.mem atom not_yet_closed ->
-          close_now_over not_yet_closed (count e1) (fun not_yet_closed ->
+      | EAssign ({ node = EOpen (_, atom); _ }, e_rhs) when AtomMap.mem atom not_yet_closed ->
+          close_now_over not_yet_closed (count e_rhs) (fun not_yet_closed ->
               (* Combined "close now" (above) + let-binding insertion in lieu of the assignment *)
               assert (b.node.meta = Some MetaSequence);
               let e2 = snd (open_binder b e2) in
-              let name, typ = AtomMap.find atom not_yet_closed in
+              let name, typ, meta = AtomMap.find atom not_yet_closed in
               let b =
                 {
                   node =
@@ -174,11 +186,13 @@ let remove_assignments =
                       attempt_inline = false;
                     };
                   typ;
+                  meta = meta @ e1.meta;
                 }
               in
               let not_yet_closed = AtomMap.remove atom not_yet_closed in
+              (* Krml.(KPrint.bprintf "rebuilt: %a\n" PrintAst.Ops.pexpr (with_type TUnit (ELet (b, e_rhs, e2)))); *)
               let e2 = self#visit_expr_w not_yet_closed (close_binder b e2) in
-              ELet (b, e1, e2))
+              ELet (b, e_rhs, e2))
       | EIfThenElse (e, e', e'') ->
           assert (b.node.meta = Some MetaSequence);
           close_now_over not_yet_closed
@@ -192,7 +206,7 @@ let remove_assignments =
               (fun not_yet_closed ->
               ELet
                 ( b,
-                  with_type e1.typ
+                  mk e1.typ e1.meta
                     (EIfThenElse
                        ( self#visit_expr_w not_yet_closed e,
                          recurse_or_close not_yet_closed e',
@@ -207,7 +221,7 @@ let remove_assignments =
             (fun not_yet_closed ->
               ELet
                 ( b,
-                  with_type TUnit
+                  mk e1.typ e1.meta
                     (EWhile (self#visit_expr_w not_yet_closed e, recurse_or_close not_yet_closed e')),
                   recurse_or_close not_yet_closed e2 ))
       | ESwitch (e, branches) ->
@@ -227,7 +241,7 @@ let remove_assignments =
               (fun not_yet_closed ->
               ELet
                 ( b,
-                  with_type e1.typ
+                  mk e1.typ e1.meta
                     (ESwitch
                        ( self#visit_expr_w not_yet_closed e,
                          List.map (fun (p, e) -> p, recurse_or_close not_yet_closed e) branches )),
@@ -248,7 +262,7 @@ let remove_assignments =
               (fun not_yet_closed ->
               ELet
                 ( b,
-                  with_type e1.typ
+                  mk e1.typ e1.meta
                     (EMatch
                        ( c,
                          self#visit_expr_w not_yet_closed e,

lib/Cleanup3.ml

@@ -113,6 +113,7 @@ let bonus_cleanups =
           _,
           ESequence [ { node = EAssign ({ node = EBound 0; _ }, e3); _ }; { node = EBound 0; _ } ] )
         -> (DeBruijn.subst Helpers.eunit 0 e3).node
+
       (* let uu; memcpy(uu, ..., src, ...); e2  -->  let copy_of_src; ... *)
       | ( EAny,
           TArray (_, (_, n)),
@@ -131,13 +132,12 @@ let bonus_cleanups =
               _;
             ] )
         when n = n' && Krml.Helpers.is_uu b.node.name ->
-          EComment
-            ( "Passing arrays by value in Rust generates a copy in C",
-              with_type e2.typ
-                (super#visit_ELet env
-                   { b with node = { b.node with name = "copy_of_" ^ List.nth bs (src - 1) } }
-                   e1 e2),
-              "" )
+          super#visit_ELet env
+             { b with
+                node = { b.node with name = "copy_of_" ^ List.nth bs (src - 1) };
+                meta = [ CommentBefore "Passing arrays by value in Rust generates a copy in C" ] }
+             e1 e2
+
       (* let uu = f(e); y = uu; e2  -->  let y = f(e); e2 *)
       | ( EApp ({ node = EQualified _; _ }, es),
           _,
@@ -183,6 +183,7 @@ let add_extra_type_to_slice_index =
             {
               node = EFlat [ (Some "start", e_start); (Some "end", e_end) ];
               typ = TQualified ([ "core"; "ops"; "range" ], id);
+              _
             };
           ] )
         when KString.starts_with id "Range" ->
@@ -204,6 +205,7 @@ let add_extra_type_to_slice_index =
             {
               node = EFlat [ (Some "start", e_start); (Some "end", e_end) ];
               typ = TQualified ([ "core"; "ops"; "range" ], id);
+              _
             };
           ] )
         when KString.starts_with id "Range" ->

test/array/src/main.rs

@@ -33,13 +33,16 @@ fn mk_incr<const K: usize>() -> [ u32; K ] {
 }
 
 fn main() {
+    // XXX1
     let Foo { x, y } = mk_foo2();
     let expected = 0u32;
     mut_array(x);
+    // XXX2
     mut_foo(Foo { x, y });
     assert_eq!(x[0], expected);
     let a: [ u32; 10 ] = mk_incr();
     let expected = 9;
+    // XXX3
     assert_eq!(a[9], expected);
     // let a: [ u32; 10 ] = mk_incr2();
     // let expected = 10;