feat: BitVec.sshiftRight' in bv_decide (#5995)

src/Init/Data/BitVec/Bitblast.lean

@@ -1092,8 +1092,8 @@ def sshiftRightRec (x : BitVec w₁) (y : BitVec w₂) (n : Nat) : BitVec w₁ :
 
 @[simp]
 theorem sshiftRightRec_zero_eq (x : BitVec w₁) (y : BitVec w₂) :
-    sshiftRightRec x y 0 = x.sshiftRight' (y &&& 1#w₂) := by
-  simp only [sshiftRightRec, twoPow_zero]
+    sshiftRightRec x y 0 = x.sshiftRight' (y &&& twoPow w₂ 0) := by
+  simp only [sshiftRightRec]
 
 @[simp]
 theorem sshiftRightRec_succ_eq (x : BitVec w₁) (y : BitVec w₂) (n : Nat) :

src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reflect.lean

@@ -65,6 +65,8 @@ where
     mkApp4 (mkConst ``BVExpr.shiftLeft) (toExpr m) (toExpr n) (go lhs) (go rhs)
   | .shiftRight (m := m) (n := n) lhs rhs =>
     mkApp4 (mkConst ``BVExpr.shiftRight) (toExpr m) (toExpr n) (go lhs) (go rhs)
+  | .arithShiftRight (m := m) (n := n) lhs rhs =>
+    mkApp4 (mkConst ``BVExpr.arithShiftRight) (toExpr m) (toExpr n) (go lhs) (go rhs)
 
 instance : ToExpr BVBinPred where
   toExpr x :=

src/Lean/Elab/Tactic/BVDecide/Frontend/BVDecide/Reify.lean

@@ -60,8 +60,8 @@ where
           ``BVUnOp.shiftLeftConst
           ``Std.Tactic.BVDecide.Reflect.BitVec.shiftLeftNat_congr
       else
+        let_expr BitVec _ := β | return none
         shiftReflection
-          β
           distanceExpr
           innerExpr
           .shiftLeft
@@ -78,8 +78,8 @@ where
           ``BVUnOp.shiftRightConst
           ``Std.Tactic.BVDecide.Reflect.BitVec.shiftRightNat_congr
       else
+        let_expr BitVec _ := β | return none
         shiftReflection
-          β
           distanceExpr
           innerExpr
           .shiftRight
@@ -92,6 +92,13 @@ where
         innerExpr
         .arithShiftRightConst
         ``BVUnOp.arithShiftRightConst
+        ``Std.Tactic.BVDecide.Reflect.BitVec.arithShiftRightNat_congr
+    | BitVec.sshiftRight' _ _ innerExpr distanceExpr =>
+      shiftReflection
+        distanceExpr
+        innerExpr
+        .arithShiftRight
+        ``BVExpr.arithShiftRight
         ``Std.Tactic.BVDecide.Reflect.BitVec.arithShiftRight_congr
     | BitVec.zeroExtend _ newWidthExpr innerExpr =>
       let some newWidth ← getNatValue? newWidthExpr | return none
@@ -258,11 +265,10 @@ where
     let some distance ← ReifiedBVExpr.getNatOrBvValue? β distanceExpr | return none
     shiftConstLikeReflection distance innerExpr shiftOp shiftOpName congrThm
 
-  shiftReflection (β : Expr) (distanceExpr : Expr) (innerExpr : Expr)
+  shiftReflection (distanceExpr : Expr) (innerExpr : Expr)
       (shiftOp : {m n : Nat} → BVExpr m → BVExpr n → BVExpr m) (shiftOpName : Name)
       (congrThm : Name) :
       LemmaM (Option ReifiedBVExpr) := do
-    let_expr BitVec _ ← β | return none
     let some inner ← goOrAtom innerExpr | return none
     let some distance ← goOrAtom distanceExpr | return none
     let bvExpr : BVExpr inner.width := shiftOp inner.bvExpr distance.bvExpr

src/Std/Tactic/BVDecide/Bitblast/BVExpr/Basic.lean

