From a82aebc804c340e8e57c1dc31cb90c39167fe62d Mon Sep 17 00:00:00 2001
From: Siddharth Bhat <siddu.druid@gmail.com>
Date: Mon, 22 Jul 2024 18:41:41 +0100
Subject: [PATCH] feat: getLsb_replicate

This allows bitblasting `BitVec.replicate`.
---
 src/Init/Data/BitVec/Basic.lean  |  6 ++---
 src/Init/Data/BitVec/Lemmas.lean | 41 ++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/Init/Data/BitVec/Basic.lean b/src/Init/Data/BitVec/Basic.lean
index 9566fe2ec8bc..0e335ad371ae 100644
--- a/src/Init/Data/BitVec/Basic.lean
+++ b/src/Init/Data/BitVec/Basic.lean
@@ -583,11 +583,9 @@ instance : HAppend (BitVec w) (BitVec v) (BitVec (w + v)) := ⟨.append⟩
 -- TODO: write this using multiplication
 /-- `replicate i x` concatenates `i` copies of `x` into a new vector of length `w*i`. -/
 def replicate : (i : Nat) → BitVec w → BitVec (w*i)
-  | 0,   _ => 0
+  | 0,   _ => 0#0
   | n+1, x =>
-    have hEq : w + w*n = w*(n + 1) := by
-      rw [Nat.mul_add, Nat.add_comm, Nat.mul_one]
-    hEq ▸ (x ++ replicate n x)
+    (x ++ replicate n x).cast (by rw [Nat.mul_succ]; omega)
 
 /-!
 ### Cons and Concat
diff --git a/src/Init/Data/BitVec/Lemmas.lean b/src/Init/Data/BitVec/Lemmas.lean
index 936bdd68ac08..ffe3c4373080 100644
--- a/src/Init/Data/BitVec/Lemmas.lean
+++ b/src/Init/Data/BitVec/Lemmas.lean
@@ -1504,4 +1504,45 @@ theorem zeroExtend_truncate_succ_eq_zeroExtend_truncate_or_twoPow_of_getLsb_true
     simp [hx]
   · by_cases hik' : k < i + 1 <;> simp [hik, hik'] <;> omega
 
+@[simp]
+theorem replicate_zero_eq {x : BitVec w} : x.replicate 0 = 0#0 := by
+  simp [replicate]
+
+@[simp]
+theorem replicate_succ_eq {x : BitVec w} :
+    x.replicate (n + 1) =
+    (x ++ replicate n x).cast (by rw [Nat.mul_succ]; omega) := by
+  simp [replicate]
+
+/--
+If a number `w * n ≤ i < w * (n + 1)`, then `i - w * n` equals `i % w`.
+Intuitively, this is true by subtracting `w * n` from the inequality, giving
+`0 ≤ i - w * n < w`, which uniquely identifies `i % w`.
+-/
+private theorem Nat.sub_mul_eq_mod_of_lt_of_le (hlo : w * n ≤ i) (hhi : i < w * (n + 1)) :
+    i - w * n = i % w := by
+  rw [Nat.mul_succ] at hhi
+  rw [Nat.sub_eq_of_eq_add]
+  suffices i / w = n by rw [← this, Nat.add_comm, Nat.div_add_mod]
+  apply Nat.div_eq_of_lt_le
+    (by rw [Nat.mul_comm]; omega)
+    (by rw [Nat.mul_comm, Nat.mul_succ]; omega)
+
+theorem getLsb_replicate {n w : Nat} (x : BitVec w) :
+    (x.replicate n).getLsb i =
+    ((decide (i < w * n)) && (x.getLsb (i % w))) := by
+  induction n generalizing x
+  case zero => simp
+  case succ n ih =>
+    simp only [replicate_succ_eq, getLsb_cast, getLsb_append]
+    by_cases hi : i < w * (n + 1)
+    · simp only [hi, decide_True, Bool.true_and]
+      by_cases hi' : i < w * n
+      · simp [hi', ih]
+      · simp only [hi', decide_False, cond_false]
+        rw [Nat.sub_mul_eq_mod_of_lt_of_le] <;> omega
+    · rw [Nat.mul_succ] at hi ⊢
+      simp only [show ¬i < w * n by omega, decide_False, cond_false, hi, Bool.false_and]
+      apply BitVec.getLsb_ge (x := x) (i := i - w * n) (ge := by omega)
+
 end BitVec