Merge pull request #546 from hacspec/spec-traits

Spec traits

engine/backends/fstar/fstar_backend.ml

@@ -578,6 +578,10 @@ struct
       List.map ~f:(of_generic_param ~kind span) generics.params
       @ List.mapi ~f:(of_generic_constraint span) generics.constraints
 
+    let of_typ span (nth : int) typ : t =
+      let ident = F.id ("x" ^ Int.to_string nth) in
+      { kind = Explicit; ident; typ = pty span typ }
+
     let to_pattern (x : t) : F.AST.pattern =
       let subpat =
         match x.kind with
@@ -599,6 +603,9 @@ struct
     let to_typ (x : t) : F.AST.term = x.typ
     let to_ident (x : t) : F.Ident.ident = x.ident
 
+    let to_term (x : t) : F.AST.term =
+      F.term @@ F.AST.Var (FStar_Ident.lid_of_ns_and_id [] (to_ident x))
+
     let to_binder (x : t) : F.AST.binder =
       F.AST.
         {
@@ -1074,6 +1081,39 @@ struct
                              [],
                              F.mk_e_app base args ))
                          bounds
+                | TIFn (TArrow (inputs, output))
+                  when Attrs.find_unique_attr i.ti_attrs ~f:(function
+                         | TraitMethodNoPrePost -> Some ()
+                         | _ -> None)
+                       |> Option.is_none ->
+                    let inputs =
+                      List.mapi ~f:(FStarBinder.of_typ e.span) inputs
+                    in
+                    let inputs = generics @ inputs in
+                    let output = pty e.span output in
+                    let ty_pre_post =
+                      let inputs = List.map ~f:FStarBinder.to_term inputs in
+                      let pre =
+                        F.mk_e_app (F.term_of_lid [ name ^ "_pre" ]) inputs
+                      in
+                      let result = F.term_of_lid [ "result" ] in
+                      let post =
+                        F.mk_e_app
+                          (F.term_of_lid [ name ^ "_post" ])
+                          (inputs @ [ result ])
+                      in
+                      let post =
+                        F.mk_e_abs
+                          [ F.pat @@ F.AST.PatVar (F.id "result", None, []) ]
+                          post
+                      in
+                      F.mk_e_app
+                        (F.term_of_lid [ "Prims"; "Pure" ])
+                        [ output; pre; post ]
+                    in
+                    let inputs = List.map ~f:FStarBinder.to_binder inputs in
+                    let ty = F.term @@ F.AST.Product (inputs, ty_pre_post) in
+                    [ (F.id name, None, [], ty) ]
                 | TIFn ty ->
                     let ty = pty e.span ty in
                     let ty =
@@ -1295,6 +1335,7 @@ module TransformToInputLanguage =
   |> Phases.Reject.EarlyExit
   |> Phases.Functionalize_loops
   |> Phases.Reject.As_pattern
+  |> Phases.Traits_specs
   |> SubtypeToInputLanguage
   |> Identity
   ]

engine/lib/ast_utils.ml

@@ -60,6 +60,23 @@ module Make (F : Features.T) = struct
           Some e
       | _ -> None
 
+    let closure (e : expr) : (pat list * expr) option =
+      match e.e with
+      | Closure { params; body; _ } -> Some (params, body)
+      | _ -> None
+
+    let app (e : expr) :
+        (expr * expr list * generic_value list * impl_expr option) option =
+      match e.e with
+      | App { f; args; generic_args; impl } -> Some (f, args, generic_args, impl)
+      | _ -> None
+
+    let pbinding_simple (p : pat) : (local_ident * ty) option =
+      match p.p with
+      | PBinding { mut = Immutable; mode = _; var; typ; subpat = None } ->
+          Some (var, typ)
+      | _ -> None
+
     let concrete_app1 (f : Concrete_ident.name) (e : expr) : expr option =
       match e.e with
       | App
