feat: allow pre/post on traits (fixes #731)

engine/backends/fstar/fstar_backend.ml

@@ -1186,7 +1186,114 @@ struct
                     let ty = F.term @@ F.AST.Product (inputs, ty_pre_post) in
                     [ (F.id name, None, [], ty) ]
                 | TIFn ty ->
+                    let weakest =
+                      let h kind =
+                        Attrs.associated_fns kind i.ti_attrs |> List.hd
+                      in
+                      Option.first_some (h Ensures) (h Requires)
+                      |> Option.map ~f:(fun (generics, params, expr) ->
+                             let dummy_self =
+                               List.find generics.params
+                                 ~f:[%matches? { kind = GPType _; _ }]
+                               |> Option.value_or_thunk ~default:(fun () ->
+                                      Error.assertion_failure i.ti_span
+                                        ("Expected a first generic of type \
+                                          `Self`. Instead generics params \
+                                          are: "
+                                        ^ [%show: generic_param list]
+                                            generics.params))
+                               |> fun x -> x.ident
+                             in
+                             let self =
+                               Local_ident.{ name = "Self"; id = mk_id Typ 0 }
+                             in
+                             let renamer =
+                               let f (id : local_ident) =
+                                 if [%eq: string] dummy_self.name id.name then
+                                   self
+                                 else id
+                               in
+                               U.Mappers.rename_local_idents f
+                             in
+                             let generics =
+                               renamer#visit_generics () generics
+                             in
+                             let params =
+                               List.map ~f:(renamer#visit_param ()) params
+                             in
+                             let expr = renamer#visit_expr () expr in
+                             (generics, params, expr))
+                    in
                     let ty = pty e.span ty in
+                    let ty =
+                      let variables =
+                        let idents_visitor = U.Reducers.collect_local_idents in
+                        idents_visitor#visit_trait_item () i
+                        :: (Option.map
+                              ~f:(fun (generics, params, expr) ->
+                                [
+                                  idents_visitor#visit_generics () generics;
+                                  idents_visitor#visit_expr () expr;
+                                ]
+                                @ List.map
+                                    ~f:(idents_visitor#visit_param ())
+                                    params)
+                              weakest
+                           |> Option.value ~default:[])
+                        |> Set.union_list (module Local_ident)
+                        |> Set.to_list |> ref
+                      in
+                      Stdlib.prerr_endline @@ "Variables: "
+                      ^ [%show: local_ident list] !variables;
+                      let mk_fresh prefix =
+                        let v = U.fresh_local_ident_in !variables prefix in
+                        variables := v :: !variables;
+                        v
+                      in
+                      let bindings = ref [] in
+                      let f (p : param) =
+                        let name =
+                          match p.pat.p with
+                          | PBinding { var; _ } -> var
+                          | _ ->
+                              let name = mk_fresh "x" in
+                              let ({ span; typ; _ } : pat) = p.pat in
+                              let expr = { e = LocalVar name; span; typ } in
+                              bindings := (p.pat, expr) :: !bindings;
+                              name
+                        in
+                        FStarBinder.of_named_typ p.pat.span name p.typ
+                      in
+                      weakest
+                      |> Option.map ~f:(List.map ~f |> map_snd3)
+                      |> Option.map
+                           ~f:(fun (generics, binders, (expr : expr)) ->
+                             let result_ident = mk_fresh "pred" in
+                             let result_bd =
+                               FStarBinder.of_named_typ expr.span result_ident
+                                 expr.typ
+                             in
+                             let typ = pty expr.span expr.typ in
+                             let expr = U.make_lets !bindings expr in
+                             let expr = pexpr expr in
+                             let result =
+                               F.term
+                               @@ F.AST.Var
+                                    (plocal_ident result_ident |> F.lid_of_id)
+                             in
+                             let result =
+                               F.AST.Refine
+                                 ( FStarBinder.to_binder result_bd,
+                                   F.implies result expr )
+                               |> F.term
+                             in
+                             F.AST.Product
+                               ( List.map ~f:FStarBinder.to_binder binders,
+                                 result )
+                             |> F.term)
+                      |> Option.value ~default:ty
+                    in
+
                     let ty =
                       F.term
                       @@ F.AST.Product

engine/lib/phases/phase_traits_specs.ml

@@ -28,8 +28,8 @@ module Make (F : Features.T) =
           let v =
             match item.v with
             | Trait { name; generics; items } ->
-                let f (item : trait_item) =
-                  let mk kind =
+                let f attrs (item : trait_item) =
+                  let mk role kind =
                     let ti_ident = mk_name item.ti_ident kind in
                     {
                       item with
@@ -38,23 +38,50 @@ module Make (F : Features.T) =
                         [
                           Attr_payloads.to_attr TraitMethodNoPrePost
                             item.ti_span;
-                        ];
+                        ]
+                        @ (List.filter
+                             ~f:
+                               [%matches?
+                                 Types.AssociatedItem { role = role'; _ }, _ when 
+                                 [%eq: Types.ha_assoc_role] role role' ]
+                             attrs
+                          |> List.map ~f:(uncurry Attr_payloads.to_attr));
                     }
                   in
                   match item.ti_v with
                   | TIFn (TArrow (inputs, output)) ->
                       [
-                        { (mk "pre") with ti_v = TIFn (TArrow (inputs, TBool)) };
                         {
-                          (mk "post") with
+                          (mk Types.Requires "pre") with
+                          ti_v = TIFn (TArrow (inputs, TBool));
+                        };
+                        {
+                          (mk Types.Ensures "post") with
                           ti_v = TIFn (TArrow (inputs @ [ output ], TBool));
                         };
                       ]
                   | TIFn _ -> [ (* REFINEMENTS FOR CONSTANTS? *) ]
                   | TIType _ -> [ (* TODO REFINEMENTS FOR TYPES *) ]
                 in
                 let items =
-                  List.concat_map ~f:(fun item -> f item @ [ item ]) items
+                  List.concat_map
+                    ~f:(fun item ->
+                      let attrs =
+                        Attr_payloads.payloads item.ti_attrs
+                      in
+                      let ti_attrs =
+attrs
+                        |> List.filter
+                             ~f:
+                               (fst
+                               >> [%matches?
+                                    Types.AssociatedItem
+                                      { role = Ensures | Requires; _ }] >> not)
+|>
+                        List.map ~f:(uncurry Attr_payloads.to_attr)
+                      in
+                      f attrs item @ [ { item with ti_attrs } ])
+                    items
                 in
                 Trait { name; generics; items }
             | Impl { generics; self_ty; of_trait; items; parent_bounds } ->

hax-lib-macros/src/impl_fn_decoration.rs

@@ -0,0 +1,87 @@
+//! This module defines the `ImplFnDecoration` structure and utils
+//! around it.
+
+use crate::prelude::*;
+use crate::utils::*;
+
+/// Supporting structure that holds the data required by the internal
+/// macro `impl_fn_decoration`.
+pub struct ImplFnDecoration {
+    pub kind: FnDecorationKind,
+    pub phi: Expr,
+    pub generics: Generics,
+    pub self_ty: Type,
+}
+
+/// The various strings allowed as decoration kinds.
+pub const DECORATION_KINDS: &[&str] = &["decreases", "ensures", "requires"];
+
+/// Expects a `Path` to be a decoration kind: `::hax_lib::<KIND>`,
+/// `hax_lib::<KIND>` or `<KIND>` in (with `KIND` in
+/// `DECORATION_KINDS`).
+pub fn expects_path_decoration(path: &Path) -> Result<Option<String>> {
+    let path_span = path.span();
+    let path = path
+        .expect_simple_path()
+        .ok_or_else(|| Error::new(path_span, "Expected a simple path, with no `<...>`."))?;
+    Ok(
+        match path
+            .iter()
+            .map(|x| x.as_str())
+            .collect::<Vec<_>>()
+            .as_slice()
+        {
+            [kw] | ["", "hax_lib", kw] | ["hax_lib", kw] if DECORATION_KINDS.contains(kw) => {
+                Some(kw.to_string())
+            }
+            _ => None,
+        },
+    )
+}
+
+impl parse::Parse for ImplFnDecoration {
+    fn parse(input: parse::ParseStream) -> Result<Self> {
+        let parse_next = || -> Result<_> {
+            input.parse::<Token![,]>()?;
+            let mut generics = input.parse::<Generics>()?;
+            input.parse::<Token![,]>()?;
+            generics.where_clause = input.parse::<Option<WhereClause>>()?;
+            input.parse::<Token![,]>()?;
+            let self_ty = input.parse::<Type>()?;
+            input.parse::<Token![,]>()?;
+            Ok((generics, self_ty))
+        };
+
+        let path = input.parse::<Path>()?;
+        let path_span = path.span();
+        let kind = match expects_path_decoration(&path)? {
+            Some(s) => match s.as_str() {
+                "decreases" => FnDecorationKind::Decreases,
+                "requires" => FnDecorationKind::Requires,
+                "ensures" => {
+                    let (generics, self_ty) = parse_next()?;
+                    let ExprClosure1 { arg, body } = input.parse::<ExprClosure1>()?;
+                    input.parse::<syn::parse::Nothing>()?;
+                    return Ok(ImplFnDecoration {
+                        kind: FnDecorationKind::Ensures { ret_binder: arg },
+                        phi: body,
+                        generics,
+                        self_ty,
+                    });
+                }
+                _ => unreachable!(),
+            }
+            None => Err(Error::new(path_span, "Expected `::hax_lib::<KIND>`, `hax_lib::<KIND>` or `<KIND>` with `KIND` in {DECORATION_KINDS:?}"))?,
+        };
+
+        let (generics, self_ty) = parse_next()?;
+        let phi = input.parse::<Expr>()?;
+        input.parse::<syn::parse::Nothing>()?;
+        Ok(ImplFnDecoration {
+            kind,
+            phi,
+            generics,
+            self_ty,
+        })
+    }
+}