Make use of named syntax, do reg assign for fancy

4 files changed, 155 insertions, 27 deletions

diff --git a/src/Reflection/Named/ContextOn.v b/src/Reflection/Named/ContextOn.v
new file mode 100644
index 000000000..d32911283
--- /dev/null
+++ b/src/Reflection/Named/ContextOn.v
@@ -0,0 +1,16 @@
+(** * Transfer [Context] across an injection *)
+Require Import Crypto.Reflection.Named.Syntax.
+
+Section language.
+  Context {base_type_code Name1 Name2 : Type}
+          (f : Name2 -> Name1)
+          (f_inj : forall x y, f x = f y -> x = y)
+          {var : base_type_code -> Type}.
+
+  Definition ContextOn (Ctx : Context Name1 var) : Context Name2 var
+    := {| ContextT := Ctx;
+          lookupb ctx n t := lookupb ctx (f n) t;
+          extendb ctx n t v := extendb ctx (f n) v;
+          removeb ctx n t := removeb ctx (f n) t;
+          empty := empty |}.
+End language.
diff --git a/src/Reflection/Named/DeadCodeElimination.v b/src/Reflection/Named/DeadCodeElimination.v
index f7bbed46b..1b2fa3fc0 100644
--- a/src/Reflection/Named/DeadCodeElimination.v
+++ b/src/Reflection/Named/DeadCodeElimination.v
@@ -1,8 +1,13 @@
 (** * PHOAS → Named Representation of Gallina *)
+Require Import Coq.PArith.BinPos Coq.Lists.List.
 Require Import Crypto.Reflection.Named.Syntax.
 Require Import Crypto.Reflection.Named.Compile.
-Require Import Crypto.Reflection.FilterLive.
+Require Import Crypto.Reflection.Named.RegisterAssign.
+Require Import Crypto.Reflection.Named.EstablishLiveness.
+Require Import Crypto.Reflection.CountLets.
 Require Import Crypto.Reflection.Syntax.
+Require Import Crypto.Util.ListUtil.
+Require Import Crypto.Util.LetIn.
 
 Local Notation eta x := (fst x, snd x).
 
@@ -13,7 +18,8 @@ Section language.
   Context (base_type_code : Type)
           (interp_base_type : base_type_code -> Type)
           (op : flat_type base_type_code -> flat_type base_type_code -> Type)
-          (Name : Type).
+          (Name : Type)
+          {Context : Context Name (fun _ : base_type_code => positive)}.
 
   Local Notation flat_type := (flat_type base_type_code).
   Local Notation type := (type base_type_code).
@@ -24,12 +30,12 @@ Section language.
   Local Notation exprf := (@exprf base_type_code interp_base_type op (fun _ => Name)).
   Local Notation expr := (@expr base_type_code interp_base_type op (fun _ => Name)).
   Local Notation Expr := (@Expr base_type_code interp_base_type op).
-  Local Notation lexprf := (@Syntax.exprf base_type_code interp_base_type op (fun _ => list (option Name))).
-  Local Notation lexpr := (@Syntax.expr base_type_code interp_base_type op (fun _ => list (option Name))).
+  (*Local Notation lexprf := (@Syntax.exprf base_type_code interp_base_type op (fun _ => list (option Name))).
+  Local Notation lexpr := (@Syntax.expr base_type_code interp_base_type op (fun _ => list (option Name))).*)
   Local Notation nexprf := (@Named.exprf base_type_code interp_base_type op Name).
   Local Notation nexpr := (@Named.expr base_type_code interp_base_type op Name).
 
-  Definition get_live_namesf (names : list Name) {t} (e : lexprf t) : list (option Name)
+  (*Definition get_live_namesf (names : list (option Name)) {t} (e : lexprf t) : list (option Name)
     := filter_live_namesf
          base_type_code interp_base_type op
          (option Name) None
@@ -38,8 +44,8 @@ Section language.
                      | _, Some y => Some y
                      | None, None => None
                      end)
-         nil (List.map (@Some _) names) e.
-  Definition get_live_names (names : list Name) {t} (e : lexpr t) : list (option Name)
+         nil names e.
+  Definition get_live_names (names : list (option Name)) {t} (e : lexpr t) : list (option Name)
     := filter_live_names
          base_type_code interp_base_type op
          (option Name) None
