Split off a-normal form from flattening

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

1 files changed, 29 insertions, 13 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.