feat(engine): drop only _some_ bodies

As @franziskuskiefer noted, some signature-relevant pieces of
expressions were dropped:
 - constants within types (e.g. a function that returns `[u8; 4]`, `4` was dropped);
 - pre or post-condition on functions;
 - methods in implementations (see #616)

This commit:
 - keeps expr translations on types and methods;
 - disables body drops when a function is marked with hax' attribute
   `AttrPayload::NeverDropBody` (see `hax-lib-macros/types/lib.rs`).

engine/lib/import_thir.ml

@@ -200,6 +200,7 @@ let cast_name_for_type typ_name =
 
 module type EXPR = sig
   val c_expr : Thir.decorated_for__expr_kind -> expr
+  val c_expr_drop_body : Thir.decorated_for__expr_kind -> expr
   val c_ty : Thir.span -> Thir.ty -> ty
   val c_generic_value : Thir.span -> Thir.generic_arg -> generic_value
   val c_generics : Thir.generics -> generics
@@ -215,7 +216,6 @@ end
 
 module Make (CTX : sig
   val is_core_item : bool
-  val drop_body : bool
 end) : EXPR = struct
   let c_binop (op : Thir.bin_op) (lhs : expr) (rhs : expr) (span : span)
       (typ : ty) =
@@ -317,7 +317,7 @@ end) : EXPR = struct
     U.call' (`Concrete name) [ lhs; rhs ] span typ
 
   let rec c_expr (e : Thir.decorated_for__expr_kind) : expr =
-    try (if CTX.drop_body then c_expr_drop_body else c_expr_unwrapped) e
+    try c_expr_unwrapped e
     with Diagnostics.SpanFreeError.Exn (Data (ctx, kind)) ->
       let typ : ty =
         try c_ty e.span e.ty
@@ -1086,7 +1086,6 @@ end
 include struct
   open Make (struct
     let is_core_item = false
-    let drop_body = false
   end)
 
   let import_ty : Types.span -> Types.ty -> Ast.Rust.ty = c_ty
@@ -1099,21 +1098,18 @@ include struct
     c_predicate_kind
 end
 
-(** Instantiate the functor for translating expressions. Two options
-can configured: the crate name (there are special handling related to
-`core`) and `drop_body` (if enabled, original expressions are not
-translated, and the dummy global variable `dropped_body` is used
-instead) *)
-let make ~krate ~drop_body : (module EXPR) =
+(** Instantiate the functor for translating expressions. The crate
+name can be configured (there are special handling related to `core`)
+*)
+let make ~krate : (module EXPR) =
   let is_core_item = String.(krate = "core" || krate = "core_hax_model") in
   let module M : EXPR = Make (struct
     let is_core_item = is_core_item
-    let drop_body = drop_body
   end) in
   (module M)
 
-let c_trait_item (item : Thir.trait_item) ~drop_body : trait_item =
-  let open (val make ~krate:item.owner_id.krate ~drop_body : EXPR) in
+let c_trait_item (item : Thir.trait_item) : trait_item =
+  let open (val make ~krate:item.owner_id.krate : EXPR) in
   let { params; constraints } = c_generics item.generics in
   (* TODO: see TODO in impl items *)
   let ti_ident = Concrete_ident.of_def_id Field item.owner_id in
@@ -1250,15 +1246,21 @@ let rec c_item ~ident ~drop_body (item : Thir.item) : item list =
     [ make_hax_error_item span ident error ]
 
 and c_item_unwrapped ~ident ~drop_body (item : Thir.item) : item list =
-  let open (val make ~krate:item.owner_id.krate ~drop_body : EXPR) in
+  let open (val make ~krate:item.owner_id.krate : EXPR) in
   if should_skip item.attributes then []
   else
     let span = Span.of_thir item.span in
-    let mk_one v =
-      let attrs = c_item_attrs item.attributes in
-      { span; v; ident; attrs }
-    in
+    let attrs = c_item_attrs item.attributes in
+    let mk_one v = { span; v; ident; attrs } in
     let mk v = [ mk_one v ] in
+    let drop_body =
+      drop_body
+      && Attr_payloads.payloads attrs
+         |> List.exists
+              ~f:(fst >> [%matches? (NeverDropBody : Types.ha_payload)])
+         |> not
+    in
+    let c_body = if drop_body then c_expr_drop_body else c_expr in
     (* TODO: things might be unnamed (e.g. constants) *)
     match (item.kind : Thir.item_kind) with
     | Const (_, body) ->
@@ -1268,7 +1270,7 @@ and c_item_unwrapped ~ident ~drop_body (item : Thir.item) : item list =
                name =
                  Concrete_ident.of_def_id Value (Option.value_exn item.def_id);
                generics = { params = []; constraints = [] };
-               body = c_expr body;
+               body = c_body body;
                params = [];
              }
     | TyAlias (ty, generics) ->
