Split off a-normal form from flattening

Now we can flatten let binders without putting operations in a-normal form

7 files changed, 132 insertions, 62 deletions

diff --git a/src/Compilers/Linearize.v b/src/Compilers/Linearize.v
index fc6957bf4..540f705eb 100644
--- a/src/Compilers/Linearize.v
+++ b/src/Compilers/Linearize.v
@@ -5,7 +5,8 @@ Require Import Crypto.Compilers.SmartMap.
 
 Local Open Scope ctype_scope.
 Section language.
-  Context {base_type_code : Type}
+  Context (let_bind_op_args : bool)
+          {base_type_code : Type}
           {op : flat_type base_type_code -> flat_type base_type_code -> Type}.
 
   Local Notation flat_type := (flat_type base_type_code).
@@ -33,31 +34,46 @@ Section language.
            end.
     End under_lets.
 
-    Fixpoint linearizef {t} (e : exprf t) : exprf t
+    Fixpoint linearizef_gen {t} (e : exprf t) : exprf t
       := match e in Syntax.exprf _ _ t return exprf t with
          | LetIn _ ex _ eC
-           => under_letsf (@linearizef _ ex) (fun x => @linearizef _ (eC x))
+           => under_letsf (@linearizef_gen _ ex) (fun x => @linearizef_gen _ (eC x))
          | TT => TT
          | Var _ x => Var x
          | Op _ _ op args
-           => under_letsf (@linearizef _ args) (fun args => LetIn (Op op (SmartVarf args)) SmartVarf)
+           => if let_bind_op_args
+              then under_letsf (@linearizef_gen _ args) (fun args => LetIn (Op op (SmartVarf args)) SmartVarf)
+              else Op op (@linearizef_gen _ args)
          | Pair A ex B ey
