Better handling of comments

lib/Ast.ml

@@ -120,12 +120,17 @@ and branch_t =
 and fields_t =
   (ident * (typ * bool)) list
 
+type node_meta =
+  | NodeComment of string
+
+and node_meta' = node_meta [@visitors.opaque] [@@deriving show]
 
 (* This type, by virtue of being separated from the recursive definition of expr
  * and pattern, generates no implementation. We provide our own below. *)
 type 'a with_type = {
   node: 'a;
-  mutable typ: typ
+  mutable typ: typ;
+  meta: node_meta' list;
     (** Filled in by [Checker] *)
 }
   [@@deriving show]
@@ -163,7 +168,7 @@ class ['self] map_typ_adapter = object (self: 'self)
     fun f (env, _) x ->
       let typ = self#visit_typ env x.typ in
       let node = f (env, typ) x.node in
-      { node; typ }
+      { node; typ; meta = x.meta }
 end
 
 class ['self] iter_typ_adapter = object (self: 'self)
@@ -381,7 +386,7 @@ class ['self] map_expr_adapter = object (self: 'self)
     fun f env x ->
       let typ = self#visit_typ env x.typ in
       let node = f (env, typ) x.node in
-      { node; typ }
+      { node; typ; meta = x.meta }
 
   method visit_expr_w env e =
     self#visit_expr (env, e.typ) e
@@ -679,7 +684,7 @@ let map_decls f files =
   List.map (fun (file, decls) -> file, List.map f decls) files
 
 let with_type typ node =
-  { typ; node }
+  { typ; node; meta = [] }
 
 let flatten_tapp t =
   let rec flatten_tapp cgs t =

lib/AstToMiniRust.ml

@@ -567,7 +567,7 @@ and translate_array (env: env) is_toplevel (init: Ast.expr): env * MiniRust.expr
 (* However, generally, we will have a type provided by the context that may
    necessitate the insertion of conversions *)
 and translate_expr_with_type (env: env) (e: Ast.expr) (t_ret: MiniRust.typ): env * MiniRust.expr =
