Some typeclass performance improvements

src/tactics/FStar.Tactics.V2.Interpreter.fst

@@ -99,8 +99,8 @@ let __primitive_steps_ref : ref (list PO.primitive_step) =
   BU.mk_ref []
 
 let primitive_steps () : list PO.primitive_step =
-    (!__primitive_steps_ref)
-    @ (native_tactics_steps ())
+    (native_tactics_steps ())
+    @ (!__primitive_steps_ref)
 
 let register_tactic_primitive_step (s : PO.primitive_step) : unit =
   __primitive_steps_ref := s :: !__primitive_steps_ref

src/typechecker/FStar.TypeChecker.Tc.fst

@@ -1110,7 +1110,7 @@ let tc_decls env ses =
     let r =
       Profiling.profile
                  (fun () -> process_one_decl acc se)
-                 (Some (Print.sigelt_to_string_short se))
+                 (Some (Ident.string_of_lid (Env.current_module env)))
                  "FStar.TypeChecker.Tc.process_one_decl"
       // ^ See a special case for this phase in FStar.Options. --timing
       // enables it.

ulib/FStar.Tactics.NamedView.fsti

@@ -215,3 +215,23 @@ let inspect_bv       = id #bv
 let pack_bv          = id #bv
 let inspect_comp     = id #comp
 let pack_comp        = id #comp
+
+let tag_of (t:term) : Tac string =
+  match inspect t with
+  | Tv_Var bv -> "Tv_Var"
+  | Tv_BVar fv -> "Tv_BVar"
+  | Tv_FVar fv -> "Tv_FVar"
+  | Tv_UInst _ _ -> "Tv_UInst"
+  | Tv_App f x -> "Tv_App"
+  | Tv_Abs x t -> "Tv_Abs"
+  | Tv_Arrow x t -> "Tv_Arrow"
+  | Tv_Type _ -> "Tv_Type"
+  | Tv_Refine x t -> "Tv_Refine"
+  | Tv_Const cst -> "Tv_Const"
+  | Tv_Uvar i t -> "Tv_Uvar"
+  | Tv_Let r attrs b t1 t2 -> "Tv_Let"
+  | Tv_Match t _ branches -> "Tv_Match"
+  | Tv_AscribedT _ _ _ _ -> "Tv_AscribedT"
+  | Tv_AscribedC _ _ _ _ -> "Tv_AscribedC"
+  | Tv_Unknown -> "Tv_Unknown"
+  | Tv_Unsupp -> "Tv_Unsupp"

ulib/FStar.Tactics.Typeclasses.fst

@@ -24,6 +24,11 @@ open FStar.Tactics.V2.Derived
 open FStar.Tactics.V2.SyntaxCoercions
 open FStar.Tactics.NamedView
 
+(* Thunked version of debug *)
+let debug (f : unit -> Tac string) : Tac unit =
+  if debugging () then
+    print (f ())
+
 module L = FStar.List.Tot.Base
 let (@) = L.op_At
 
@@ -40,51 +45,124 @@ let tcinstance : unit = ()
 irreducible
 let no_method : unit = ()
 
