Add extract_as attribute.

src/extraction/FStar.Extraction.ML.Modul.fst

@@ -751,13 +751,22 @@ let mark_sigelt_erased (se:sigelt) (g:uenv) : uenv =
   List.fold_right (fun lid g -> extend_erased_fv g (S.lid_as_fv lid None))
                   (U.lids_of_sigelt se) g
 
+// If the definition has an [@@extract_as impl] attribute,
+// replace the lbdef with the specified impl:
+let fixup_sigelt_extract_as se =
+  match se.sigel, find_map se.sigattrs N.is_extract_as_attr with
+  | Sig_let {lids; lbs=(_, [lb])}, Some impl ->
+    {se with sigel = Sig_let {lids; lbs=(true, [{lb with lbdef = impl}])}}
+  | _ -> se
+
 (*  The top-level extraction of a sigelt to an interface *)
 let rec extract_sigelt_iface (g:uenv) (se:sigelt) : uenv & iface =
     if sigelt_has_noextract se then
       let g = mark_sigelt_erased se g in
       g, empty_iface
     else
     let se = karamel_fixup_qual se in
+    let se = fixup_sigelt_extract_as se in
 
     match se.sigel with
     | Sig_bundle _
@@ -997,6 +1006,7 @@ let rec extract_sig (g:env_t) (se:sigelt) : env_t & list mlmodule1 =
     g, []
   else begin
     let se = karamel_fixup_qual se in
+    let se = fixup_sigelt_extract_as se in
 
     match se.sigel with
     | Sig_bundle _

src/parser/FStar.Parser.Const.fst

@@ -598,4 +598,5 @@ let tref_lid        = p2l ["FStar"; "Stubs"; "Tactics"; "Types"; "tref"]
 let document_lid = p2l ["FStar"; "Stubs"; "Pprint"; "document"]
 let issue_lid = p2l ["FStar"; "Issue"; "issue"]
 
+let extract_as_lid = p2l ["FStar"; "Pervasives"; "extract_as"]
 let extract_as_impure_effect_lid = p2l ["FStar"; "Pervasives"; "extract_as_impure_effect"]

src/typechecker/FStar.TypeChecker.Env.fst

@@ -768,15 +768,15 @@ let num_datacon_non_injective_ty_params env lid =
       if injective_type_params then Some 0 else Some num_ty_params
     | _ -> None
 
+let visible_with delta_levels quals =
+  delta_levels |> BU.for_some (fun dl -> quals |> BU.for_some (visible_at dl))
+
 let lookup_definition_qninfo_aux rec_ok delta_levels lid (qninfo : qninfo) =
-  let visible quals =
-      delta_levels |> BU.for_some (fun dl -> quals |> BU.for_some (visible_at dl))
-  in
   match qninfo with
   | Some (Inr (se, None), _) ->
     begin match se.sigel with
       | Sig_let {lbs=(is_rec, lbs)}