-  (* KPrint.bprintf "translate_expr_with_type: %a @@ %a\n" PrintMiniRust.ptyp t_ret PrintAst.Ops.pexpr e; *)
+  KPrint.bprintf "translate_expr_with_type: %a @@ %a\n" PrintMiniRust.ptyp t_ret PrintAst.Ops.pexpr e;
 
   let erase_lifetime_info = (object(self)
     inherit [_] MiniRust.DeBruijn.map
@@ -641,9 +641,10 @@ and translate_expr_with_type (env: env) (e: Ast.expr) (t_ret: MiniRust.typ): env
         if erase_lifetime_info t = erase_lifetime_info t_ret then
           x
         else
-          Warn.failwith "type mismatch;\n  e=%a\n  t=%a\n  t_ret=%a\n  x=%a"
+          Warn.failwith "type mismatch;\n  e=%a\n  t=%a (verbose: %s)\n  t_ret=%a\n  x=%a"
             PrintAst.Ops.pexpr e
-            PrintMiniRust.ptyp t PrintMiniRust.ptyp t_ret
+            PrintMiniRust.ptyp t (MiniRust.show_typ t)
+            PrintMiniRust.ptyp t_ret
             PrintMiniRust.pexpr x;
     end
   in
@@ -721,7 +722,7 @@ and translate_expr_with_type (env: env) (e: Ast.expr) (t_ret: MiniRust.typ): env
       assert (cgs @ cgs' = []);
       let es =
         match es with
-        | [ { typ = TUnit; node } ] -> assert (node = EUnit); []
+        | [ { typ = TUnit; node; _ } ] -> assert (node = EUnit); []
         | _ -> es
       in
       let env, e = translate_expr env e in
@@ -739,7 +740,7 @@ and translate_expr_with_type (env: env) (e: Ast.expr) (t_ret: MiniRust.typ): env
   | EApp (e0, es) ->
       let es =
         match es with
-        | [ { typ = TUnit; node } ] -> assert (node = EUnit); []
+        | [ { typ = TUnit; node; _ } ] -> assert (node = EUnit); []
         | _ -> es
       in
       let env, e0 = translate_expr env e0 in
@@ -1173,9 +1174,10 @@ let bind_decl env (d: Ast.decl): env =
           (* These sets are mutually exclusive, so we don't box *and* introduce a
              lifetime at the same time *)
           let box = Idents.LidSet.mem lid env.heap_structs in
-          KPrint.bprintf "%a: lifetime=%b\n" PrintAst.Ops.plid lid (Idents.LidSet.mem lid env.pointer_holding_structs);
+          let lifetime = Idents.LidSet.mem lid env.pointer_holding_structs in
+          KPrint.bprintf "%a (FLAT): lifetime=%b box=%b\n" PrintAst.Ops.plid lid lifetime box;
           let lifetime =
-            if Idents.LidSet.mem lid env.pointer_holding_structs then
+            if lifetime then
               Some (MiniRust.Label "a")
             else
               None
@@ -1186,6 +1188,9 @@ let bind_decl env (d: Ast.decl): env =
           ) fields in
           { env with struct_fields = LidMap.add lid fields env.struct_fields }
       | Variant branches ->
+          let box = Idents.LidSet.mem lid env.heap_structs in
+          let lifetime = Idents.LidSet.mem lid env.pointer_holding_structs in
+          KPrint.bprintf "%a (VARIANT): lifetime=%b box=%b\n" PrintAst.Ops.plid lid lifetime box;
           List.fold_left (fun env (cons, fields) ->
             { env with struct_fields = LidMap.add (fst lid, snd lid ^ "_" ^
             cons) (List.map (fun (f, (t, _)) ->

lib/DataTypes.ml

@@ -137,7 +137,7 @@ let remove_unused_type_arguments files =
             chop (kept + 1) (i + 1) (def, binders, ret)
           else
             let def = DeBruijn.subst_te TAny (n - i - 1) def in
-            let binders = List.map (fun { node; typ } -> { node; typ = DeBruijn.subst_t TAny (n - i - 1) typ }) binders in
+            let binders = List.map (fun { node; typ; _ } -> { node; typ = DeBruijn.subst_t TAny (n - i - 1) typ; meta = [] }) binders in
             let ret = DeBruijn.subst_t TAny (n - i - 1) ret in
             chop kept (i + 1) (def, binders, ret)
       in
@@ -511,17 +511,17 @@ let compile_simple_matches (map, enums) = object(self)
   method! visit_expr_w env x =
     let node = self#visit_expr' (env, x.typ) x.node in
     let typ = self#visit_typ env x.typ in
-    { node; typ }
+    { node; typ; meta = [] }
 
   method! visit_pattern_w env x =
     let node = self#visit_pattern' (env, x.typ) x.node in
     let typ = self#visit_typ env x.typ in
-    { node; typ }
+    { node; typ; meta = [] }
 
   method! visit_with_type f (env, _) x =
       let node = f (env, x.typ) x.node in
       let typ = self#visit_typ env x.typ in
-      { node; typ }
+      { node; typ; meta = [] }
 
   method! visit_TQualified _ lid =
     match Hashtbl.find map lid with

lib/DeBruijn.ml

@@ -234,7 +234,7 @@ let close_branch bs p e =
 let opening_binder b =
   let a = Atom.fresh () in
   let b = { b with node = { b.node with atom = a } } in
-  b, subst { node = EOpen (b.node.name, a); typ = b.typ } 0
+  b, subst { node = EOpen (b.node.name, a); typ = b.typ; meta = [] } 0
 
 let open_binder b e1 =
   let b, f = opening_binder b in
@@ -254,7 +254,7 @@ let open_branch bs pat expr =
   List.fold_right (fun binder (bs, pat, expr) ->
     let b, expr = open_binder binder expr in
     let pat =
-      subst_p { node = POpen (b.node.name, b.node.atom); typ = b.typ } 0 pat
+      subst_p { node = POpen (b.node.name, b.node.atom); typ = b.typ; meta = [] } 0 pat
     in
     b :: bs, pat, expr
   ) bs ([], pat, expr)

lib/Helpers.ml

@@ -71,7 +71,7 @@ let pwild = with_type TAny PWild
 
 let mk_op op w =
   { node = EOp (op, w);
-    typ = type_of_op op w }
+    typ = type_of_op op w; meta = [] }
 
 (* @0 < <finish> *)
 let mk_lt w finish =
@@ -160,13 +160,13 @@ let unused_binding = sequence_binding
 let mk_binding ?(mut=false) name t =
   let b = fresh_binder name t in
   { b with node = { b.node with mut } },
-  { node = EOpen (b.node.name, b.node.atom); typ = t }
+  { node = EOpen (b.node.name, b.node.atom); typ = t; meta = [] }
 
 (** Generates "let [[name]]: [[t]] = [[e]] in [[name]]" *)
 let mk_named_binding name t e =
   let b, ref = mk_binding name t in
   b,
-  { node = e; typ = t },
+  { node = e; typ = t; meta = [] },
   ref
 
 
@@ -411,7 +411,7 @@ let rec is_initializer_constant e =
   match e with
   | { node = EAddrOf { node = EQualified _; _ }; _ } ->
       true
-  | { node = EQualified _; typ = t } ->
+  | { node = EQualified _; typ = t; _ } ->
       is_address t
   | { node = EEnum _; _ } ->
       true
@@ -688,7 +688,7 @@ let rec strip_cast e =
 let rec nest bs t e2 =
   match bs with
   | (b, e1) :: bs ->
-      { node = ELet (b, e1, close_binder b (nest bs t e2)); typ = t }
+      { node = ELet (b, e1, close_binder b (nest bs t e2)); typ = t; meta = [] }
   | [] ->
       e2
 
@@ -715,30 +715,30 @@ let rec nest_in_return_pos i typ f e =
   match e.node with
   | ELet (b, e1, e2) ->
       let e2 = nest_in_return_pos (i + 1) typ f e2 in
-      { node = ELet (b, e1, e2); typ }
+      { node = ELet (b, e1, e2); typ; meta = [] }
   | EIfThenElse (e1, e2, e3) ->
       let e2 = nest_in_return_pos i typ f e2 in
       let e3 = nest_in_return_pos i typ f e3 in
-      { node = EIfThenElse (e1, e2, e3); typ }
+      { node = EIfThenElse (e1, e2, e3); typ; meta = [] }
   | ESwitch (e, branches) ->
       let branches = List.map (fun (t, e) ->
         t, nest_in_return_pos i typ f e
       ) branches in
-      { node = ESwitch (e, branches); typ }
+      { node = ESwitch (e, branches); typ; meta = [] }
   | EMatch (c, e, branches) ->
       { node =
         EMatch (c, e, List.map (fun (bs, p, e) ->
           bs, p, nest_in_return_pos (i + List.length bs) typ f e
         ) branches);
-        typ }
+        typ; meta = [] }
   | ESequence es ->
       let l = List.length es in
       { node = ESequence (List.mapi (fun j e ->
           if j = l - 1 then
             nest_in_return_pos i typ f e
           else
             e
-        ) es); typ }
+        ) es); typ; meta = [] }
   | _ ->
       f i e
 