@@ -231,6 +248,23 @@ module Make (F : Features.T) = struct
           | _ -> super#visit_item' () item'
       end
 
+    let replace_local_variables (map : (local_ident, expr, _) Map.t) =
+      object
+        inherit [_] Visitors.map as super
+
+        method! visit_expr () e =
+          match e.e with
+          | LocalVar var -> Map.find map var |> Option.value ~default:e
+          | _ -> super#visit_expr () e
+      end
+
+    (** [replace_local_variable var replacement] returns a visitor
+      that maps any type of the AST replacing every occurence of the
+      expression [LocalVar var] by [replacement]. *)
+    let replace_local_variable (var : local_ident) (replacement : expr) =
+      replace_local_variables
+        (Map.of_alist_exn (module Local_ident) [ (var, replacement) ])
+
     let rename_local_idents (f : local_ident -> local_ident) =
       object
         inherit [_] Visitors.map as _super
@@ -643,7 +677,21 @@ module Make (F : Features.T) = struct
         e
     | _ -> e
 
-  (* let rec remove_empty_tap *)
+  (** See [beta_reduce_closure]'s documentation. *)
+  let beta_reduce_closure_opt (e : expr) : expr option =
+    let* f, args, _, _ = Expect.app e in
+    let* pats, body = Expect.closure f in
+    let* vars = List.map ~f:Expect.pbinding_simple pats |> sequence in
+    let vars = List.map ~f:fst vars in
+    let replacements =
+      List.zip_exn vars args |> Map.of_alist_exn (module Local_ident)
+    in
+    Some ((Mappers.replace_local_variables replacements)#visit_expr () body)
+
+  (** Reduces a [(|x1, ..., xN| body)(e1, ..., eN)] to [body[x1/e1, ..., xN/eN]].
+        This assumes the arities are right: [(|x, y| ...)(e1)]. *)
+  let beta_reduce_closure (e : expr) : expr =
+    beta_reduce_closure_opt e |> Option.value ~default:e
 
   let is_unit_typ : ty -> bool =
     remove_tuple1 >> [%matches? TApp { ident = `TupleType 0; _ }]

engine/lib/attr_payloads.ml

@@ -1,6 +1,7 @@
 open! Prelude
 open Ast
 
+(** Parse [_hax::json] attributes *)
 let payloads : attrs -> (Types.ha_payload * span) list =
   let parse =
     (* we have to parse ["JSON"]: first a string, then a ha_payload *)
@@ -19,6 +20,16 @@ let payloads : attrs -> (Types.ha_payload * span) list =
       | _ -> None)
   >> List.map ~f:(map_fst (Yojson.Safe.from_string >> parse))
 
+(** Create a attribute out of a [payload] *)
+let to_attr (payload : Types.ha_payload) (span : span) : attr =
+  let json =
+    `String (Yojson.Safe.to_string (Types.to_json_ha_payload payload))
+  in
+  let kind : attr_kind =
+    Tool { path = "_hax::json"; tokens = Yojson.Safe.to_string json }
+  in
+  { kind; span }
+
 module UId = struct
   module T = struct
     type t = UId of string [@@deriving show, yojson, compare, sexp, eq]
@@ -135,7 +146,23 @@ module Make (F : Features.T) (Error : Phase_utils.ERROR) = struct
     val associated_expr :
       ?keep_last_args:int -> AssocRole.t -> attrs -> expr option
 
+    val associated_expr_rebinding :
+      pat list -> AssocRole.t -> attrs -> expr option
+    (** Looks up an expression but takes care of rebinding free variables. *)
+
     val associated_refinement_in_type : string list -> attrs -> expr option
+    (** For type, there is a special treatment. The name of fields are
+        global identifiers, and thus are subject to rewriting by
+        [Concrete_ident] at the moment of printing. In contrast, in the
+        refinement `fn` item generated by the proc-macros, the
+        arguments are local identifiers, and thus are rewrited in a
+        different manner.
+
+        Thus, [associated_refinement_in_type] takes a list of
+        [free_variables]: those are already formatted strings as
+        printed by the backend. Then, we rewrite identities in the
+        refinement formula to match exactly this print policy, using
+        *final* local identifiers (see `Local_ident.make_final`). *)
 
     include module type of MakeBase (Error)
   end
@@ -198,30 +225,41 @@ module Make (F : Features.T) (Error : Phase_utils.ERROR) = struct
                Error.assertion_failure span
                  "this associated item was expected to be a `fn` item")
 
+    (** Looks up an associated expression, optionally keeping `keep_last_args` last arguments. If keep_last_args is negative, then all arguments are kept. *)
     let associated_expr ?(keep_last_args = 0) (role : AssocRole.t) :
         attrs -> expr option =
       associated_fn role
       >> Option.map ~f:(fun (_generics, params, body) ->
-             let params =
-               List.drop params (List.length params - keep_last_args)
-               |> List.map ~f:(fun p -> p.pat)
+             let n =
+               if keep_last_args < 0 then 0
+               else List.length params - keep_last_args
              in
+             let params = List.drop params n |> List.map ~f:(fun p -> p.pat) in
              match params with
              | [] -> body
              | _ -> { body with e = Closure { params; body; captures = [] } })
 
-    (* For type, there is a special treatment. The name of fields are
-       global identifiers, and thus are subject to rewriting by
-       [Concrete_ident] at the moment of printing. In contrast, in the
-       refinement `fn` item generated by the proc-macros, the
-       arguments are local identifiers, and thus are rewrited in a
-       different manner.
-
-       Thus, [associated_refinement_in_type] takes a list of
-       [free_variables]: those are already formatted strings as
-       printed by the backend. Then, we rewrite identities in the
-       refinement formula to match exactly this print policy, using
-       *final* local identifiers (see `Local_ident.make_final`). *)
+    let associated_expr_rebinding (params : pat list) (role : AssocRole.t)
+        (attrs : attrs) : expr option =
+      let* _, original_params, body = associated_fn role attrs in
+      let original_params =
+        List.map ~f:(fun param -> param.pat) original_params
+      in
+      let vars_of_pat =
+        U.Reducers.collect_local_idents#visit_pat () >> Set.to_list
+      in
+      let original_vars = List.concat_map ~f:vars_of_pat original_params in
+      let target_vars = List.concat_map ~f:vars_of_pat params in
+      let replacements =
+        List.zip_exn original_vars target_vars
+        |> Map.of_alist_exn (module Local_ident)
+      in
+      Some
+        ((U.Mappers.rename_local_idents (fun v ->
+              Map.find replacements v |> Option.value ~default:v))
+           #visit_expr
+           () body)
+
     let associated_refinement_in_type (free_variables : string list) :
         attrs -> expr option =
       associated_fn Refine

engine/lib/concrete_ident/concrete_ident.ml

@@ -64,15 +64,20 @@ module Imported = struct
         | _ -> path);
     }
 
+  let map_def_path_item_string ~(f : string -> string) :
+      def_path_item -> def_path_item = function
+    | TypeNs s -> TypeNs (f s)
+    | ValueNs s -> ValueNs (f s)
+    | MacroNs s -> MacroNs (f s)
+    | LifetimeNs s -> LifetimeNs (f s)
+    | other -> other
+
+  let map_disambiguated_def_path_item_string ~(f : string -> string)
+      (x : disambiguated_def_path_item) : disambiguated_def_path_item =
+    { x with data = map_def_path_item_string ~f x.data }
+
   let map_path_strings ~(f : string -> string) (did : def_id) : def_id =
-    let f : def_path_item -> def_path_item = function
-      | TypeNs s -> TypeNs (f s)
-      | ValueNs s -> ValueNs (f s)
-      | MacroNs s -> MacroNs (f s)
-      | LifetimeNs s -> LifetimeNs (f s)
-      | other -> other
-    in
-    let f x = { x with data = f x.data } in
+    let f = map_disambiguated_def_path_item_string ~f in
     { did with path = List.map ~f did.path }
 
   module AssociatedItem : sig
@@ -595,6 +600,14 @@ module Create = struct
           path = new_parent.path @ [ List.last_exn old.def_id.path ];
         };
     }