@@ -245,6 +245,10 @@ inductive BVExpr : Nat → Type where
   shift right by another BitVec expression. For constant shifts there exists a `BVUnop`.
   -/
   | shiftRight (lhs : BVExpr m) (rhs : BVExpr n) : BVExpr m
+  /--
+  shift right arithmetically by another BitVec expression. For constant shifts there exists a `BVUnop`.
+  -/
+  | arithShiftRight (lhs : BVExpr m) (rhs : BVExpr n) : BVExpr m
 
 namespace BVExpr
 
@@ -260,6 +264,7 @@ def toString : BVExpr w → String
   | .signExtend v expr => s!"(sext {v} {expr.toString})"
   | .shiftLeft lhs rhs => s!"({lhs.toString} << {rhs.toString})"
   | .shiftRight lhs rhs => s!"({lhs.toString} >> {rhs.toString})"
+  | .arithShiftRight lhs rhs => s!"({lhs.toString} >>a {rhs.toString})"
 
 
 instance : ToString (BVExpr w) := ⟨toString⟩
@@ -299,6 +304,7 @@ def eval (assign : Assignment) : BVExpr w → BitVec w
   | .signExtend v expr => BitVec.signExtend v (eval assign expr)
   | .shiftLeft lhs rhs => (eval assign lhs) <<< (eval assign rhs)
   | .shiftRight lhs rhs => (eval assign lhs) >>> (eval assign rhs)
+  | .arithShiftRight lhs rhs => BitVec.sshiftRight' (eval assign lhs) (eval assign rhs)
 
 @[simp]
 theorem eval_var : eval assign ((.var idx) : BVExpr w) = (assign.getD idx).bv.truncate _ := by
@@ -343,6 +349,11 @@ theorem eval_shiftLeft : eval assign (.shiftLeft lhs rhs) = (eval assign lhs) <<
 theorem eval_shiftRight : eval assign (.shiftRight lhs rhs) = (eval assign lhs) >>> (eval assign rhs) := by
   rfl
 
+@[simp]
+theorem eval_arithShiftRight :
+    eval assign (.arithShiftRight lhs rhs) = BitVec.sshiftRight' (eval assign lhs) (eval assign rhs) := by
+  rfl
+
 end BVExpr
 
 /--

src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Impl/Expr.lean

@@ -214,6 +214,18 @@ where
         dsimp only at hlaig hraig
         omega
       ⟨res, this⟩
+    | .arithShiftRight lhs rhs =>
+      let ⟨⟨aig, lhs⟩, hlaig⟩ := go aig lhs
+      let ⟨⟨aig, rhs⟩, hraig⟩ := go aig rhs
+      let lhs := lhs.cast <| by
+        dsimp only at hlaig hraig
+        omega
+      let res := bitblast.blastArithShiftRight aig ⟨_, lhs, rhs⟩
+      have := by
+        apply AIG.LawfulVecOperator.le_size_of_le_aig_size (f := bitblast.blastArithShiftRight)
+        dsimp only at hlaig hraig
+        omega
+      ⟨res, this⟩
 
 theorem bitblast.go_decl_eq (aig : AIG BVBit) (expr : BVExpr w) :
     ∀ (idx : Nat) (h1) (h2), (go aig expr).val.aig.decls[idx]'h2 = aig.decls[idx]'h1 := by
@@ -300,6 +312,16 @@ theorem bitblast.go_decl_eq (aig : AIG BVBit) (expr : BVExpr w) :
     · omega
     · apply Nat.lt_of_lt_of_le h1
       apply Nat.le_trans <;> assumption
+  | arithShiftRight lhs rhs lih rih =>
+    dsimp only [go]
+    have := (bitblast.go aig lhs).property
+    have := (bitblast.go aig lhs).property
+    have := (go (go aig lhs).1.aig rhs).property
+    rw [AIG.LawfulVecOperator.decl_eq (f := blastArithShiftRight)]
+    rw [rih, lih]
+    · omega
+    · apply Nat.lt_of_lt_of_le h1
+      apply Nat.le_trans <;> assumption
 
 instance : AIG.LawfulVecOperator BVBit (fun _ w => BVExpr w) bitblast where
   le_size := by