lib/InputAstToAst.ml

@@ -9,7 +9,7 @@ open Common
 module I = InputAst
 
 let mk (type a) (node: a): a with_type =
-  { node; typ = TAny }
+  { node; typ = TAny; meta = [] }
 
 let rec binders_of_pat p =
   let open I in
@@ -147,13 +147,13 @@ and mk_expr = function
   | I.EApp (I.ETApp (I.EQualified ([ "Steel"; "Reference" ], "null"), [ t ]), [ I.EUnit ])
   | I.EApp (I.ETApp (I.EQualified ([ "C"; "Nullity" ], "null"), [ t ]), [ I.EUnit ])
   | I.EBufNull t ->
-      { node = EBufNull; typ = TBuf (mk_typ t, false) }
+      { node = EBufNull; typ = TBuf (mk_typ t, false); meta = [] }
 
   | I.EApp (I.ETApp (I.EQualified ( [ "LowStar"; "Monotonic"; "Buffer" ], "is_null"), [ t; _; _ ]), [ e ])
   | I.EApp (I.ETApp (I.EQualified ( [ "Steel"; "Reference" ], "is_null"), [ t ]), [ e ]) ->
       mk (EApp (mk (EPolyComp (K.PEq, TBuf (mk_typ t, false))), [
         mk_expr e;
-        { node = EBufNull; typ = TBuf (mk_typ t, false) }]))
+        { node = EBufNull; typ = TBuf (mk_typ t, false); meta = [] }]))
 
   | I.EApp (e, es) ->
       mk (EApp (mk_expr e, List.map mk_expr es))
