Merge pull request #790 from hacspec/pre-post-traits

Allow `ensures` and `requires` on traits

engine/backends/fstar/fstar_ast.ml

@@ -117,3 +117,4 @@ let decl_of_string s =
   match decls_of_string s with [ d ] -> d | _ -> failwith "decl_of_string"
 
 let ascribe t e = term @@ AST.Ascribed (e, t, None, false)
+let implies p q = AST.Op (id "==>", [ p; q ]) |> term

engine/backends/fstar/fstar_backend.ml

@@ -688,6 +688,10 @@ struct
       let ident = F.id ("x" ^ Int.to_string nth) in
       { kind = Explicit; ident; typ = pty span typ }
 
+    let of_named_typ span name typ : t =
+      let ident = plocal_ident name in
+      { kind = Explicit; ident; typ = pty span typ }
+
     let to_pattern (x : t) : F.AST.pattern =
       let subpat =
         match x.kind with
@@ -740,87 +744,6 @@ struct
         Error.assertion_failure span "pgeneric_constraint_bd:LIFETIME"
     | GCType { goal; name = _ } -> c_trait_goal span goal
 
-  let get_attr (type a) (name : string) (map : string -> a) (attrs : attrs) :
-      a option =
-    List.find_map
-      ~f:(fun attr ->
-        match attr.kind with
-        | Tool { path; tokens } when [%eq: string] path name ->
-            Some (map tokens)
-        | _ -> None)
-      attrs
-
-  module UUID : sig
-    type t
-
-    val of_attrs : attrs -> t option
-    val associated_items : ?kind:string -> t option -> item list
-    val associated_item : ?kind:string -> t option -> item option
-  end = struct
-    (* TODO: parse_quoted_string is incorrect *)
-    let parse_quoted_string = String.strip ~drop:([%eq: char] '"')
-
-    let parse_associated_with s =
-      let uuid, kind = String.lsplit2 ~on:',' s |> Option.value_exn in
-      let uuid = parse_quoted_string uuid in
-      let kind = String.strip ~drop:([%eq: char] ' ') kind in
-      (uuid, kind)
-
-    type t = string
-
-    let of_attrs : attrs -> t option = get_attr "_hax::uuid" parse_quoted_string
-
-    let associated_items ?kind (uuid : t option) : item list =
-      let ( let* ) x f = Option.bind ~f x in
-      Option.value ~default:[]
-        (let* uuid = uuid in
-         List.filter
-           ~f:(fun item ->
-             Option.value ~default:false
-               (let* uuid', kind' =
-                  get_attr "_hax::associated_with" parse_associated_with
-                    item.attrs
-                in
-                let kind_eq =
-                  match kind with
-                  | Some kind -> String.equal kind' kind
-                  | None -> true
-                in
-                Some (kind_eq && String.equal uuid' uuid)))
-           ctx.items
-         |> Option.some)
-
-    let associated_item ?kind (uuid : t option) : item option =
-      match associated_items ?kind uuid with
-      | [ i ] -> Some i
-      | [] -> None
-      | _ -> failwith "expected 0-1 items"
-  end
-
-  (* TODO: incorrect *)
-  let fvar_of_params = function
-    | { pat = { p = PBinding { var; _ }; _ }; _ } -> var
-    | _ -> failwith "?? todo"
-
-  let associated_refinement (free_variables : string list) (attrs : attrs) :
-      expr option =
-    UUID.associated_item ~kind:"refinement" (UUID.of_attrs attrs)
-    |> Option.map ~f:(function
-         | { v = Fn { params; body; _ }; _ } ->
-             let substs =
-               List.zip_exn
-                 (List.map ~f:fvar_of_params params)
-                 (List.map ~f:Local_ident.make_final free_variables)
-             in
-             let v =
-               U.Mappers.rename_local_idents (fun i ->
-                   match List.find ~f:(fst >> [%eq: local_ident] i) substs with
-                   | None -> i
-                   | Some (_, i) -> i)
-             in
-             v#visit_expr () body
-         | _ -> failwith "expected associated_refinement")
-
   let pmaybe_refined_ty span (free_variables : string list) (attrs : attrs)
       (binder_name : string) (ty : ty) : F.AST.term =
     match Attrs.associated_refinement_in_type free_variables attrs with
@@ -1285,7 +1208,112 @@ 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
+                      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/ast_utils.ml

@@ -378,13 +378,6 @@ module Make (F : Features.T) = struct
         method! visit_local_ident () x = Set.singleton (module Local_ident) x
       end
 
-    let collect_local_idents =
-      object
-        inherit [_] Visitors.reduce as _super
-        inherit [_] Sets.Local_ident.monoid as _m
-        method! visit_local_ident () x = Set.singleton (module Local_ident) x
-      end
-
     include struct
       open struct
         type env = Local_ident.t list
@@ -742,6 +735,11 @@ module Make (F : Features.T) = struct
     if pat_is_expr lhs body then rhs
     else { body with e = Let { monadic = None; lhs; rhs; body } }
 
+  let make_lets (lbs : (pat * expr) list) (body : expr) =
+    List.fold_right ~init:body
+      ~f:(fun (pat, expr) body -> make_let pat expr body)
+      lbs
+
   let make_var_pat (var : local_ident) (typ : ty) (span : span) : pat =
     {
       p = PBinding { mut = Immutable; mode = ByValue; var; typ; subpat = None };

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,48 @@ 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 } ->

engine/lib/print_rust.ml

@@ -450,6 +450,8 @@ module Raw = struct
     let generics = pgeneric_params ti.ti_generics.params in
     let bounds = pgeneric_constraints ti.ti_span ti.ti_generics.constraints in
     let ident = !(Concrete_ident_view.to_definition_name ti.ti_ident) in
+    pattrs ti.ti_attrs
+    &
     match ti.ti_v with
     | TIType _ -> !"type " & ident & !": TodoPrintRustBoundsTyp;"
     | TIFn ty ->
@@ -481,6 +483,8 @@ module Raw = struct
     let generics = pgeneric_params ii.ii_generics.params in
     let bounds = pgeneric_constraints span ii.ii_generics.constraints in
     let ident = !(Concrete_ident_view.to_definition_name ii.ii_ident) in
+    pattrs ii.ii_attrs
+    &
     match ii.ii_v with
     | IIType _ -> !"type " & ident & !": TodoPrintRustBoundsTyp;"
     | IIFn { body; params } ->