1 files changed, 88 insertions, 0 deletions
diff --git a/src/Compilers/Named/RegisterAssign.v b/src/Compilers/Named/RegisterAssign.v
new file mode 100644
index 000000000..d1e871fd0
--- /dev/null
+++ b/src/Compilers/Named/RegisterAssign.v
@@ -0,0 +1,88 @@
+(** * Reassign registers *)
+Require Import Crypto.Compilers.Syntax.
+Require Import Crypto.Compilers.Named.PositiveContext.
+Require Import Crypto.Compilers.Named.Syntax.
+Require Import Crypto.Compilers.Named.NameUtil.
+Require Import Crypto.Util.Decidable.
+
+Local Notation eta x := (fst x, snd x).
+
+Local Open Scope ctype_scope.
+Delimit Scope nexpr_scope with nexpr.
+Section language.
+  Context (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 exprf := (@exprf base_type_code op).
+  Local Notation expr := (@expr base_type_code op).
+
+  Section internal.
+    Context (InName OutName : Type)
+            {InContext : Context InName (fun _ : base_type_code => OutName)}
+            {ReverseContext : Context OutName (fun _ : base_type_code => InName)}
+            (InName_beq : InName -> InName -> bool).
+
+    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
+         | TT => Some TT
+         | Var t' name => match lookupb ctxi name t' with
+                          | Some new_name
+                            => match lookupb ctxr new_name t' with
+                               | Some name'
+                                 => if InName_beq name name'
+                                    then Some (Var new_name)
+                                    else None
+                               | None => None
+                               end
+                          | None => None
+                          end
+         | Op _ _ op args
+           => option_map (Op op) (@register_reassignf ctxi ctxr _ args new_names)
+         | LetIn tx n ex _ eC
+           => let '(n', new_names') := eta (split_onames tx new_names) in
+              match n', @register_reassignf ctxi ctxr _ ex nil with
+              | Some n', Some x
+                => 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')
+              | _, _
+                => let ctxi := remove ctxi n in
+                   @register_reassignf ctxi ctxr _ eC new_names'
+              end
+         | Pair _ ex _ ey
+           => match @register_reassignf ctxi ctxr _ ex nil, @register_reassignf ctxi ctxr _ ey nil with
+              | Some x, Some y
+                => Some (Pair x y)
+              | _, _ => None
+              end
+         end.
+    Fixpoint register_reassignf ctxi ctxr {t} e new_names
+      := @register_reassignf_step (@register_reassignf) ctxi ctxr t e new_names.
+
+    Definition register_reassign (ctxi : InContext) (ctxr : ReverseContext)
+             {t} (e : expr InName t) (new_names : list (option OutName))
+      : option (expr OutName t)
+      := match e in Named.expr _ _ _ t return option (expr _ t) with
+         | Abs src _ n f
+           => let '(n', new_names') := eta (split_onames src new_names) in
+              match n' with
+              | Some n'
+                => let ctxi := extend (t:=src) ctxi n n' in
+                   let ctxr := extend (t:=src) ctxr n' n in
+                   option_map (Abs n') (register_reassignf ctxi ctxr f new_names')
+              | None => None
+              end
+         end.
+  End internal.
+End language.
+
+Global Arguments register_reassign {_ _ _ _ _ _} _ ctxi ctxr {t} e _.
+Global Arguments register_reassignf {_ _ _ _ _ _} _ ctxi ctxr {t} e _.