src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Impl/Operations/ShiftRight.lean

@@ -126,14 +126,14 @@ instance : AIG.LawfulVecOperator α AIG.ShiftTarget blastArithShiftRightConst wh
     unfold blastArithShiftRightConst
     simp
 
-namespace blastShiftRight
-
 structure TwoPowShiftTarget (aig : AIG α) (w : Nat) where
   n : Nat
   lhs : AIG.RefVec aig w
   rhs : AIG.RefVec aig n
   pow : Nat
 
+namespace blastShiftRight
+
 def twoPowShift (aig : AIG α) (target : TwoPowShiftTarget aig w) : AIG.RefVecEntry α w :=
   let ⟨n, lhs, rhs, pow⟩ := target
   if h : pow < n then
@@ -246,6 +246,120 @@ instance : AIG.LawfulVecOperator α AIG.ArbitraryShiftTarget blastShiftRight whe
       apply AIG.LawfulVecOperator.lt_size_of_lt_aig_size (f := blastShiftRight.twoPowShift)
       assumption
 
+namespace blastArithShiftRight
+
+def twoPowShift (aig : AIG α) (target : TwoPowShiftTarget aig w) : AIG.RefVecEntry α w :=
+  let ⟨n, lhs, rhs, pow⟩ := target
+  if h : pow < n then
+    let res := blastArithShiftRightConst aig ⟨lhs, 2 ^ pow⟩
+    let aig := res.aig
+    let shifted := res.vec
+
+    have := AIG.LawfulVecOperator.le_size (f := blastArithShiftRightConst) ..
+    let rhs := rhs.cast this
+    let lhs := lhs.cast this
+    AIG.RefVec.ite aig ⟨rhs.get pow h, shifted, lhs⟩
+  else
+    ⟨aig, lhs⟩
+
+instance : AIG.LawfulVecOperator α TwoPowShiftTarget twoPowShift where
+  le_size := by
+    intros
+    unfold twoPowShift
+    dsimp only
+    split
+    · apply AIG.LawfulVecOperator.le_size_of_le_aig_size (f := AIG.RefVec.ite)
+      apply AIG.LawfulVecOperator.le_size (f := blastArithShiftRightConst)
+    · simp
+  decl_eq := by
+    intros
+    unfold twoPowShift
+    dsimp only
+    split
+    · rw [AIG.LawfulVecOperator.decl_eq (f := AIG.RefVec.ite)]
+      rw [AIG.LawfulVecOperator.decl_eq (f := blastArithShiftRightConst)]
+      apply AIG.LawfulVecOperator.lt_size_of_lt_aig_size (f := blastArithShiftRightConst)
+      assumption
+    · simp
+
+end blastArithShiftRight
+
+def blastArithShiftRight (aig : AIG α) (target : AIG.ArbitraryShiftTarget aig w) :
+    AIG.RefVecEntry α w :=
+  let ⟨n, input, distance⟩ := target
+  if n = 0 then
+    ⟨aig, input⟩
+  else
+    let res := blastArithShiftRight.twoPowShift aig ⟨_, input, distance, 0⟩
+    let aig := res.aig
+    let acc := res.vec
+    have := AIG.LawfulVecOperator.le_size (f := blastArithShiftRight.twoPowShift) ..
+    let distance := distance.cast this
+    go aig distance 0 acc
+where
+  go {n : Nat} (aig : AIG α) (distance : AIG.RefVec aig n) (curr : Nat)
+      (acc : AIG.RefVec aig w) :
+      AIG.RefVecEntry α w :=
+    if curr < n - 1 then
+      let res := blastArithShiftRight.twoPowShift aig ⟨_, acc, distance, curr + 1⟩
+      let aig := res.aig
+      let acc := res.vec
+      have := AIG.LawfulVecOperator.le_size (f := blastArithShiftRight.twoPowShift) ..
+      let distance := distance.cast this
+      go aig distance (curr + 1) acc
+    else
+      ⟨aig, acc⟩
+  termination_by n - 1 - curr
+
+theorem blastArithShiftRight.go_le_size (aig : AIG α) (distance : AIG.RefVec aig n) (curr : Nat)
+    (acc : AIG.RefVec aig w) :
+    aig.decls.size ≤ (go aig distance curr acc).aig.decls.size := by
+  unfold go
+  dsimp only
+  split
+  · refine Nat.le_trans ?_ (by apply go_le_size)
+    apply AIG.LawfulVecOperator.le_size (f := blastArithShiftRight.twoPowShift)
+  · simp
+termination_by n - 1 - curr
+
+theorem blastArithShiftRight.go_decl_eq (aig : AIG α) (distance : AIG.RefVec aig n) (curr : Nat)
+    (acc : AIG.RefVec aig w) :
+    ∀ (idx : Nat) (h1) (h2),
+        (go aig distance curr acc).aig.decls[idx]'h2 = aig.decls[idx]'h1 := by
+  generalize hgo : go aig distance curr acc = res
+  unfold go at hgo
+  dsimp only at hgo
+  split at hgo
+  · rw [← hgo]
+    intros
+    rw [blastArithShiftRight.go_decl_eq]
+    rw [AIG.LawfulVecOperator.decl_eq (f := blastArithShiftRight.twoPowShift)]
+    apply AIG.LawfulVecOperator.lt_size_of_lt_aig_size (f := blastArithShiftRight.twoPowShift)
+    assumption
+  · simp [← hgo]
+termination_by n - 1 - curr
+
+
+instance : AIG.LawfulVecOperator α AIG.ArbitraryShiftTarget blastArithShiftRight where
+  le_size := by
+    intros
+    unfold blastArithShiftRight
+    dsimp only
+    split
+    · simp
+    · refine Nat.le_trans ?_ (by apply blastArithShiftRight.go_le_size)
+      apply AIG.LawfulVecOperator.le_size (f := blastArithShiftRight.twoPowShift)
+  decl_eq := by
+    intros
+    unfold blastArithShiftRight
+    dsimp only
+    split
+    · simp
+    · rw [blastArithShiftRight.go_decl_eq]
+      rw [AIG.LawfulVecOperator.decl_eq (f := blastArithShiftRight.twoPowShift)]
+      apply AIG.LawfulVecOperator.lt_size_of_lt_aig_size (f := blastArithShiftRight.twoPowShift)
+      assumption
+
 end bitblast
 end BVExpr
 

