Merge pull request #3466 from FStarLang/gebner_extract_as

Add extract_as attribute.

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

@@ -751,13 +751,24 @@ 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 ->
+    // The specified implementation can be recursive,
+    // to be on the safe side we always mark the replaced sigelt as recursive.
+    {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 +1008,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"; "ExtractAs"; "extract_as"]
 let extract_as_impure_effect_lid = p2l ["FStar"; "Pervasives"; "extract_as_impure_effect"]

src/syntax/FStar.Syntax.Resugar.fst

@@ -1389,8 +1389,7 @@ let mk_decl r q d' =
     d = d' ;
     drange = r ;
     quals = List.choose resugar_qualifier q ;
-    (* TODO : are these stocked up somewhere ? *)
-    attrs = [] ;
+    attrs = [] ; // We fill in the attrs in resugar_sigelt'
     interleaved = false;
   }
 
@@ -1481,7 +1480,7 @@ let resugar_eff_decl' env ed =
   mk_decl r q (A.NewEffect(DefineEffect(eff_name, eff_binders, eff_typ, decls)))
 
 let resugar_sigelt' env se : option A.decl =
-  match se.sigel with
+  let d = (match se.sigel with
   | Sig_bundle {ses} ->
     let decl_typ_ses, datacon_ses = ses |> List.partition
       (fun se -> match se.sigel with
@@ -1601,7 +1600,11 @@ let resugar_sigelt' env se : option A.decl =
     Some (decl'_to_decl se (A.Polymonadic_bind (m, n, p, resugar_term' env t)))
 
   | Sig_polymonadic_subcomp {m_lid=m; n_lid=n; tm=(_, t)} ->
-    Some (decl'_to_decl se (A.Polymonadic_subcomp (m, n, resugar_term' env t)))
+    Some (decl'_to_decl se (A.Polymonadic_subcomp (m, n, resugar_term' env t)))) in
+
+  match d with
+  | Some d -> Some { d with attrs = List.map (resugar_term' env) se.sigattrs }
+  | None -> None
 
 (* Old interface: no envs *)
 

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,20 @@ 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_fvar fv, [t, _] when Syntax.fv_eq_lid fv PC.extract_as_lid ->
+    (match (Subst.compress t).n with
+    | Tm_quoted(impl, _) -> Some impl
+    | _ -> None)
+  | _ -> 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 +1512,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,18 @@
+module ExtractAs
+open FStar.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: nat) -> 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: nat) -> 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.ExtractAs.fst

@@ -0,0 +1,37 @@
+(*
+   Copyright 2024 Microsoft Research
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*)
+
+module FStar.ExtractAs
+open FStar.Stubs.Reflection.Types
+
+(** 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.
+
+    Note that the argument needs to be a literal quotation.
+    The implementation can be recursive,
+    but then you need to construct the attribute via a tactic.
+ *)
+let extract_as (impl: term) = ()

ulib/LowStar.Monotonic.Buffer.fst

@@ -1877,7 +1877,8 @@ let blit #a #rrel1 #rrel2 #rel1 #rel2 src idx_src dst idx_dst len =
       g_upd_seq_as_seq dst (Seq.slice s2' 0 (U32.v length2)) h  //for modifies clause
   | _, _ -> ()
 
-#push-options "--z3rlimit 128 --max_fuel 0 --max_ifuel 1 --initial_ifuel 1 --z3cliopt smt.qi.EAGER_THRESHOLD=4"
+#restart-solver
+#push-options "--z3rlimit 256 --max_fuel 0 --max_ifuel 1 --initial_ifuel 1 --z3cliopt smt.qi.EAGER_THRESHOLD=4"
 let fill' (#t:Type) (#rrel #rel: srel t)
   (b: mbuffer t rrel rel)
   (z:t)