Improve cps_id tactics, add add_var_types_cps_id, unify_extracted_cps_id tactics

2 files changed, 38 insertions, 20 deletions

diff --git a/src/Experiments/NewPipeline/RewriterWf1.v b/src/Experiments/NewPipeline/RewriterWf1.v
index 72d62e8fa..a244fa1b4 100644
--- a/src/Experiments/NewPipeline/RewriterWf1.v
+++ b/src/Experiments/NewPipeline/RewriterWf1.v
@@ -19,6 +19,7 @@ Require Import Crypto.Util.Tactics.RewriteHyp.
 Require Import Crypto.Util.Tactics.UniquePose.
 Require Import Crypto.Util.Tactics.Head.
 Require Import Crypto.Util.Tactics.RewriteHyp.
+Require Import Crypto.Util.Tactics.CPSId.
 Require Import Crypto.Util.FMapPositive.Equality.
 Require Import Crypto.Util.MSetPositive.Equality.
 Require Import Crypto.Util.Prod.
@@ -59,6 +60,9 @@ Module Compilers.
           reflexivity.
         Qed.
 
+        Ltac add_var_types_cps_id :=
+          cps_id_with_option (@add_var_types_cps_id _ _ _ _ _).
+
         Lemma unify_extracted_cps_id {pt et T k}
           : @pattern.base.unify_extracted_cps pt et T k = k (@pattern.base.unify_extracted_cps pt et _ id).
         Proof using Type.
@@ -69,6 +73,9 @@ Module Compilers.
                          | break_innermost_match_step
                          | rewrite_hyp !* ].
         Qed.
+
+        Ltac unify_extracted_cps_id :=
+          cps_id_with_option (@unify_extracted_cps_id _ _ _ _).
       End base.
 
       Module type.
@@ -82,6 +89,9 @@ Module Compilers.
           reflexivity.
         Qed.
 
+        Ltac add_var_types_cps_id :=
+          cps_id_with_option (@add_var_types_cps_id _ _ _ _ _).
+
         Lemma unify_extracted_cps_id {pt et T k}
           : @pattern.type.unify_extracted_cps pt et T k = k (@pattern.type.unify_extracted_cps pt et _ id).
         Proof using Type.
@@ -94,6 +104,9 @@ Module Compilers.
                          | rewrite_hyp !* ].
         Qed.
 
+        Ltac unify_extracted_cps_id :=
+          cps_id_with_option (@unify_extracted_cps_id _ _ _ _).
+
         Lemma app_forall_vars_under_forall_vars_relation
               {p k1 k2 F v1 v2 evm}
           : @pattern.type.under_forall_vars_relation p k1 k2 F v1 v2
diff --git a/src/Util/Tactics/CPSId.v b/src/Util/Tactics/CPSId.v
index 076eed95a..92b694396 100644
--- a/src/Util/Tactics/CPSId.v
+++ b/src/Util/Tactics/CPSId.v
@@ -16,26 +16,31 @@ Ltac ensure_complex_continuation allow_option k :=
        end
   end.
 
-Ltac cps_id_step allow_option match_to_cont_lem :=
-  match goal with
-  | [ |- context[?e] ]
-    => let res := match_to_cont_lem e in
-       lazymatch res with
-       | (?k, ?lem)
-         => ensure_complex_continuation allow_option k;
-            (* sometimes reduction mismatches happen, so we cast [lem] to fix them *)
-            (rewrite lem || rewrite (lem : e = _))
-       end
-  | [ H : context[?e] |- _ ]
-    => let res := match_to_cont_lem e in
-       lazymatch res with
-       | (?k, ?lem)
-         => ensure_complex_continuation allow_option k;
-            (* sometimes reduction mismatches happen, so we cast [lem] to fix them *)
-            (rewrite lem in H || rewrite (lem : e = _) in H)
+Ltac cps_id_step allow_option lem :=
+  let lem := open_constr:(lem) in
+  lazymatch type of lem with
+  | ?T -> _ => cps_id_step allow_option open_constr:(lem _)
+  | ?lhs = ?k ?rhs
+    => match goal with
+       | [ |- context[?e] ]
+         => (* take advantage of lazymatch treating evars as holes/wildcards *)
+         lazymatch e with
+         | lhs
+           => let lem := constr:(lem : e = _) in
+              ensure_complex_continuation allow_option k;
+              rewrite lem
+         end
+       | [ H : context[?e] |- _ ]
+         => (* take advantage of lazymatch treating evars as holes/wildcards *)
+         lazymatch e with
+         | lhs
+           => let lem := constr:(lem : e = _) in
+              ensure_complex_continuation allow_option k;
+              rewrite lem in H
+         end
        end
   end.
 
-Ltac cps_id allow_option match_to_cont_lem := repeat cps_id_step allow_option match_to_cont_lem.
-Ltac cps_id_with_option match_to_cont_lem := cps_id true match_to_cont_lem.
-Ltac cps_id_no_option match_to_cont_lem := cps_id false match_to_cont_lem.
+Ltac cps_id allow_option lem := repeat cps_id_step allow_option lem.
+Tactic Notation "cps_id_with_option" uconstr(lem) := cps_id true lem.
+Tactic Notation "cps_id_no_option" uconstr(lem) := cps_id false lem.