Break up MapCast into separate pieces for easier debugging

3 files changed, 68 insertions, 31 deletions

diff --git a/src/Reflection/MapCastByDeBruijn.v b/src/Reflection/MapCastByDeBruijn.v
index 61af3a1ca..7998b0a37 100644
--- a/src/Reflection/MapCastByDeBruijn.v
+++ b/src/Reflection/MapCastByDeBruijn.v
@@ -6,6 +6,7 @@ Require Import Crypto.Reflection.Named.Compile.
 Require Import Crypto.Reflection.Named.PositiveContext.
 Require Import Crypto.Reflection.Named.PositiveContext.Defaults.
 Require Import Crypto.Reflection.Linearize.
+Require Import Crypto.Reflection.Eta.
 Require Import Crypto.Reflection.Syntax.
 
 Section language.
@@ -21,23 +22,41 @@ Section language.
   Context (cast_op : forall t tR (opc : op t tR) args_bs,
               op (pick_type args_bs) (pick_type (interp_op_bounds t tR opc args_bs))).
 
-  Definition MapCast {t} (e : Expr base_type_code op t)
-             (input_bounds : interp_flat_type interp_base_type_bounds (domain t))
-    : option { output_bounds : interp_flat_type interp_base_type_bounds (codomain t)
-                               & Expr base_type_code op (Arrow (pick_type input_bounds) (pick_type output_bounds)) }
-    := let Context var := PositiveContext _ _ _ base_type_code_bl_transparent in
-       match option_map
-               (fun e' => map_cast
-                            interp_op_bounds pick_typeb cast_op
-                            (BoundsContext:=Context _)
-                            empty
-                            e'
-                            input_bounds)
-               (let e := Linearize e in
-                compile (e _) (DefaultNamesFor e))
-       with
-       | Some (Some (existT output_bounds e'))
-         => Some (existT _ output_bounds (InterpToPHOAS (Context:=Context) failb e'))
-       | Some None | None => None
-       end.
+  Local Notation PContext var := (PositiveContext _ var _ base_type_code_bl_transparent).
+
+  Section MapCast.
+    Context {t} (e : Expr base_type_code op t)
+            (input_bounds : interp_flat_type interp_base_type_bounds (domain t)).
+
+    Definition MapCastPreProcess
+      := ExprEta (Linearize e).
+    Definition MapCastCompile (e' : Expr base_type_code op (domain t -> codomain t))
+      := compile (e' _) (DefaultNamesFor e').
+    Definition MapCastDoCast (e' : option (Named.expr base_type_code op BinNums.positive (domain t -> codomain t)))
+      := option_map
+           (fun e'' => map_cast
+                         (t:=Arrow (domain t) (codomain t))
+                         interp_op_bounds pick_typeb cast_op
+                         (BoundsContext:=PContext _)
+                         empty
+                         e''
+                         input_bounds)
+           e'.
+    Definition MapCastDoInterp
+               (e' : option
+                       (option
+                          { output_bounds : interp_flat_type interp_base_type_bounds (codomain t) &
+                                           Named.expr _ _ _ (Arrow (pick_type input_bounds) (pick_type output_bounds)) }))
+      : option { output_bounds : interp_flat_type interp_base_type_bounds (codomain t)
+                                 & Expr base_type_code op (Arrow (pick_type input_bounds) (pick_type output_bounds)) }
+      := match e' with
+         | Some (Some (existT output_bounds e''))
+           => Some (existT _ output_bounds (Eta.ExprEta (InterpToPHOAS (Context:=fun var => PContext var) failb e'')))
+         | Some None | None => None
+         end.
+    Definition MapCast
+      : option { output_bounds : interp_flat_type interp_base_type_bounds (codomain t)
+                                 & Expr base_type_code op (Arrow (pick_type input_bounds) (pick_type output_bounds)) }
+      := MapCastDoInterp (MapCastDoCast (MapCastCompile MapCastPreProcess)).
+  End MapCast.
 End language.
diff --git a/src/Reflection/MapCastByDeBruijnInterp.v b/src/Reflection/MapCastByDeBruijnInterp.v
index 0e425886b..7fa043791 100644
--- a/src/Reflection/MapCastByDeBruijnInterp.v
+++ b/src/Reflection/MapCastByDeBruijnInterp.v
@@ -11,6 +11,8 @@ Require Import Crypto.Reflection.Named.CompileWf.
 Require Import Crypto.Reflection.Named.PositiveContext.
 Require Import Crypto.Reflection.Named.PositiveContext.Defaults.
 Require Import Crypto.Reflection.Named.PositiveContext.DefaultsProperties.
+Require Import Crypto.Reflection.EtaInterp.
+Require Import Crypto.Reflection.EtaWf.
 Require Import Crypto.Reflection.LinearizeInterp.
 Require Import Crypto.Reflection.LinearizeWf.
 Require Import Crypto.Reflection.Syntax.
@@ -70,6 +72,7 @@ Section language.
       abstract (intro; erewrite base_type_code_lb in pf by eassumption; congruence). }
   Defined.
 
+  Local Arguments Compile.compile : simpl never.
   Lemma MapCastCorrect
         {t} (e : Expr base_type_code op t)
         (Hwf : Wf e)
@@ -80,7 +83,7 @@ Section language.
       /\ @inbounds _ b (Interp interp_op e v)
       /\ cast_back _ _ (Interp interp_op e' v') = (Interp interp_op e v).
   Proof.
-    unfold MapCastByDeBruijn.MapCast, option_map; intros b e'.
+    unfold MapCastByDeBruijn.MapCast, MapCastCompile, MapCastDoCast, MapCastDoInterp, MapCastPreProcess, option_map; intros b e'.
     break_innermost_match; try congruence; intros ? v v'.
     inversion_option; inversion_sigma; subst; simpl in *; intros.
     lazymatch goal with
@@ -93,23 +96,30 @@ Section language.
                 | repeat first [ rewrite !lookupb_empty by (apply PositiveContextOk; assumption)
                                | intro
                                | congruence ] ];
-      unfold Interp, InterpretToPHOAS.Named.InterpToPHOAS, InterpretToPHOAS.Named.InterpToPHOAS_gen;
+      unfold Interp;
       [ match goal with
         | [ H : ?y = Some ?b |- ?x = ?b ]
           => cut (y = Some x); [ congruence | ]
         end
       |
-      | rewrite <- interp_interp_to_phoas; [ reflexivity | ] ].
-    { erewrite (interp_compile (ContextOk:=PositiveContextOk)), interp_linearize;
-        [ reflexivity | apply wf_linearize; auto | .. | eassumption ];
+      | change (interp interp_op (?e ?var) ?v') with (Interp interp_op e v');
+        rewrite InterpExprEta;
+        unfold Interp, InterpretToPHOAS.Named.InterpToPHOAS, InterpretToPHOAS.Named.InterpToPHOAS_gen;
+        rewrite <- interp_interp_to_phoas; [ reflexivity | ] ].
+    { erewrite (interp_compile (ContextOk:=PositiveContextOk)) with (e':=Eta.ExprEta (Linearize.Linearize e) _);
+        cbv [Eta.ExprEta Linearize.Linearize];
+        [ rewrite interp_expr_eta, interp_linearize | .. ];
+        [ reflexivity | eapply wf_expr_eta, wf_linearize; auto | .. | eassumption ];
         auto using name_list_unique_DefaultNamesFor. }
-    { erewrite (interp_compile (ContextOk:=PositiveContextOk)), interp_linearize;
-        [ reflexivity | apply wf_linearize; auto | .. | eassumption ];
+    { erewrite (interp_compile (ContextOk:=PositiveContextOk)) with (e':=Eta.ExprEta (Linearize.Linearize e) _);
+        cbv [Eta.ExprEta Linearize.Linearize];
+        [ rewrite interp_expr_eta, interp_linearize | .. ];
+        [ reflexivity | eapply wf_expr_eta, wf_linearize; auto | .. | eassumption ];
         auto using name_list_unique_DefaultNamesFor. }
     { intro; eapply wf_map_cast with (t := Arrow _ _) (fValues := empty); eauto using PositiveContextOk with typeclass_instances.
-      { eapply (wf_compile (ContextOk:=PositiveContextOk));
-          [ apply wf_linearize; auto | .. | eassumption ].
-        auto using name_list_unique_DefaultNamesFor. }
+      { eapply (wf_compile (ContextOk:=PositiveContextOk)) with (e':=Eta.ExprEta (Linearize.Linearize e) _);
+          [ eapply wf_expr_eta, wf_linearize; auto | .. | eassumption ];
+          auto using name_list_unique_DefaultNamesFor. }
       { intros ???; rewrite lookupb_empty by apply PositiveContextOk; congruence. } }
   Qed.
 End language.
diff --git a/src/Reflection/MapCastByDeBruijnWf.v b/src/Reflection/MapCastByDeBruijnWf.v
index c5fbfa695..4bf5d7c29 100644
--- a/src/Reflection/MapCastByDeBruijnWf.v
+++ b/src/Reflection/MapCastByDeBruijnWf.v
@@ -10,6 +10,7 @@ Require Import Crypto.Reflection.Named.PositiveContext.
 Require Import Crypto.Reflection.Named.PositiveContext.Defaults.
 Require Import Crypto.Reflection.Named.PositiveContext.DefaultsProperties.
 Require Import Crypto.Reflection.LinearizeWf.
+Require Import Crypto.Reflection.EtaWf.
 Require Import Crypto.Reflection.Syntax.
 Require Import Crypto.Reflection.MapCastByDeBruijn.
 Require Import Crypto.Util.Decidable.
@@ -67,6 +68,7 @@ Section language.
       abstract (intro; erewrite base_type_code_lb in pf by eassumption; congruence). }
   Defined.
 
+  Local Arguments Compile.compile : simpl never.
   Lemma Wf_MapCast
         {t} (e : Expr base_type_code op t)
         (Hwf : Wf e)
@@ -74,10 +76,12 @@ Section language.
     : forall {b} e' (He':MapCast e input_bounds = Some (existT _ b e')),
       Wf e'.
   Proof.
-    unfold MapCastByDeBruijn.MapCast, option_map; intros b e'.
+    unfold MapCastByDeBruijn.MapCast, MapCastCompile, MapCastDoCast, MapCastDoInterp, MapCastPreProcess, option_map; intros b e'.
     break_innermost_match; try congruence; intros ? v v'.
     inversion_option; inversion_sigma; subst; simpl in *; intros.
     unfold InterpretToPHOAS.Named.InterpToPHOAS, InterpretToPHOAS.Named.InterpToPHOAS_gen.
+    destruct t as [src dst].
+    apply (proj2 (wf_expr_eta (t:=Arrow _ _))).
     eapply (@wf_interp_to_phoas
               base_type_code op FMapPositive.PositiveMap.key _ _ _ _
               (PositiveContext base_type_code _ base_type_code_beq base_type_code_bl_transparent)
@@ -86,7 +90,11 @@ Section language.
               (failb _) (failb _) _ e1);
       (eapply wf_map_cast with (fValues:=empty); eauto using PositiveContextOk with typeclass_instances);
       try (eapply (wf_compile (ContextOk:=PositiveContextOk));
-           [ apply wf_linearize; eauto | .. | eassumption ]);
+           [ eapply (proj2 (wf_expr_eta (t:=Arrow _ _)));
+             [ eapply wf_linearize | .. ];
+             eauto
+           | ..
+           | eassumption ]);
       try solve [ auto using name_list_unique_DefaultNamesFor
                 | intros ???; rewrite lookupb_empty by apply PositiveContextOk; congruence ].
   Qed.