-        when visible se.sigquals
+        when visible_with delta_levels se.sigquals
           && (not is_rec || rec_ok) ->
           BU.find_map lbs (fun lb ->
               let fv = right lb.lbname in

src/typechecker/FStar.TypeChecker.Env.fsti

@@ -330,6 +330,7 @@ val lookup_and_inst_datacon: env -> universes -> lident -> typ
 (* the boolean tells if the lident was actually a inductive *)
 val datacons_of_typ        : env -> lident -> (bool & list lident)
 val typ_of_datacon         : env -> lident -> lident
+val visible_with           : list delta_level -> list qualifier -> bool
 val lookup_definition_qninfo : list delta_level -> lident -> qninfo -> option (univ_names & term)
 val lookup_definition      : list delta_level -> env -> lident -> option (univ_names & term)
 val lookup_nonrec_definition: list delta_level -> env -> lident -> option (univ_names & term)

src/typechecker/FStar.TypeChecker.Normalize.fst

@@ -863,6 +863,19 @@ let is_forall_const cfg (phi : term) : option term =
 
     | _ -> None
 
+let is_extract_as_attr (attr: attribute) : option term =
+  let head, args = head_and_args attr in
+  match (Subst.compress head).n, args with
+  | Tm_uinst({n=Tm_fvar fv}, [u]), [_; (t, _)]
+        when Syntax.fv_eq_lid fv PC.extract_as_lid ->
+      Some t
+  | _ -> None
+
+let has_extract_as_attr (g: Env.env) (lid: I.lid) : option term =
+  match Env.lookup_attrs_of_lid g lid with
+  | Some attrs -> find_map attrs is_extract_as_attr
+  | None -> None
+
 (* GM: Please consider this function private outside of this recursive
  * group, and call `normalize` instead. `normalize` will print timing
  * information when --debug NormTop is given, which makes it a
@@ -1498,7 +1511,20 @@ let rec norm : cfg -> env -> stack -> term -> term =
 (* NOTE: we do not need any environment here, since an fv does not
  * have any free indices. Hence, we use empty_env as environment when needed. *)
 and do_unfold_fv (cfg:Cfg.cfg) stack (t0:term) (qninfo : qninfo) (f:fv) : term =
-    match Env.lookup_definition_qninfo cfg.delta_level f.fv_name.v qninfo with
+    // Second, try to unfold to the definition itself.
+    let defn () = Env.lookup_definition_qninfo cfg.delta_level f.fv_name.v qninfo in
+    // First, try to unfold to the implementation specified in the extract_as attribute (when doing extraction)
+    let defn () =
+      if cfg.steps.for_extraction then
+        match qninfo with
+        | Some (Inr (se, None), _) when Env.visible_with cfg.delta_level se.sigquals ->
+          (match find_map se.sigattrs is_extract_as_attr with
+          | Some impl -> Some ([], impl)
+          | None -> defn ())
+        | _ -> defn ()
+      else
+        defn () in
+    match defn () with
        | None ->
          log_unfolding cfg (fun () ->
            BU.print2 " >> No definition found for %s (delta_level = %s)\n"

src/typechecker/FStar.TypeChecker.Normalize.fsti

@@ -66,6 +66,9 @@ val remove_uvar_solutions: Env.env -> term -> term
 val unfold_head_once: Env.env -> term -> option term
 val unembed_binder_knot : ref (option (FStar.Syntax.Embeddings.embedding binder))
 
+val is_extract_as_attr : attribute -> option term
+val has_extract_as_attr : Env.env -> Ident.lid -> option term
+
 val reflection_env_hook : ref (option Env.env)
 
 (* Destructs the term as an arrow type and returns its binders and

tests/extraction/ExtractAs.fst

@@ -0,0 +1,17 @@
+module ExtractAs
+let fail_unless (b: bool) = if b then "ok" else magic ()
+
+// Test that extract_as works when extracting the definition itself.
+
+[@@extract_as (fun x -> x + 10)]
+let frob y = 2 + y
+
+let _ = fail_unless (frob 1 = 11)
+
+// Test that extract_as works when inlining the definition.
+
+inline_for_extraction noextract [@@extract_as (fun x -> x + 10)]
+let bar_2 y = 2 + y
+let bar z = bar_2 z
+
+let _ = fail_unless (bar 1 = 11)

tests/extraction/Makefile

@@ -2,7 +2,7 @@ FSTAR_EXAMPLES=$(realpath ../../examples)
 include $(FSTAR_EXAMPLES)/Makefile.include
 include $(FSTAR_ULIB)/ml/Makefile.include
 
-all: inline_let all_cmi Eta_expand.test Div.test
+all: inline_let all_cmi Eta_expand.test Div.test ExtractAs.test
 
 %.exe: %.fst | out
 	$(FSTAR) $(FSTAR_DEFAULT_ARGS) --odir out --codegen OCaml $<

ulib/FStar.Pervasives.fst

@@ -185,6 +185,8 @@ let primitive_extraction = ()
 
 let extract_as_impure_effect = ()
 
+let extract_as _ = ()
+
 let strictly_positive = ()
 
 let unused = ()

ulib/FStar.Pervasives.fsti

@@ -1146,6 +1146,20 @@ val primitive_extraction : unit
  *)
 val extract_as_impure_effect : unit
 
+(** Replaces the annotated definition
+    by the specified implementation during extraction.
+    There are no checks whether the implementation
+    has the same semantics, or even the same type.
+
+    For example, if you have:
+
+    [@@extract_as (fun (x: nat) -> "not a number")]
+    let add_one (x: nat) : nat = x + 42
+
+    Then `add_one` will extract to `let add_one x = "not a number"`,
+    and most likely cause the extracted program to crash.
+ *)
+val extract_as (#t: Type u#a) (impl: t) : unit
 
 (** A binder in a definition/declaration may optionally be annotated as strictly_positive
     When the let definition is used in a data constructor type in an inductive