@@ -48,24 +54,17 @@ Section language.
                      | _, Some y => Some y
                      | None, None => None
                      end)
-         nil (List.map (@Some _) names) e.
-
-  Definition compile_and_eliminate_dead_codef
-             {t} (e : exprf t) (e' : lexprf t) (ls : list Name)
-    : option (nexprf t)
-    := ocompilef e (get_live_namesf ls e').
-
-  Definition compile_and_eliminate_dead_code
-             {t} (e : expr t) (e' : lexpr t) (ls : list Name)
-    : option (nexpr t)
-    := ocompile e (get_live_names ls e').
+         nil names e.*)
 
   Definition CompileAndEliminateDeadCode
              {t} (e : Expr t) (ls : list Name)
     : option (nexpr t)
-    := compile_and_eliminate_dead_code (e _) (e _) ls.
+    := let e := compile (Name:=positive) (e _) (List.map Pos.of_nat (seq 1 (CountBinders e))) in
+       match e with
+       | Some e => Let_In (insert_dead_names None e ls) (* help vm_compute by factoring this out *)
+                          (fun names => register_reassign Pos.eqb empty empty e names)
+       | None => None
+       end.
 End language.
 
-Global Arguments compile_and_eliminate_dead_codef {_ _ _ _ t} _ _ ls.
-Global Arguments compile_and_eliminate_dead_code {_ _ _ _ t} _ _ ls.
-Global Arguments CompileAndEliminateDeadCode {_ _ _ _ t} _ ls.
+Global Arguments CompileAndEliminateDeadCode {_ _ _ _ _ t} e ls.
diff --git a/src/Reflection/Named/EstablishLiveness.v b/src/Reflection/Named/EstablishLiveness.v
new file mode 100644
index 000000000..2301eb6a1
--- /dev/null
+++ b/src/Reflection/Named/EstablishLiveness.v
@@ -0,0 +1,109 @@
+(** * Compute a list of liveness values for each binding *)
+Require Import Coq.Lists.List.
+Require Import Crypto.Reflection.Syntax.
+Require Import Crypto.Reflection.Named.Syntax.
+Require Import Crypto.Reflection.CountLets.
+Require Import Crypto.Util.ListUtil.
+
+Local Notation eta x := (fst x, snd x).
+
+Local Open Scope ctype_scope.
+Delimit Scope nexpr_scope with nexpr.
+
+Inductive liveness := live | dead.
+Fixpoint merge_liveness (ls1 ls2 : list liveness) :=
+  match ls1, ls2 with
+  | cons x xs, cons y ys
+    => cons match x, y with
+            | live, _
+            | _, live
+              => live
+            | dead, dead
+              => dead
+            end
+            (@merge_liveness xs ys)
+  | nil, ls
+  | ls, nil
+    => ls
+  end.
+
+Section language.
+  Context (base_type_code : Type)
+          (interp_base_type : 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).
+  Local Notation type := (type base_type_code).
+  Let Tbase := @Tbase base_type_code.
+  Local Coercion Tbase : base_type_code >-> Syntax.flat_type.
+  Local Notation interp_type := (interp_type interp_base_type).
+  Local Notation interp_flat_type := (interp_flat_type_gen interp_base_type).
+  Local Notation exprf := (@exprf base_type_code interp_base_type op).
+  Local Notation expr := (@expr base_type_code interp_base_type op).
+
+  Section internal.
+    Context (Name : Type)
+            (OutName : Type)
+            {Context : Context Name (fun _ : base_type_code => list liveness)}.
+
+    Definition compute_livenessf_step
+               (compute_livenessf : forall (ctx : Context) {t} (e : exprf Name t) (prefix : list liveness), list liveness)
+               (ctx : Context)
+               {t} (e : exprf Name t) (prefix : list liveness)
+      : list liveness
+      := match e with
+         | Const _ x => prefix
+         | Var t' name => match lookup ctx t' name with
+                          | Some ls => ls
+                          | _ => nil
+                          end
+         | Op _ _ op args
+           => @compute_livenessf ctx _ args prefix
+         | LetIn tx n ex _ eC
+           => let lx := @compute_livenessf ctx _ ex prefix in
+              let lx := merge_liveness lx (prefix ++ repeat live (count_pairs tx)) in
+              let ctx := extend ctx n (SmartVal _ (fun _ => lx) tx) in
+              @compute_livenessf ctx _ eC (prefix ++ repeat dead (count_pairs tx))
+         | Pair _ ex _ ey
+           => merge_liveness (@compute_livenessf ctx _ ex prefix)
+                             (@compute_livenessf ctx _ ey prefix)
+         end.
+
+    Fixpoint compute_livenessf ctx {t} e prefix
+      := @compute_livenessf_step (@compute_livenessf) ctx t e prefix.
+
+    Fixpoint compute_liveness (ctx : Context)
+             {t} (e : expr Name t) (prefix : list liveness)
+      : list liveness
+      := match e with
+         | Return _ x => compute_livenessf ctx x prefix
+         | Abs src _ n f
+           => let prefix := prefix ++ (live::nil) in
+              let ctx := extendb (t:=src) ctx n prefix in
+              @compute_liveness ctx _ f prefix
+         end.
+
+    Section insert_dead.
+      Context (default_out : option OutName).
+
+      Fixpoint insert_dead_names_gen (ls : list liveness) (lsn : list OutName)
+        : list (option OutName)
+        := match ls with
+           | nil => nil
+           | cons live xs
+             => match lsn with
+                | cons n lsn' => Some n :: @insert_dead_names_gen xs lsn'
+                | nil => default_out :: @insert_dead_names_gen xs nil
+                end
+           | cons dead xs
+             => None :: @insert_dead_names_gen xs lsn
+           end.
+      Definition insert_dead_names {t} (e : expr Name t)
+        := insert_dead_names_gen (compute_liveness empty e nil).
+    End insert_dead.
+  End internal.
+End language.
+
+Global Arguments compute_livenessf {_ _ _ _ _} ctx {t} e prefix.
+Global Arguments compute_liveness {_ _ _ _ _} ctx {t} e prefix.
+Global Arguments insert_dead_names {_ _ _ _ _ _} default_out {t} e lsn.
diff --git a/src/Reflection/Named/RegisterAssign.v b/src/Reflection/Named/RegisterAssign.v
index 05b9609d3..5736d01a3 100644
--- a/src/Reflection/Named/RegisterAssign.v
+++ b/src/Reflection/Named/RegisterAssign.v
@@ -1,5 +1,4 @@
 (** * Reassign registers *)
-Require Import Coq.Strings.String Coq.Classes.RelationClasses.
 Require Import Coq.FSets.FMapPositive Coq.PArith.BinPos.
 Require Import Crypto.Reflection.Syntax.
 Require Import Crypto.Reflection.Named.Syntax.
@@ -30,8 +29,12 @@ Section language.
             {ReverseContext : Context OutName (fun _ : base_type_code => InName)}
             (InName_beq : InName -> InName -> bool).
 
-    Fixpoint register_reassignf (ctxi : InContext) (ctxr : ReverseContext)
-             {t} (e : exprf InName t) (new_names : list (option OutName))
+    Definition register_reassignf_step
+               (register_reassignf : forall (ctxi : InContext) (ctxr : ReverseContext)
+                                            {t} (e : exprf InName t) (new_names : list (option OutName)),
+                   option (exprf OutName t))
+               (ctxi : InContext) (ctxr : ReverseContext)
+               {t} (e : exprf InName t) (new_names : list (option OutName))
       : option (exprf OutName t)
       := match e in Named.exprf _ _ _ _ t return option (exprf _ t) with
          | Const _ x => Some (Const x)
@@ -55,10 +58,9 @@ Section language.
                 => let ctxi := extend ctxi n n' in
                    let ctxr := extend ctxr n' n in
                    option_map (LetIn tx n' x) (@register_reassignf ctxi ctxr _ eC new_names')
-              | None, Some x
+              | _, _
                 => let ctxi := remove ctxi n in
                    @register_reassignf ctxi ctxr _ eC new_names'
-              | _, None => None
               end
          | Pair _ ex _ ey
            => match @register_reassignf ctxi ctxr _ ex nil, @register_reassignf ctxi ctxr _ ey nil with
@@ -67,6 +69,8 @@ Section language.
               | _, _ => None
               end
          end.
+    Fixpoint register_reassignf ctxi ctxr {t} e new_names
+      := @register_reassignf_step (@register_reassignf) ctxi ctxr t e new_names.
 
     Fixpoint register_reassign (ctxi : InContext) (ctxr : ReverseContext)
              {t} (e : expr InName t) (new_names : list (option OutName))