@@ -211,18 +211,18 @@ and mk_expr = function
   | I.EAbortS s ->
       mk (EAbort (None, Some s))
   | I.EAbortT (s, t) ->
-      { node = EAbort (Some (mk_typ t), Some s); typ = mk_typ t }
+      { node = EAbort (Some (mk_typ t), Some s); typ = mk_typ t; meta = [] }
   | I.EReturn e ->
       mk (EReturn (mk_expr e))
   | I.EFlat (tname, fields) ->
-      { node = EFlat (mk_fields fields); typ = mk_typ tname }
+      { node = EFlat (mk_fields fields); typ = mk_typ tname; meta = [] }
   | I.EField (tname, e, field) ->
       let e = { (mk_expr e) with typ = mk_typ tname } in
       mk (EField (e, field))
   | I.ETuple es ->
       mk (ETuple (List.map mk_expr es))
   | I.ECons (lid, id, es) ->
-      { node = ECons (id, List.map mk_expr es); typ = mk_typ lid }
+      { node = ECons (id, List.map mk_expr es); typ = mk_typ lid; meta = [] }
   | I.EFun (bs, e, t) ->
       mk (EFun (mk_binders bs, mk_expr e, mk_typ t))
   | I.EComment (before, e, after) ->

lib/Monomorphization.ml

@@ -612,8 +612,8 @@ let functions files =
                     (* binders are the remaining binders after the cg-binders have been eliminated *)
                     let diff = List.length binders - List.length cgs in
                     let _cg_binders, binders = KList.split (List.length cgs + List.length cgs') binders in
-                    let binders = List.map (fun { node; typ } ->
-                      { node; typ = DeBruijn.(subst_ctn diff cgs (subst_tn ts typ)) }
+                    let binders = List.map (fun { node; typ; _ } ->
+                      { node; typ = DeBruijn.(subst_ctn diff cgs (subst_tn ts typ)); meta = [] }
                     ) binders in
                     (* KPrint.bprintf "about to substitute:\n  e=%a\n  cgs=%a\n cgs'=%a\n  ts=%a\n  head type=%a\n%a\n" *)
                     (*   pexpr e *)

lib/PrintAst.ml

@@ -156,7 +156,7 @@ and print_flag = function
   | Target s ->
       string ("__attribute__((target = "^s^"))")
 
-and print_binder { typ; node = { name; mut; meta; mark; _ }} =
+and print_binder { typ; node = { name; mut; meta; mark; _ }; _ } =
   let o, u = !mark in
   (if mut then string "mutable" ^^ break 1 else empty) ^^
   group (group (string name ^^ lparen ^^ string (Mark.show_occurrence o) ^^ comma ^^
@@ -212,8 +212,12 @@ and print_let_binding (binder, e1) =
   group (group (string "let" ^/^ print_binder binder ^/^ equals) ^^
   jump (print_expr e1))
 
-and print_expr { node; typ } =
+and print_expr { node; typ; meta } =
   (* print_typ typ ^^ colon ^^ space ^^ parens @@ *)
+  begin match List.filter_map (function NodeComment s -> Some s (* | _ -> None*) ) meta with
+  | [] -> empty
+  | s -> string (String.concat "\n" s) ^^ break1
+  end ^^
   match node with
   | EComment (s, e, s') ->
       surround 2 1 (string s) (print_expr e) (string s')

lib/PrintMiniRust.ml

@@ -184,6 +184,7 @@ let rec print_typ env (t: typ): document =
       let env = push_n_type env n in
       group @@
       group (parens (separate_map (comma ^^ break1) (print_typ env) ts)) ^/^
+      minus ^^ rangle ^^
       print_typ env t
   | Bound n ->
       begin try