src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Lemmas/Expr.lean

@@ -147,6 +147,18 @@ theorem go_denote_eq (aig : AIG BVBit) (expr : BVExpr w) (assign : Assignment) :
       · simp [Ref.hgate]
     · intros
       rw [← rih]
+  | arithShiftRight lhs rhs lih rih =>
+    simp only [go, eval_arithShiftRight]
+    apply denote_blastArithShiftRight
+    · intros
+      dsimp only
+      rw [go_denote_mem_prefix]
+      rw [← lih (aig := aig)]
+      · simp
+      · assumption
+      · simp [Ref.hgate]
+    · intros
+      rw [← rih]
   | bin lhs op rhs lih rih =>
     cases op with
     | and =>

src/Std/Tactic/BVDecide/Bitblast/BVExpr/Circuit/Lemmas/Operations/ShiftRight.lean

@@ -391,6 +391,139 @@ theorem denote_blastShiftRight (aig : AIG α) (target : ArbitraryShiftTarget aig
       · simp [hright]
       · simp [Ref.hgate]
 
+namespace blastArithShiftRight
+
+theorem twoPowShift_eq (aig : AIG α) (target : TwoPowShiftTarget aig w) (lhs : BitVec w)
+    (rhs : BitVec target.n) (assign : α → Bool)
+    (hleft : ∀ (idx : Nat) (hidx : idx < w), ⟦aig, target.lhs.get idx hidx, assign⟧ = lhs.getLsbD idx)
+    (hright : ∀ (idx : Nat) (hidx : idx < target.n), ⟦aig, target.rhs.get idx hidx, assign⟧ = rhs.getLsbD idx) :
+    ∀ (idx : Nat) (hidx : idx < w),
+        ⟦
+          (twoPowShift aig target).aig,
+          (twoPowShift aig target).vec.get idx hidx,
+          assign
+        ⟧
+          =
+        (BitVec.sshiftRight' lhs (rhs &&& BitVec.twoPow target.n target.pow)).getLsbD idx := by
+  intro idx hidx
+  generalize hg : twoPowShift aig target = res
+  rcases target with ⟨n, lvec, rvec, pow⟩
+  simp only [BitVec.and_twoPow]
+  unfold twoPowShift at hg
+  dsimp only at hg
+  split at hg
+  · split
+    · next hif1 =>
+      rw [← hg]
+      simp only [RefVec.denote_ite, RefVec.get_cast, Ref.cast_eq,
+        denote_blastArithShiftRightConst]
+      rw [AIG.LawfulVecOperator.denote_mem_prefix (f := blastArithShiftRightConst)]
+      rw [hright]
+      simp only [hif1, ↓reduceIte]
+      have hmod : 2 ^ pow % 2 ^ n = 2 ^ pow := by
+        apply Nat.mod_eq_of_lt
+        apply Nat.pow_lt_pow_of_lt <;> omega
+      split
+      · next hlt =>
+        rw [hleft]
+        simp [hmod, BitVec.getLsbD_sshiftRight, hlt, hidx]
+      · next hlt =>
+        rw [hleft]
+        simp [BitVec.getLsbD_sshiftRight, hmod, hlt, hidx, BitVec.msb_eq_getLsbD_last]
+    · next hif1 =>
+      simp only [Bool.not_eq_true] at hif1
+      rw [← hg]
+      simp only [RefVec.denote_ite, RefVec.get_cast, Ref.cast_eq,
+        denote_blastArithShiftRightConst]
+      rw [AIG.LawfulVecOperator.denote_mem_prefix (f := blastArithShiftRightConst)]
+      rw [hright]
+      simp only [hif1, Bool.false_eq_true, ↓reduceIte]
+      rw [AIG.LawfulVecOperator.denote_mem_prefix (f := blastArithShiftRightConst)]
+      rw [hleft]
+      simp
+  · have : rhs.getLsbD pow = false := by
+      apply BitVec.getLsbD_ge
+      dsimp only
+      omega
+    simp only [this, Bool.false_eq_true, ↓reduceIte]
+    rw [← hg]
+    rw [hleft]
+    simp
+
+theorem go_denote_eq (aig : AIG α) (distance : AIG.RefVec aig n) (curr : Nat)
+      (hcurr : curr ≤ n - 1) (acc : AIG.RefVec aig w)
+    (lhs : BitVec w) (rhs : BitVec n) (assign : α → Bool)
+    (hacc : ∀ (idx : Nat) (hidx : idx < w), ⟦aig, acc.get idx hidx, assign⟧ = (BitVec.sshiftRightRec lhs rhs curr).getLsbD idx)
+    (hright : ∀ (idx : Nat) (hidx : idx < n), ⟦aig, distance.get idx hidx, assign⟧ = rhs.getLsbD idx) :
+    ∀ (idx : Nat) (hidx : idx < w),
+        ⟦
+          (go aig distance curr acc).aig,
+          (go aig distance curr acc).vec.get idx hidx,
+          assign
+        ⟧
+          =
+        (BitVec.sshiftRightRec lhs rhs (n - 1)).getLsbD idx := by
+  intro idx hidx
+  generalize hgo : go aig distance curr acc = res
+  unfold go at hgo
+  dsimp only at hgo
+  split at hgo
+  · rw [← hgo]
+    rw [go_denote_eq]
+    · omega
+    · intro idx hidx
+      simp only [BitVec.sshiftRightRec_succ_eq]
+      rw [twoPowShift_eq (lhs := BitVec.sshiftRightRec lhs rhs curr)]
+      · simp [hacc]
+      · simp [hright]
+    · intro idx hidx
+      rw [AIG.LawfulVecOperator.denote_mem_prefix (f := twoPowShift)]
+      · simp [hright]
+      · simp [Ref.hgate]
+  · have : curr = n - 1 := by omega
+    rw [← hgo]
+    simp [hacc, this]
+termination_by n - 1 - curr
+
+end blastArithShiftRight
+
+theorem denote_blastArithShiftRight (aig : AIG α) (target : ArbitraryShiftTarget aig w0)
+    (lhs : BitVec w0) (rhs : BitVec target.n) (assign : α → Bool)
+    (hleft : ∀ (idx : Nat) (hidx : idx < w0), ⟦aig, target.target.get idx hidx, assign⟧ = lhs.getLsbD idx)
+    (hright : ∀ (idx : Nat) (hidx : idx < target.n), ⟦aig, target.distance.get idx hidx, assign⟧ = rhs.getLsbD idx) :
+    ∀ (idx : Nat) (hidx : idx < w0),
+        ⟦
+          (blastArithShiftRight aig target).aig,
+          (blastArithShiftRight aig target).vec.get idx hidx,
+          assign
+        ⟧
+          =
+        (BitVec.sshiftRight' lhs rhs).getLsbD idx := by
+  intro idx hidx
+  rw [BitVec.sshiftRight_eq_sshiftRightRec]
+  generalize hres : blastArithShiftRight aig target = res
+  rcases target with ⟨n, target, distance⟩
+  unfold blastArithShiftRight at hres
+  dsimp only at hres
+  split at hres
+  · next hzero =>
+    dsimp only
+    subst hzero
+    rw [← hres]
+    simp [hleft, BitVec.and_twoPow]
+  · rw [← hres]
+    rw [blastArithShiftRight.go_denote_eq]
+    · omega
+    · intro idx hidx
+      simp only [BitVec.sshiftRightRec_zero_eq]
+      rw [blastArithShiftRight.twoPowShift_eq]
+      · simp [hleft]
+      · simp [hright]
+    · intros
+      rw [AIG.LawfulVecOperator.denote_mem_prefix (f := blastArithShiftRight.twoPowShift)]
+      · simp [hright]
+      · simp [Ref.hgate]
+
 end bitblast
 end BVExpr
 

src/Std/Tactic/BVDecide/Reflect.lean

@@ -54,10 +54,15 @@ theorem shiftRight_congr (m n : Nat) (lhs : BitVec m) (rhs : BitVec n) (lhs' : B
     lhs >>> rhs = lhs' >>> rhs' := by
   simp[*]
 
-theorem arithShiftRight_congr (n : Nat) (w : Nat) (x x' : BitVec w) (h : x = x') :
+theorem arithShiftRightNat_congr (n : Nat) (w : Nat) (x x' : BitVec w) (h : x = x') :
     BitVec.sshiftRight x' n = BitVec.sshiftRight x n := by
   simp[*]
 
+theorem arithShiftRight_congr (m n : Nat) (lhs : BitVec m) (rhs : BitVec n) (lhs' : BitVec m)
+    (rhs' : BitVec n) (h1 : lhs' = lhs) (h2 : rhs' = rhs) :
+    BitVec.sshiftRight' lhs rhs = BitVec.sshiftRight' lhs' rhs' := by
+  simp[*]
+
 theorem add_congr (w : Nat) (lhs rhs lhs' rhs' : BitVec w) (h1 : lhs' = lhs) (h2 : rhs' = rhs) :
     lhs' + rhs' = lhs + rhs := by
   simp[*]

tests/lean/run/bv_shift.lean

@@ -52,3 +52,6 @@ theorem shift_unit_7 {x : BitVec 64} : BitVec.sshiftRight x 65 = BitVec.sshiftRi
 
 theorem shift_unit_8 {x : BitVec 64} (h : BitVec.slt 0 x) : BitVec.sshiftRight x 32 = x >>> 32 := by
   bv_decide
+
+theorem shift_unit_9 {x y : BitVec 32} (h : BitVec.slt 0 x) : BitVec.sshiftRight' x y = x >>> y := by
+  bv_decide