From e6cc5d83e9b96cb1ecf69c0409c6887b8c2ebaea Mon Sep 17 00:00:00 2001
From: Benjamin Gregoire <Benjamin.Gregoire@inria.fr>
Date: Tue, 26 Sep 2023 21:26:29 +0200
Subject: [PATCH] use RSB for x = (y << n) - z;

---
 .../tests/success/arm-m4/sub_using_RSB.jazz   | 33 +++++++++++++
 proofs/compiler/arm_lowering.v                | 10 +++-
 proofs/compiler/arm_lowering_proof.v          | 46 ++++++++++++++-----
 3 files changed, 75 insertions(+), 14 deletions(-)
 create mode 100644 compiler/tests/success/arm-m4/sub_using_RSB.jazz

diff --git a/compiler/tests/success/arm-m4/sub_using_RSB.jazz b/compiler/tests/success/arm-m4/sub_using_RSB.jazz
new file mode 100644
index 000000000..7082a1c48
--- /dev/null
+++ b/compiler/tests/success/arm-m4/sub_using_RSB.jazz
@@ -0,0 +1,33 @@
+u32 MINUS_Q = -8380417;
+
+inline fn __montgomery_reduce_8380417(
+  reg u32 a_low a_high
+) -> reg u32
+{
+  reg u32 a_low_x7 a_low_x7169 a_low_x58728449;
+  reg u32 minus_q;
+  reg u32 res;
+
+    a_low_x7 = (a_low << 3) - a_low;
+    a_low_x7169 = a_low + (a_low_x7 << 10);
+
+  minus_q = MINUS_Q;
+
+  a_low_x58728449 = a_low + (a_low_x7169 << 13);
+
+    // (a_low_x58728449 * minus_q) + (a_low + a_high * 2**32);
+  _, res = #SMLAL(a_low, a_high, a_low_x58728449, minus_q);
+
+  res = res;
+  return res;
+}
+
+
+export fn montgomery_reduce_8380417(reg u32 al ah) -> reg u32
+{
+  reg u32 r;
+
+  r = __montgomery_reduce_8380417(al, ah);
+
+  return r;
+}
\ No newline at end of file
diff --git a/proofs/compiler/arm_lowering.v b/proofs/compiler/arm_lowering.v
index cb9b379cb..6a8e6b0dc 100644
--- a/proofs/compiler/arm_lowering.v
+++ b/proofs/compiler/arm_lowering.v
@@ -130,6 +130,9 @@ Definition get_arg_shift
   else
     None.
 
+Definition is_shift (ws : wsize) (e : pexpr) := 
+  if get_arg_shift ws [::e] is None then true else false. 
+           
 Definition arg_shift
   (mn : arm_mnemonic) (ws : wsize) (e : pexprs) : arm_op * seq pexpr :=
   let '(osh, es) :=
@@ -236,11 +239,14 @@ Definition lower_Papp2_op
       Some (ADD, e0, [:: e1 ])
   | Omul (Op_w _) =>
       Some (MUL, e0, [:: e1 ])
-  | Osub (Op_w _) =>
+  | Osub (Op_w ws) =>
       if is_mul e1 is Some (x, y)
       then Some (MLS, x, [:: y; e0 ])
       else
-      Some (SUB, e0, [:: e1 ])
+      if is_shift ws e1 && ~~ (is_shift ws e0) then 
+        Some (RSB, e1, [:: e0])
+      else
+        Some (SUB, e0, [:: e1 ])
   | Odiv (Cmp_w Signed U32) =>
       Some (SDIV, e0, [:: e1 ])
   | Odiv (Cmp_w Unsigned U32) =>
diff --git a/proofs/compiler/arm_lowering_proof.v b/proofs/compiler/arm_lowering_proof.v
index ef119aaa9..e34e4df8e 100644
--- a/proofs/compiler/arm_lowering_proof.v
+++ b/proofs/compiler/arm_lowering_proof.v
@@ -803,7 +803,7 @@ Proof.
           match goal with
           | [ |- context[ is_mul ] ] => case: is_mulP => [???|]; subst
           end.
-
+      6: case: ifP => _.
       all: move=> [? ? ?] hsemop; subst mn e0' e1'.
       all: discriminate hhas_shift || clear hhas_shift.
 
