Add some util files for reflective let bindings

4 files changed, 175 insertions, 0 deletions

diff --git a/_CoqProject b/_CoqProject
index a14a022f2..a61d2ea0c 100644
--- a/_CoqProject
+++ b/_CoqProject
@@ -62,6 +62,8 @@ src/ModularArithmetic/Montgomery/Z.v
 src/ModularArithmetic/Montgomery/ZBounded.v
 src/ModularArithmetic/Montgomery/ZProofs.v
 src/Reflection/CommonSubexpressionElimination.v
+src/Reflection/CountLets.v
+src/Reflection/FilterLive.v
 src/Reflection/Inline.v
 src/Reflection/InlineInterp.v
 src/Reflection/InlineWf.v
@@ -76,6 +78,7 @@ src/Reflection/TestCase.v
 src/Reflection/WfProofs.v
 src/Reflection/WfReflective.v
 src/Reflection/WfRel.v
+src/Reflection/Named/NameUtil.v
 src/Spec/CompleteEdwardsCurve.v
 src/Spec/EdDSA.v
 src/Spec/Encoding.v
diff --git a/src/Reflection/CountLets.v b/src/Reflection/CountLets.v
new file mode 100644
index 000000000..43f152da1
--- /dev/null
+++ b/src/Reflection/CountLets.v
@@ -0,0 +1,55 @@
+(** * Counts how many binders there are *)
+Require Import Crypto.Reflection.Syntax.
+
+Local Open Scope ctype_scope.
+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).
+  Local Notation interp_type := (interp_type interp_base_type).
+  Local Notation interp_flat_type := (interp_flat_type_gen interp_base_type).
+  Local Notation Expr := (@Expr base_type_code interp_base_type op).
+
+  Fixpoint count_pairs (t : flat_type) : nat
+    := match t with
+       | Tbase _ => 1
+       | Prod A B => count_pairs A + count_pairs B
+       end%nat.
+
+  Section with_var.
+    Context {var : base_type_code -> Type}
+            (mkVar : forall t, var t).
+
+    Local Notation exprf := (@exprf base_type_code interp_base_type op var).
+    Local Notation expr := (@expr base_type_code interp_base_type op var).
+
+    Section gen.
+      Context (count_type : flat_type -> nat).
+
+      Fixpoint count_lets_genf {t} (e : exprf t) : nat
+        := match e with
+           | LetIn tx _ _ eC
+             => count_type tx + @count_lets_genf _ (eC (SmartVal var mkVar tx))
+           | _ => 0
+           end.
+      Fixpoint count_lets_gen {t} (e : expr t) : nat
+        := match e with
+           | Return _ x
+             => count_lets_genf x
+           | Abs tx _ f => @count_lets_gen _ (f (mkVar tx))
+           end.
+    End gen.
+
+    Definition count_letsf {t} (e : exprf t) : nat
+      := count_lets_genf count_pairs e.
+    Definition count_lets {t} (e : expr t) : nat
+      := count_lets_gen count_pairs e.
+  End with_var.
+
+  Definition CountLetsGen (count_type : flat_type -> nat) {t} (e : Expr t) : nat
+    := count_lets_gen (fun _ => tt) count_type (e _).
+  Definition CountLets {t} (e : Expr t) : nat := CountLetsGen count_pairs e.
+End language.
diff --git a/src/Reflection/FilterLive.v b/src/Reflection/FilterLive.v
new file mode 100644
index 000000000..a16ff19e3
--- /dev/null
+++ b/src/Reflection/FilterLive.v
@@ -0,0 +1,76 @@
+(** * Computes a list of live variables *)
+Require Import Crypto.Reflection.Syntax.
+Require Import Crypto.Reflection.Named.NameUtil.
+Require Import Crypto.Util.ListUtil.
+
+Local Notation eta x := (fst x, snd x).
+
+Local Open Scope ctype_scope.
+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)
+          (dead_name : Name)
+          (merge_names : Name -> Name -> Name)
+          (* equations:
+             - [merge_names x dead_name = merge_names dead_name x = x]
+             - [merge_names x x = x] *).
+
+  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 var := (fun t : base_type_code => list Name).
+  Local Notation exprf := (@exprf base_type_code interp_base_type op var).
+  Local Notation expr := (@expr base_type_code interp_base_type op var).
+  Local Notation Expr := (@Expr base_type_code interp_base_type op var).
+
+  Fixpoint merge_name_lists (ls1 ls2 : list Name) : list Name :=
+    match ls1, ls2 with
+    | cons x xs, cons y ys => cons (merge_names x y) (merge_name_lists xs ys)
+    | ls1, nil => ls1
+    | nil, ls2 => ls2
+    end.
+
+  Fixpoint count_pairs (t : flat_type) : nat
+    := match t with
+       | Syntax.Tbase _ => 1
+       | Prod A B => count_pairs A + count_pairs B
+       end%nat.
+
+  Definition names_to_list {t} : interp_flat_type_gen (fun _ => Name) t -> list Name
+    := smart_interp_flat_map base_type_code (g:=fun _ => list Name) (fun _ x => x :: nil)%list (fun _ _ x y => x ++ y)%list.
+
+  Fixpoint filter_live_namesf (prefix remaining : list Name) {t} (e : exprf t) : list Name
+    := match e with
+       | Const _ x => prefix
+       | Var _ x => x
+       | Op _ _ op args => @filter_live_namesf prefix remaining _ args
+       | LetIn tx ex _ eC
+         => let namesx := @filter_live_namesf prefix nil _ ex in
+            let '(ns, remaining') := eta (split_names tx remaining) in
+            match ns with
+            | Some n =>
+              @filter_live_namesf
+                (prefix ++ repeat dead_name (count_pairs tx))%list remaining' _
+                (eC (SmartVal (fun _ => list Name) (fun _ => namesx ++ names_to_list n)%list _))
+            | None => nil
+            end
+       | Pair _ ex _ ey => merge_name_lists (@filter_live_namesf prefix remaining _ ex)
+                                            (@filter_live_namesf prefix remaining _ ey)
+       end.
+
+  Fixpoint filter_live_names (prefix remaining : list Name) {t} (e : expr t) : list Name
+    := match e, remaining with
+       | Return _ x, _ => filter_live_namesf prefix remaining x
+       | Abs src _ ef, cons n remaining'
+         => let prefix' := (prefix ++ (n :: nil))%list in
+            @filter_live_names
+              prefix' remaining' _
+              (ef prefix')
+       | Abs _ _ _, nil => nil
+       end.
+End language.
diff --git a/src/Reflection/Named/NameUtil.v b/src/Reflection/Named/NameUtil.v
new file mode 100644
index 000000000..bab860585
--- /dev/null
+++ b/src/Reflection/Named/NameUtil.v
@@ -0,0 +1,41 @@
+Require Import Crypto.Reflection.Syntax.
+
+Local Open Scope core_scope.
+Local Notation eta x := (fst x, snd x).
+
+Section language.
+  Context {base_type_code : Type}
+          {Name : Type}.
+
+  Section monad.
+    Context (MName : Type) (force : MName -> option Name).
+    Fixpoint split_mnames
+             (t : flat_type base_type_code) (ls : list MName)
+      : option (interp_flat_type_gen (fun _ => Name) t) * list MName
+      := match t return option (@interp_flat_type_gen base_type_code (fun _ => Name) t) * _ with
+         | Syntax.Tbase _
+           => match ls with
+              | cons n ls'
+                => match force n with
+                   | Some n => (Some n, ls')
+                   | None => (None, ls')
+                   end
+              | nil => (None, nil)
+              end
+         | Prod A B
+           => let '(a, ls) := eta (@split_mnames A ls) in
+              let '(b, ls) := eta (@split_mnames B ls) in
+              (match a, b with
+               | Some a', Some b' => Some (a', b')
+               | _, _ => None
+               end,
+               ls)
+         end.
+  End monad.
+  Definition split_onames := @split_mnames (option Name) (fun x => x).
+  Definition split_names := @split_mnames Name (@Some _).
+End language.
+
+Global Arguments split_mnames {_ _ MName} force _ _, {_ _} MName force _ _.
+Global Arguments split_onames {_ _} _ _.
+Global Arguments split_names {_ _} _ _.