@@ -1291,7 +1293,7 @@ and c_item_unwrapped ~ident ~drop_body (item : Thir.item) : item list =
                name =
                  Concrete_ident.of_def_id Value (Option.value_exn item.def_id);
                generics = c_generics generics;
-               body = c_expr body;
+               body = c_body body;
                params;
              }
     | Enum (variants, generics, repr) ->
@@ -1357,7 +1359,6 @@ and c_item_unwrapped ~ident ~drop_body (item : Thir.item) : item list =
         let def_id = Option.value_exn item.def_id in
         let is_struct = true in
         (* repeating the attributes of the item in the variant: TODO is that ok? *)
-        let attrs = c_item_attrs item.attributes in
         let v =
           let kind = Concrete_ident.Kind.Constructor { is_struct } in
           let name = Concrete_ident.of_def_id kind def_id in
@@ -1407,7 +1408,7 @@ and c_item_unwrapped ~ident ~drop_body (item : Thir.item) : item list =
         in
         let params = self :: params in
         let generics = { params; constraints } in
-        let items = List.map ~f:(c_trait_item ~drop_body) items in
+        let items = List.map ~f:c_trait_item items in
         mk @@ Trait { name; generics; items }
     | Trait (Yes, _, _, _, _) -> unimplemented [ item.span ] "Auto trait"
     | Trait (_, Unsafe, _, _, _) -> unimplemented [ item.span ] "Unsafe trait"
@@ -1552,7 +1553,6 @@ and c_item_unwrapped ~ident ~drop_body (item : Thir.item) : item list =
                 ];
           }
         in
-        let attrs = c_item_attrs item.attributes in
         [ { span; v; ident = Concrete_ident.of_def_id Value def_id; attrs } ]
     | ExternCrate _ | Static _ | Macro _ | Mod _ | ForeignMod _ | GlobalAsm _
     | OpaqueTy _ | Union _ | TraitAlias _ ->

hax-lib-macros/src/lib.rs

@@ -144,7 +144,8 @@ pub fn lemma(attr: pm::TokenStream, item: pm::TokenStream) -> pm::TokenStream {
             );
         }
     }
-    quote! { #attr #item }.into()
+    use AttrPayload::NeverDropBody;
+    quote! { #attr #NeverDropBody #item }.into()
 }
 
 /*

hax-lib-macros/src/utils.rs

@@ -187,6 +187,7 @@ pub fn make_fn_decoration(
         } else {
             quote! {}
         };
+        use AttrPayload::NeverDropBody;
         quote! {
             #[cfg(#DebugOrHaxCfgExpr)]
             #late_skip
@@ -197,6 +198,7 @@ pub fn make_fn_decoration(
                 #uid_attr
                 #late_skip
                 #[allow(unused)]
+                #NeverDropBody
                 #decoration_sig {
                     #phi
                 }

hax-lib-macros/types/src/lib.rs

@@ -71,6 +71,10 @@ pub enum AttrPayload {
         item: ItemUid,
     },
     Uid(ItemUid),
+    /// Mark an item so that hax never drop its body (this is useful
+    /// for pre- and post- conditions of a function we dropped the
+    /// body of: pre and post are part of type signature)
+    NeverDropBody,
     /// Mark an item as a lemma statement to prove in the backend
     Lemma,
     Language,

test-harness/src/snapshots/toolchain__interface-only into-fstar.snap

@@ -35,5 +35,11 @@ module Interface_only
 open Core
 open FStar.Mul
 
-val f (x: u8) : Prims.Pure Prims.unit Prims.l_True (fun _ -> Prims.l_True)
+val f (x: u8)
+    : Prims.Pure (t_Array u8 (sz 4))
+      (requires x <. 254uy)
+      (ensures
+        fun r ->
+          let r:t_Array u8 (sz 4) = r in
+          (r.[ sz 0 ] <: u8) >. x)
 '''

tests/cli/interface-only/Cargo.toml

@@ -4,6 +4,8 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+hax-lib-macros = { path = "../../../hax-lib-macros" }
+hax-lib = { path = "../../../hax-lib" }
 
 [package.metadata.hax-tests]
 into."fstar" = { include-flag = "+:** -interface_only::Foo" }

tests/cli/interface-only/src/lib.rs

@@ -1,14 +1,23 @@
 #![allow(dead_code)]
 
+use hax_lib as hax;
+
 /// This item contains unsafe blocks and raw references, two features
 /// not supported by hax. Thanks to the `-i` flag and the `+:`
 /// modifier, `f` is still extractable as an interface.
-fn f(x: u8) {
+///
+/// Expressions within type are still extracted, as well as pre- and
+/// post-conditions.
+#[hax::requires(x < 254)]
+#[hax::ensures(|r| r[0] > x)]
+fn f(x: u8) -> [u8; 4] {
     let y = x as *const i8;
 
     unsafe {
         println!("{}", *y);
     }
+
+    [x + 1, x, x, x]
 }
 
 /// This struct contains a field which uses raw pointers, which are