@@ -818,10 +818,24 @@ Proof.
       all: have {hws0} hws0 := cmp_le_trans hws hws0.
       all: have {hws1} hws1 := cmp_le_trans hws hws1.
 
+      2:{
+        have [ws2 [wbase [wsham [hws2 hbase hsham hw1 [hfvbase hfvsham]]]]] :=
+          get_arg_shiftP hget_arg_shift hfve0 hseme0.
+        have hfves := disj_fvars_read_es3 hfve1 hfvbase hfvsham.
+        split; last done.
+        clear hfve1 hfvbase hfvsham hfves.
+        case/truncate_wordP: hw1 => _ hw1.
+        rewrite /= hseme1 hbase hsham {hseme1 hbase hsham} /=.
+        eexists; first reflexivity.
+        rewrite /exec_sopn /= /sopn_sem /= !truncate_word_le // {hws1 hws2} /=.
+        rewrite (wsub_zero_extend _ _ hws) wsub_wnot1.
+        rewrite !(zero_extend_idem _ hws).
+        by rewrite zero_extend_u hw1.
+      }
       all:
         have [ws2 [wbase [wsham [hws2 hbase hsham hw1 [hfvbase hfvsham]]]]] :=
           get_arg_shiftP hget_arg_shift hfve1 hseme1.
-      all: clear hfve1 hseme1 hget_arg_shift.
+    
       all: have hfves := disj_fvars_read_es3 hfve0 hfvbase hfvsham.
       all: split; last done.
       all: clear hfve0 hfvbase hfvsham hfves.
@@ -887,6 +901,11 @@ Proof.
       match goal with
       | [ |- context[ is_mul ] ] => case: is_mulP => [???|]; subst
       end.
+  all:
+    repeat 
+      match goal with 
+      | [ |- context[ if _ then _ else _] ] => case: ifP => _
+      end.
 
   all: move=> [? ? ?] hsemop; subst mn e0' e1'.
   all: discriminate hhas_shift || clear hhas_shift.
@@ -909,16 +928,19 @@ Proof.
   all: move=> /= ?; subst w.
 
   all: have hfves := disj_fvars_read_es2 hfve0 hfve1.
-  6: have hfves' := disj_fvars_read_es2_app2 hfve1 hfve0.
-  7, 9: have hfves' := disj_fvars_read_es2_app2 hfve0 hfve1.
+  2: have hfves' := disj_fvars_read_es2 hfve1 hfve0.
+  7: have hfves' := disj_fvars_read_es2_app2 hfve1 hfve0.
+  8, 10: have hfves' := disj_fvars_read_es2_app2 hfve0 hfve1.
   all: split; last done.
+
   all: clear hfve0 hfve1 hfves.
 
   all: rewrite /=.
-  6: move: hseme0 => /=; t_xrbindP => ? -> ? -> /= hseme0.
-  7, 8: move: hseme1 => /=; t_xrbindP => ? -> ? -> /= hseme1.
+  7: move: hseme0 => /=; t_xrbindP => ? -> ? -> /= hseme0.
+  8, 9: move: hseme1 => /=; t_xrbindP => ? -> ? -> /= hseme1.
   all: try rewrite hseme0 {hseme0} /=.
   all: try rewrite hseme1 {hseme1} /=.
+
   all: eexists; first reflexivity.
 
   all: rewrite /exec_sopn /=.
@@ -940,12 +962,12 @@ Proof.
   all: rewrite /=.
 
   1: rewrite (wadd_zero_extend _ _ hws).
-  2: rewrite (wsub_zero_extend _ _ hws) wsub_wnot1.
-  3: rewrite -(wand_zero_extend _ _ hws).
-  4: rewrite -(wor_zero_extend _ _ hws).
-  5: rewrite -(wxor_zero_extend _ _ hws).
-  9: rewrite (wmul_zero_extend _ _ hws).
-  1-5, 9: by rewrite !(zero_extend_idem _ hws).
+  2,3: rewrite (wsub_zero_extend _ _ hws) wsub_wnot1.
+  4: rewrite -(wand_zero_extend _ _ hws).
+  5: rewrite -(wor_zero_extend _ _ hws).
+  6: rewrite -(wxor_zero_extend _ _ hws).
+  10: rewrite (wmul_zero_extend _ _ hws).
+  1-6, 10: by rewrite !(zero_extend_idem _ hws).
 
   1: move: hseme0.
   2, 3: move: hseme1.