+val fv_eq : fv -> fv -> Tot bool
+let fv_eq fv1 fv2 =
+  let n1 = inspect_fv fv1 in
+  let n2 = inspect_fv fv2 in
+  n1 = n2
+
+let rec head_of (t:term) : Tac (option fv) =
+  (* NB: must use `inspect` to make use of unionfind graph.
+  inspect_ln won't work. *)
+  match inspect t with
+  | Tv_FVar fv
+  | Tv_UInst fv _ -> Some fv
+  | Tv_App h _ -> head_of h
+  | v -> None
+
+let rec res_typ (t:term) : Tac term =
+  match inspect t with
+  | Tv_Arrow _ c -> (
+    match inspect_comp c with
+    | C_Total t -> res_typ t
+    | _ -> t
+  )
+  | _ -> t
+
+(* Would be good to use different exceptions for each reason
+the search stops, but it takes some work to properly account
+for them and report proper errors. *)
+exception NoInst
+
 private
 let rec first (f : 'a -> Tac 'b) (l : list 'a) : Tac 'b =
     match l with
-    | [] -> fail "no cands"
+    | [] -> raise NoInst
     | x::xs -> (fun () -> f x) `or_else` (fun () -> first f xs)
 
-(* TODO: memoization?. And better errors. *)
+(*
+ tcresolve': the main typeclass instantiation function.
+
+ seen : a list of goals we've seen already in this path of the search,
+        used to prevent loops
+ glb : a list of all global instances in scope, for all classes
+ fuel : amount of steps we allow in this path, we stop if we reach zero
+ head_fv : the head of the goal we're trying to solve, i.e. the class name
+
+ TODO: some form of memoization
+*)
 private
-let rec tcresolve' (seen:list term) (fuel:int) : Tac unit =
+let rec tcresolve' (seen : list term) (glb : list fv) (fuel : int) : Tac unit =
     if fuel <= 0 then
-        fail "out of fuel";
-    debug ("fuel = " ^ string_of_int fuel);
+      raise NoInst;
+    debug (fun () -> "fuel = " ^ string_of_int fuel);
+
     let g = cur_goal () in
-    if L.existsb (term_eq g) seen then
-      fail "loop";
-    let seen = g :: seen in
-    local seen fuel `or_else` (fun () -> global seen fuel `or_else` (fun () -> fail ("could not solve constraint: " ^ term_to_string g)))
-and local (seen:list term) (fuel:int) () : Tac unit =
+
+    (* Try to detect loops *)
+    if L.existsb (Reflection.V2.TermEq.term_eq g) seen then
+      raise NoInst;
+
+    match head_of g with
+    | None -> fail ("goal does not look like a typeclass")
+    | Some head_fv ->
+      (* ^ Maybe should check is this really is a class too? *)
+      let seen = g :: seen in
+      local head_fv seen glb fuel
+       `or_else`
+      global head_fv seen glb fuel
+
+and local (head_fv : fv) (seen : list term) (glb : list fv) (fuel : int) () : Tac unit =
     let bs = vars_of_env (cur_env ()) in
-    first (fun (b:binding) -> trywith seen fuel (pack (Tv_Var b))) bs
-and global (seen:list term) (fuel:int) () : Tac unit =
-    let cands = lookup_attr (`tcinstance) (cur_env ()) in
-    first (fun fv -> trywith seen fuel (pack (Tv_FVar fv))) cands
-and trywith (seen:list term) (fuel:int) (t:term) : Tac unit =
-    debug ("Trying to apply hypothesis/instance: " ^ term_to_string t);
-    (fun () -> apply_noinst t) `seq` (fun () -> tcresolve' seen (fuel-1))
+    first (fun (b:binding) ->
+              trywith head_fv seen glb fuel (pack (Tv_Var b)) b.sort)
+          bs
+
+and global (head_fv : fv) (seen : list term) (glb : list fv) (fuel : int) () : Tac unit =
+    first (fun fv ->
+              let typ = tc (cur_env()) (pack (Tv_FVar fv)) in // FIXME: a bit slow.. but at least it's a simple fvar
+              trywith head_fv seen glb fuel (pack (Tv_FVar fv)) typ)
+          glb
+
+and trywith (head_fv : fv) (seen:list term) (glb : list fv) (fuel:int) (t typ : term) : Tac unit =
+    match head_of (res_typ typ) with
+    | None -> raise NoInst
+    | Some fv' ->
+      if fv_eq fv' head_fv
+      then (
+        debug (fun () -> "Trying to apply hypothesis/instance: " ^ term_to_string t);
+        (fun () -> apply_noinst t) `seq` (fun () -> tcresolve' seen glb (fuel-1))
+      ) else
+        raise NoInst
 
 private
 let rec maybe_intros () : Tac unit =
   let g = cur_goal () in
-  match inspect_ln g with
-  | Reflection.V2.Tv_Arrow _ _ ->
+  match inspect g with
+  | Tv_Arrow _ _ ->
     ignore (intro ());
     maybe_intros ()
   | _ -> ()
 
 [@@plugin]
 let tcresolve () : Tac unit =
-    //we sometimes get goal type as _ -> t
-    //so intro if that's the case
-    //not using intros () directly, since that unfolds aggressively if the term is not an arrow
+    // We sometimes get goal type as _ -> t
+    // So intro if that's the case
+    // Not using intros () directly, since that unfolds aggressively if the term is not an arrow
+    // TODO: Should we..?  Why wouldn't the head always be an FV?
     maybe_intros ();
-    try tcresolve' [] 16
+
+    // Fetch a list of all instances in scope right now.
+    // TODO: turn this into a hash map per class, ideally one that can be
+    // stored.
+    let glb = lookup_attr (`tcinstance) (cur_env ()) in
+    try
+      tcresolve' [] glb 16
     with
-    | TacticFailure s -> fail ("Typeclass resolution failed: " ^ s)
+    | NoInst ->
+      fail ("Could not solve constraint " ^ term_to_string (cur_goal ()))
+    | TacticFailure s ->
+      fail ("Typeclass resolution failed: " ^ s)
     | e -> raise e
 
 (**** Generating methods from a class ****)