+
+  let map_last ~f old =
+    let last =
+      List.last_exn old.def_id.path
+      |> Imported.map_disambiguated_def_path_item_string ~f
+    in
+    let path = List.drop_last_exn old.def_id.path @ [ last ] in
+    { old with def_id = { old.def_id with path } }
 end
 
 let lookup_raw_impl_info (impl : t) : Types.impl_infos option =

engine/lib/concrete_ident/concrete_ident.mli

@@ -27,6 +27,10 @@ val to_debug_string : t -> string
 module Create : sig
   val fresh_module : from:t list -> t
   val move_under : new_parent:t -> t -> t
+
+  val map_last : f:(string -> string) -> t -> t
+  (** [map_last f ident] applies [f] on the last chunk of [ident]'s
+      path if it holds a string *)
 end
 
 type view = { crate : string; path : string list; definition : string }

engine/lib/phase_utils.ml

@@ -37,10 +37,64 @@ module MakePhaseImplemT (A : Ast.T) (B : Ast.T) = struct
   module type T = sig
     val metadata : Metadata.t
     val ditems : A.item list -> B.item list
-    val dexpr : A.expr -> B.expr
   end
 end
 
+(** Functor that produces module types of monomorphic phases *)
+module MAKE_MONOMORPHIC_PHASE (F : Features.T) = struct
+  module type ARG = sig
+    val phase_id : Diagnostics.Phase.t
+    val ditems : Ast.Make(F).item list -> Ast.Make(F).item list
+  end
+
+  module type T = sig
+    include module type of struct
+      module FB = F
+      module A = Ast.Make (F)
+      module B = Ast.Make (FB)
+      module ImplemT = MakePhaseImplemT (A) (B)
+      module FA = F
+    end
+
+    include ImplemT.T
+  end
+end
+
+(** Make a monomorphic phase: a phase that transform an AST with
+feature set [F] into an AST with the same feature set [F] *)
+module MakeMonomorphicPhase (F : Features.T) (M : MAKE_MONOMORPHIC_PHASE(F).ARG) :
+  MAKE_MONOMORPHIC_PHASE(F).T = struct
+  module FA = F
+  module FB = F
+  module A = Ast.Make (F)
+  module B = Ast.Make (FB)
+  module ImplemT = MakePhaseImplemT (A) (B)
+
+  module Implem = struct
+    let metadata = Metadata.make M.phase_id
+
+    include M
+
+    let subtype (l : A.item list) : B.item list = Stdlib.Obj.magic l
+    let ditems (l : A.item list) : B.item list = ditems l |> subtype
+  end
+
+  include Implem
+end
+
+(** Type of an unconstrainted (forall feature sets) monomorphic phases *)
+module type UNCONSTRAINTED_MONOMORPHIC_PHASE = functor (F : Features.T) -> sig
+  include module type of struct
+    module FB = F
+    module A = Ast.Make (F)
+    module B = Ast.Make (FB)
+    module ImplemT = MakePhaseImplemT (A) (B)
+    module FA = F
+  end
+
+  include ImplemT.T
+end
+
 exception ReportError of Diagnostics.kind
 
 module type ERROR = sig

engine/lib/phases.ml

@@ -10,3 +10,4 @@ module Cf_into_monads = Phase_cf_into_monads.Make
 module Functionalize_loops = Phase_functionalize_loops.Make
 module Reject = Phase_reject
 module Local_mutation = Phase_local_mutation.Make
+module Traits_specs = Phase_traits_specs.Make