-           => under_letsf (@linearizef _ ex) (fun x =>
-              under_letsf (@linearizef _ ey) (fun y =>
+           => under_letsf (@linearizef_gen _ ex) (fun x =>
+              under_letsf (@linearizef_gen _ ey) (fun y =>
               SmartVarf (t:=Prod A B) (x, y)))
          end.
 
-    Definition linearize {t} (e : expr t) : expr t
+    Definition linearize_gen {t} (e : expr t) : expr t
       := match e in Syntax.expr _ _ t return expr t with
-         | Abs _ _ f => Abs (fun x => linearizef (f x))
+         | Abs _ _ f => Abs (fun x => linearizef_gen (f x))
          end.
   End with_var.
 
-  Definition Linearize {t} (e : Expr t) : Expr t
-    := fun var => linearize (e _).
+  Definition Linearize_gen {t} (e : Expr t) : Expr t
+    := fun var => linearize_gen (e _).
 End language.
 
 Global Arguments under_letsf {_ _ _ _} _ {tC} _.
-Global Arguments linearizef {_ _ _ _} _.
-Global Arguments linearize {_ _ _ _} _.
-Global Arguments Linearize {_ _ _} _ var.
+Global Arguments linearizef_gen _ {_ _ _ _} _.
+Global Arguments linearize_gen _ {_ _ _ _} _.
+Global Arguments Linearize_gen _ {_ _ _} _ var.
+
+Definition linearizef {base_type_code op var t} e
+  := @linearizef_gen false base_type_code op var t e.
+Definition a_normalf {base_type_code op var t} e
+  := @linearizef_gen true base_type_code op var t e.
+Definition linearize {base_type_code op var t} e
+  := @linearize_gen false base_type_code op var t e.
+Definition a_normal {base_type_code op var t} e
+  := @linearize_gen true base_type_code op var t e.
+Definition Linearize {base_type_code op t} e
+  := @Linearize_gen false base_type_code op t e.
+Definition ANormal {base_type_code op t} e
+  := @Linearize_gen true base_type_code op t e.
diff --git a/src/Compilers/LinearizeInterp.v b/src/Compilers/LinearizeInterp.v
index a9404c0be..ec8ebf8f1 100644
--- a/src/Compilers/LinearizeInterp.v
+++ b/src/Compilers/LinearizeInterp.v
@@ -55,31 +55,56 @@ Section language.
     induction ex; t_fin.
   Qed.
 
-  Lemma interpf_linearizef {t} e
-    : interpf interp_op (linearizef (t:=t) e) = interpf interp_op e.
-  Proof using Type.
-    clear.
-    induction e;
-      repeat first [ progress rewrite ?interpf_under_letsf, ?interpf_SmartVarf
-                   | progress simpl
-                   | t_fin ].
-  Qed.
+  Section gen.
+    Context (let_bind_op_args : bool).
 
-  Local Hint Resolve interpf_linearizef.
+    Lemma interpf_linearizef_gen {t} e
+      : interpf interp_op (linearizef_gen let_bind_op_args (t:=t) e) = interpf interp_op e.
+    Proof using Type.
+      clear.
+      induction e;
+        repeat first [ progress rewrite ?interpf_under_letsf, ?interpf_SmartVarf
+                     | progress simpl
+                     | t_fin ].
+    Qed.
 
-  Lemma interp_linearize {t} e
-    : forall x, interp interp_op (linearize (t:=t) e) x = interp interp_op e x.
-  Proof using Type.
-    induction e; simpl; eauto.
-  Qed.
+    Local Hint Resolve interpf_linearizef_gen.
 
-  Lemma InterpLinearize {t} (e : Expr t)
-    : forall x, Interp interp_op (Linearize e) x = Interp interp_op e x.
-  Proof using Type.
-    unfold Interp, Linearize.
-    eapply interp_linearize.
-  Qed.
+    Lemma interp_linearize_gen {t} e
+      : forall x, interp interp_op (linearize_gen let_bind_op_args (t:=t) e) x = interp interp_op e x.
+    Proof using Type.
+      induction e; simpl; eauto.
+    Qed.
+
+    Lemma InterpLinearize_gen {t} (e : Expr t)
+      : forall x, Interp interp_op (Linearize_gen let_bind_op_args e) x = Interp interp_op e x.
+    Proof using Type.
+      unfold Interp, Linearize_gen.
+      eapply interp_linearize_gen.
+    Qed.
+  End gen.
+
+  Definition interpf_linearizef {t} e
+    : interpf interp_op (linearizef (t:=t) e) = interpf interp_op e
+    := interpf_linearizef_gen _ e.
+  Definition interpf_a_normalf {t} e
+    : interpf interp_op (a_normalf (t:=t) e) = interpf interp_op e
+    := interpf_linearizef_gen _ e.
+
+  Definition interp_linearize {t} e
+    : forall x, interp interp_op (linearize (t:=t) e) x = interp interp_op e x
+    := interp_linearize_gen _ e.
+  Definition interp_a_normal {t} e
+    : forall x, interp interp_op (a_normal (t:=t) e) x = interp interp_op e x
+    := interp_linearize_gen _ e.
+
+  Definition InterpLinearize {t} (e : Expr t)
+    : forall x, Interp interp_op (Linearize e) x = Interp interp_op e x
+    := InterpLinearize_gen _ e.
+  Definition InterpANormal {t} (e : Expr t)
+    : forall x, Interp interp_op (ANormal e) x = Interp interp_op e x
+    := InterpLinearize_gen _ e.
 End language.
 
 Hint Rewrite @interpf_under_letsf : reflective_interp.
-Hint Rewrite @InterpLinearize @interp_linearize @interpf_linearizef using solve [ eassumption | eauto with wf ] : reflective_interp.
+Hint Rewrite @InterpLinearize_gen @interp_linearize_gen @interpf_linearizef_gen @InterpLinearize @interp_linearize @interpf_linearizef @InterpANormal @interp_a_normal @interpf_a_normalf using solve [ eassumption | eauto with wf ] : reflective_interp.
diff --git a/src/Compilers/LinearizeWf.v b/src/Compilers/LinearizeWf.v
index 073137fd4..46bcc1cfe 100644
--- a/src/Compilers/LinearizeWf.v
+++ b/src/Compilers/LinearizeWf.v
@@ -3,7 +3,8 @@ Require Import Crypto.Compilers.Syntax.
 Require Import Crypto.Compilers.Wf.
 Require Import Crypto.Compilers.WfProofs.
 Require Import Crypto.Compilers.Linearize.
-Require Import (*Crypto.Util.Tactics*) Crypto.Util.Sigma.
+Require Import Crypto.Util.Sigma.
+Require Import Crypto.Util.Tactics.BreakMatch.
 
 Local Open Scope ctype_scope.
 Section language.
@@ -41,6 +42,7 @@ Section language.
       | _ => progress tac
       | [ |- wff _ _ _ ] => constructor
       | [ |- wf _ _ _ ] => constructor
+      | _ => break_innermost_match_step
       end.
     Local Ltac t_fin tac := repeat t_fin_step tac.
 
@@ -150,27 +152,58 @@ Section language.
     Local Hint Resolve wff_in_impl_Proper.
     Local Hint Resolve wff_SmartVarf.
 
-    Lemma wff_linearizef G {t} e1 e2
-      : @wff var1 var2 G t e1 e2
-        -> @wff var1 var2 G t (linearizef e1) (linearizef e2).
-    Proof using Type.
-      induction 1; t_fin ltac:(apply wff_under_letsf).
-    Qed.
+    Section gen.
+      Context (let_bind_op_args : bool).
+
+      Lemma wff_linearizef_gen G {t} e1 e2
+        : @wff var1 var2 G t e1 e2
+          -> @wff var1 var2 G t (linearizef_gen let_bind_op_args e1) (linearizef_gen let_bind_op_args e2).
+      Proof using Type.
+        induction 1; t_fin ltac:(apply wff_under_letsf).
+      Qed.
 
-    Local Hint Resolve wff_linearizef.
+      Local Hint Resolve wff_linearizef_gen.
+
+      Lemma wf_linearize_gen {t} e1 e2
+        : @wf var1 var2 t e1 e2
+          -> @wf var1 var2 t (linearize_gen let_bind_op_args e1) (linearize_gen let_bind_op_args e2).
+      Proof using Type.
+        destruct 1; constructor; auto.
+      Qed.
+    End gen.
+  End with_var.
 
-    Lemma wf_linearize {t} e1 e2
-      : @wf var1 var2 t e1 e2
-        -> @wf var1 var2 t (linearize e1) (linearize e2).
+  Section gen.
+    Context (let_bind_op_args : bool).
+
+    Lemma Wf_Linearize_gen {t} (e : Expr t) : Wf e -> Wf (Linearize_gen let_bind_op_args e).
     Proof using Type.
-      destruct 1; constructor; auto.
+      intros wf var1 var2; apply wf_linearize_gen, wf.
     Qed.
-  End with_var.
+  End gen.
+
+  Definition wff_linearizef {var1 var2} G {t} e1 e2
+    : @wff var1 var2 G t e1 e2
+      -> @wff var1 var2 G t (linearizef e1) (linearizef e2)
+    := wff_linearizef_gen _ G e1 e2.
+  Definition wff_a_normalf {var1 var2} G {t} e1 e2
+    : @wff var1 var2 G t e1 e2
+      -> @wff var1 var2 G t (a_normalf e1) (a_normalf e2)
+    := wff_linearizef_gen _ G e1 e2.
+
+  Definition wf_linearize {var1 var2 t} e1 e2
+    : @wf var1 var2 t e1 e2
+      -> @wf var1 var2 t (linearize e1) (linearize e2)
+    := wf_linearize_gen _ e1 e2.
+  Definition wf_a_normal {var1 var2 t} e1 e2
+    : @wf var1 var2 t e1 e2
+      -> @wf var1 var2 t (a_normal e1) (a_normal e2)
+    := wf_linearize_gen _ e1 e2.
 
-  Lemma Wf_Linearize {t} (e : Expr t) : Wf e -> Wf (Linearize e).
-  Proof using Type.
-    intros wf var1 var2; apply wf_linearize, wf.
-  Qed.
+  Definition Wf_Linearize {t} (e : Expr t) : Wf e -> Wf (Linearize e)
+    := Wf_Linearize_gen _ e.
+  Definition Wf_ANormal {t} (e : Expr t) : Wf e -> Wf (ANormal e)
+    := Wf_Linearize_gen _ e.
 End language.
 
-Hint Resolve Wf_Linearize : wf.
+Hint Resolve Wf_Linearize_gen Wf_Linearize Wf_ANormal : wf.
diff --git a/src/Compilers/TestCase.v b/src/Compilers/TestCase.v
index 5240cf9c5..a7fd81328 100644
--- a/src/Compilers/TestCase.v
+++ b/src/Compilers/TestCase.v
@@ -99,7 +99,7 @@ Import Linearize Inline.
 
 Goal True.
   let x := Reify (fun xy => let '(x, y) := xy in (let a := 1 in let '(c, d) := (2, 3) in a + x - a + c + d) + y)%nat in
-  pose (InlineConst is_const (Linearize x)) as e.
+  pose (InlineConst is_const (ANormal x)) as e.
   vm_compute in e.
 Abort.
 
@@ -190,7 +190,7 @@ Section cse.
   Definition CSE {t} e := @CSE base_type op_code base_type_beq op_code_beq internal_base_type_dec_bl op symbolicify_op (fun _ x => x) t e (fun _ => nil).
 End cse.
 
-Definition example_expr_simplified := Eval vm_compute in InlineConst is_const (Linearize example_expr).
+Definition example_expr_simplified := Eval vm_compute in InlineConst is_const (ANormal example_expr).
 Compute CSE example_expr_simplified.
 
 Definition example_expr_compiled
diff --git a/src/Compilers/Z/Bounds/Pipeline/Definition.v b/src/Compilers/Z/Bounds/Pipeline/Definition.v
index d7f289c70..0d8913234 100644
--- a/src/Compilers/Z/Bounds/Pipeline/Definition.v
+++ b/src/Compilers/Z/Bounds/Pipeline/Definition.v
@@ -73,7 +73,7 @@ Definition PostWfPipeline
   : option (@ProcessedReflectivePackage round_up)
   := Build_ProcessedReflectivePackage_from_option_sigma
        e input_bounds
-       (let e := Linearize e in
+       (let e := ANormal e in
         let e := InlineConst e in
         let e := MapCast _ e input_bounds in
         option_map
@@ -126,7 +126,7 @@ Section with_round_up_list.
     rewrite InterpExprEta_arrow, InterpInlineConst
       by eauto with wf.
     (** Now handle all the transformations that come before the word-size selection *)
-    rewrite <- !InterpLinearize with (e:=e), <- !(@InterpInlineConst _ _ _ (Linearize e))
+    rewrite <- !InterpANormal with (e:=e), <- !(@InterpInlineConst _ _ _ (ANormal e))
       by eauto with wf.
     (** Now handle word-size selection itself *)
     eapply MapCastCorrect; eauto with wf.
diff --git a/src/Compilers/Z/Reify.v b/src/Compilers/Z/Reify.v
index 44ac44a03..c5f31c935 100644
--- a/src/Compilers/Z/Reify.v
+++ b/src/Compilers/Z/Reify.v
@@ -5,10 +5,6 @@ Require Import Crypto.Compilers.Z.Syntax.Equality.
 Require Import Crypto.Compilers.Z.Syntax.Util.
 Require Import Crypto.Compilers.WfReflective.
 Require Import Crypto.Compilers.Reify.
-Require Import Crypto.Compilers.Inline.
-Require Import Crypto.Compilers.InlineInterp.
-Require Import Crypto.Compilers.Linearize.
-Require Import Crypto.Compilers.LinearizeInterp.
 Require Import Crypto.Compilers.Eta.
 Require Import Crypto.Compilers.EtaInterp.
 
diff --git a/src/Specific/FancyMachine256/Core.v b/src/Specific/FancyMachine256/Core.v
index 6fe27b2e4..517644572 100644
--- a/src/Specific/FancyMachine256/Core.v
+++ b/src/Specific/FancyMachine256/Core.v
@@ -157,7 +157,7 @@ Ltac base_reify_type T ::=
 Ltac Reify' e := Reify.Reify' base_type (@interp_base_type _) (@op _) e.
 Ltac Reify e :=
   let v := Reify.Reify base_type (@interp_base_type _) (@op _) (@OPconst _) e in
-  constr:(Inline ((*CSE _*) (InlineConst (@is_const _) (Linearize v)))).
+  constr:(Inline ((*CSE _*) (InlineConst (@is_const _) (ANormal v)))).
 (*Ltac Reify_rhs := Reify.Reify_rhs base_type (interp_base_type _) op (interp_op _).*)
 
 (** ** Raw Syntax Trees *)