Handle IdWithAlt in the simplifier

It now knows how to deal with _ + 0 involving words, and also eliminates
IdWithAlt when both the first argument and the output are the same size word.

4 files changed, 50 insertions, 3 deletions

diff --git a/src/Compilers/Z/ArithmeticSimplifier.v b/src/Compilers/Z/ArithmeticSimplifier.v
index 8da303f4b..21b208f33 100644
--- a/src/Compilers/Z/ArithmeticSimplifier.v
+++ b/src/Compilers/Z/ArithmeticSimplifier.v
@@ -3,6 +3,7 @@ Require Import Coq.ZArith.ZArith.
 Require Import Crypto.Compilers.Syntax.
 Require Import Crypto.Compilers.Rewriter.
 Require Import Crypto.Compilers.Z.Syntax.
+Require Import Crypto.Compilers.Z.Syntax.Equality.
 
 Section language.
   Context (convert_adc_to_sbb : bool).
@@ -84,6 +85,25 @@ Section language.
                    => Op (Sub _ _ _) (Pair ep en)
                  | _ => Op opc args
                  end
+         | Add T1 T2 Tout as opc
+           => fun args
+              => match interp_as_expr_or_const args with
+                 | Some (const_of v, gen_expr e)
+                   => if (v =? 0)%Z
+                      then match base_type_eq_semidec_transparent T2 Tout with
+                           | Some pf => eq_rect _ (fun t => exprf (Tbase t)) e _ pf
+                           | None => Op opc args
+                           end
+                      else Op opc args
+                 | Some (gen_expr e, const_of v)
+                   => if (v =? 0)%Z
+                      then match base_type_eq_semidec_transparent T1 Tout with
+                           | Some pf => eq_rect _ (fun t => exprf (Tbase t)) e _ pf
+                           | None => Op opc args
+                           end
+                      else Op opc args
+                 | _ => Op opc args
+                 end
          | Sub TZ TZ TZ as opc
            => fun args
               => match interp_as_expr_or_const args with
@@ -190,6 +210,23 @@ Section language.
                  | _
                    => Op opc args
                  end
+         | IdWithAlt (TWord _ as T1) _ (TWord _ as Tout) as opc
+           => fun args
+              => match base_type_eq_semidec_transparent T1 Tout with
+                 | Some pf
+                   => match interp_as_expr_or_const args with
+                      | Some (const_of c, _)
+                        => Op (OpConst c) TT
+                      | Some (neg_expr e, _)
+                        => Op (Opp _ _) e
+                      | Some (gen_expr e, _)
+                        => eq_rect _ (fun t => exprf (Tbase t)) e _ pf
+                      | None
+                        => Op opc args
+                      end
+                 | None
+                   => Op opc args
+                 end
          | Zselect TZ TZ TZ TZ as opc
            => fun args
               => match interp_as_expr_or_const args with
@@ -323,7 +360,6 @@ Section language.
                       | _ => Op opc args
                       end
                  else Op opc args
-         | Add _ _ _ as opc
          | Sub _ _ _ as opc
          | Mul _ _ _ as opc
          | Shl _ _ _ as opc
diff --git a/src/Compilers/Z/ArithmeticSimplifierInterp.v b/src/Compilers/Z/ArithmeticSimplifierInterp.v
index 6db34f1b8..06143255b 100644
--- a/src/Compilers/Z/ArithmeticSimplifierInterp.v
+++ b/src/Compilers/Z/ArithmeticSimplifierInterp.v
@@ -151,19 +151,24 @@ Proof.
   apply InterpRewriteOp; intros; unfold simplify_op_expr.
   break_innermost_match;
     repeat first [ fin_t
-                 | progress cbv [LetIn.Let_In Z.zselect]
+                 | progress cbv [LetIn.Let_In Z.zselect IdfunWithAlt.id_with_alt]
                  | progress simpl in *
                  | progress subst
                  | progress subst_prod
                  | rewrite_interp_as_expr_or_const_correct_base ()
                  | rewrite_interp_as_expr_or_const_correct_prod_base ()
                  | rewrite_interp_as_expr_or_const_correct_prod3_base ()
+                 | rewrite FixedWordSizesEquality.ZToWord_wordToZ
+                 | rewrite FixedWordSizesEquality.ZToWord_wordToZ_ZToWord by reflexivity
+                 | rewrite FixedWordSizesEquality.wordToZ_ZToWord_0
                  | progress unfold interp_op, lift_op
                  | progress Z.ltb_to_lt
                  | progress rewrite ?Z.land_0_l, ?Z.land_0_r, ?Z.lor_0_l, ?Z.lor_0_r
                  | rewrite !Z.sub_with_borrow_to_add_get_carry
                  | progress autorewrite with zsimplify_fast
-                 | break_innermost_match_step ].
+                 | break_innermost_match_step
+                 | inversion_base_type_constr_step
+                 | progress cbv [cast_const ZToInterp interpToZ] ].
 Qed.
 
 Hint Rewrite @InterpSimplifyArith : reflective_interp.
diff --git a/src/Compilers/Z/ArithmeticSimplifierWf.v b/src/Compilers/Z/ArithmeticSimplifierWf.v
index 56dda81ab..176301e06 100644
--- a/src/Compilers/Z/ArithmeticSimplifierWf.v
+++ b/src/Compilers/Z/ArithmeticSimplifierWf.v
@@ -241,6 +241,7 @@ Proof.
            | _ => pose_wff_prod3
            | _ => progress destruct_head'_and
            | _ => progress subst
+           | _ => inversion_base_type_constr_step
            | [ H1 : _ = Some _, H2 : _ = None, Hwf : wff _ _ _ |- _ ]
              => pose proof (interp_as_expr_or_const_Some_None Hwf H1 H2); clear H1 H2
            | [ H1 : _ = None, H2 : _ = Some _, Hwf : wff _ _ _ |- _ ]
@@ -251,6 +252,8 @@ Proof.
            | [ |- wff _ (LetIn _ _) (Pair _ _) ] => exfalso
            | [ |- wff _ _ _ ] => constructor; intros
            | [ |- List.In _ _ ] => progress (simpl; auto)
+           | _ => rewrite FixedWordSizesEquality.ZToWord_wordToZ
+           | _ => rewrite FixedWordSizesEquality.ZToWord_wordToZ_ZToWord by reflexivity
            | _ => break_innermost_match_step; simpl in *
            end.
 Qed.
diff --git a/src/Compilers/Z/Bounds/Pipeline/Definition.v b/src/Compilers/Z/Bounds/Pipeline/Definition.v
index 88492074b..ace15b4d7 100644
--- a/src/Compilers/Z/Bounds/Pipeline/Definition.v
+++ b/src/Compilers/Z/Bounds/Pipeline/Definition.v
@@ -114,6 +114,9 @@ Definition PostWfBoundsPipeline
           (projT2_map
              (fun b e'
               => let e' := InlineConst e' in
+                 let e' := InlineConst (SimplifyArith false e') in
+                 let e' := InlineConst (SimplifyArith false e') in
+                 let e' := if opts.(anf) then InlineConst (ANormal e') else e' in
                  let e' := ExprEta e' in
                  e'))
           e).