feat: add BitVec.flattenList and lemmas (#3727)

This PR adds `BitVec.flattenList`, which concatenates a list of
bitvectors of a common width into a single bitvector, together with
lemmas describing its bits: `getLsbD_flattenList` and
`getMsbD_flattenList` compute an individual bit in terms of the
corresponding element of the list, and `extractLsb_flattenList`
describes extracting a contiguous range that falls within a single
element. For efficient execution, `flattenList` is replaced at runtime
via `@[csimp]` with a divide-and-conquer implementation costing `O(n * L
* log L)` rather than the `O(n * L²)` of a naive left fold (≈900x faster
at a million elements), while keeping `O(log L)` recursion depth so it
remains stack-safe.

---------

Co-authored-by: Mario Carneiro <di.gama@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: 3 people

src/Init/Data/BitVec/Basic.lean

@@ -651,6 +651,23 @@ def replicate : (i : Nat) → BitVec w → BitVec (w*i)
   | n+1, x =>
     (x ++ replicate n x).cast (by rw [Nat.mul_succ]; omega)
 
+/-! ### flatten -/
+
+/--
+Flatten a list of bitvectors into one bitvector.
+
+This is `noncomputable` so that the only compiled implementation is the
+divide-and-conquer `BitVec.flattenListFast`, installed via `@[csimp]`; see there
+for the cost analysis.
+-/
+protected noncomputable def flattenList {n : Nat} (xs : List (BitVec n)) : BitVec (n * xs.length) :=
+  match xs with
+  | [] => 0#0
+  | x :: rest =>
+    have h : n + n * List.length rest = n * List.length (x :: rest) := by
+      simp [List.length_cons, Nat.mul_add, Nat.add_comm]
+    BitVec.cast h (x ++ (BitVec.flattenList rest))
+
 /-!
 ### Cons and Concat
 We give special names to the operations of adding a single bit to either end of a bitvector.

src/Init/Data/BitVec/Lemmas.lean

@@ -11,6 +11,8 @@ import all Init.Data.BitVec.BasicAux
 public import Init.Data.Fin.Lemmas
 public import Init.Data.List.BasicAux
 import Init.Data.List.Lemmas
+import Init.Data.List.TakeDrop
+import Init.Data.List.Nat.TakeDrop
 public import Init.Data.BitVec.Basic
 import Init.ByCases
 import Init.Data.BitVec.Bootstrap
@@ -2934,6 +2936,130 @@ theorem setWidth_append {x : BitVec w} {y : BitVec v} :
   simp only [getElem_xor, getElem_append]
   split <;> simp
 
+/-! ## flattenList -/
+
+@[simp] theorem getMsbD_flattenList {w : Nat} (vs : List (BitVec w)) (i) :
+    (BitVec.flattenList vs).getMsbD i = (vs[i / w]?.getD 0#w).getMsbD (i % w) := by
+  by_cases p : 0 < w
+  · match vs with
+    | [] => simp
+    | x :: rest =>
+      simp only [BitVec.flattenList, getMsbD_cast, getMsbD_append]
+      split <;> rename_i h
+      · obtain ⟨j, rfl⟩ := Nat.exists_eq_add_of_le h
+        rw [getMsbD_flattenList]
+        simp [p, Nat.add_div_left, Nat.add_sub_cancel_left]
+      · have hi0 : i / w = 0 := Nat.div_eq_of_lt (by omega)
+        have him : i % w = i := Nat.mod_eq_of_lt (by omega)
+        rw [hi0, him]
+        simp
+  · simp at p
+    subst p
+    simp
+
+@[simp] theorem getLsbD_flattenList {w : Nat} (vs : List (BitVec w)) (i) :
+    (BitVec.flattenList vs).getLsbD i =
+      (decide (i < w * vs.length) && (vs[vs.length - 1 - i / w]?.getD 0#w).getLsbD (i % w)) := by
+  by_cases p : 0 < w
+  · rw [getLsbD_eq_getMsbD, getMsbD_flattenList, getLsbD_eq_getMsbD]
+    have h₁ : i % w < w := Nat.mod_lt i p
+    simp only [h₁, decide_true, Bool.true_and]
+    by_cases h₂ : i < w * vs.length
+    · simp only [h₂, decide_true, Bool.true_and]
+      congr
+      · rw [Nat.sub_sub, Nat.add_comm, Nat.mul_sub_div _ _ _ h₂, Nat.sub_sub, Nat.add_comm]
+      · rw [Nat.sub_sub, Nat.add_comm, Nat.mul_sub_mod h₂, Nat.sub_sub, Nat.add_comm]
+    · simp only [h₂, decide_false, Bool.false_and]
+  · simp_all
+
+@[simp] theorem flattenList_zero_length (vs : List (BitVec 0)) : BitVec.flattenList vs = 0 := by
+  apply eq_of_getLsbD_eq
+  intro i hi
+  simp at hi
+
+/--
+Divide-and-conquer worker for `BitVec.flattenList`, returning the underlying
+natural number value with the head of the list in the most significant bits.
+
+Splitting the list in half means each bit region is shifted `O(log xs.length)`
+times rather than once per element, so flattening a list of length `L` costs
+`O(n * L * log L)` instead of the `O(n * L²)` of a left fold (which reshifts the
+whole growing accumulator at every step). The recursion depth is `O(log L)`, so
+this also does not overflow the stack.
+-/
+def flattenList.toNatAux {n : Nat} : List (BitVec n) → Nat
+  | [] => 0
+  | [x] => x.toNat
+  | x :: y :: rest =>
+    let xs := x :: y :: rest
+    let mid := xs.length / 2
+    flattenList.toNatAux (xs.take mid) <<< (n * (xs.drop mid).length)
+      ||| flattenList.toNatAux (xs.drop mid)
+termination_by l => l.length
+decreasing_by
+  all_goals simp only [List.length_take, List.length_drop, List.length_cons]
+  all_goals omega
+
+/--
+Divide-and-conquer implementation of `BitVec.flattenList`, swapped in at runtime
+via `@[csimp]`. See `BitVec.flattenList.toNatAux` for the cost analysis.
+-/
+def flattenListFast {n : Nat} (xs : List (BitVec n)) : BitVec (n * xs.length) :=
+  .ofNat (n * xs.length) (BitVec.flattenList.toNatAux xs)
+
+theorem toNat_flattenList_append {n : Nat} (l r : List (BitVec n)) :
+    (BitVec.flattenList (l ++ r)).toNat
+      = (BitVec.flattenList l).toNat <<< (n * r.length) ||| (BitVec.flattenList r).toNat := by
+  induction l with
+  | nil => simp [BitVec.flattenList]
+  | cons x l ih =>
+    rw [List.cons_append]
+    simp only [BitVec.flattenList, toNat_cast, toNat_append, ih, List.length_append, Nat.mul_add]
+    rw [Nat.shiftLeft_or_distrib, ← Nat.shiftLeft_add, Nat.or_assoc]
+
+theorem flattenList.toNatAux_eq {n : Nat} (xs : List (BitVec n)) :
+    BitVec.flattenList.toNatAux xs = (BitVec.flattenList xs).toNat := by
+  induction xs using BitVec.flattenList.toNatAux.induct with
+  | case1 => simp [BitVec.flattenList.toNatAux, BitVec.flattenList]
+  | case2 x => simp [BitVec.flattenList.toNatAux, BitVec.flattenList]
+  | case3 x y rest xs mid ih1 ih2 =>
+    rw [BitVec.flattenList.toNatAux, ih1, ih2, ← toNat_flattenList_append, List.take_append_drop]
+
+@[csimp] theorem flattenList_eq_flattenListFast :
+    @BitVec.flattenList = @BitVec.flattenListFast := by
+  funext n xs
+  rw [BitVec.flattenListFast, BitVec.flattenList.toNatAux_eq, ofNat_toNat, setWidth_eq]
+
+/-! ## extractLsb -/
+
+/--
+Extracting a bitvector from `flattenList`, when we only extract from a single bitvector.
+-/
+theorem extractLsb_flattenList (hi lo : Nat) {w  : Nat} (vs : List (BitVec w))
+    (w₁ : lo ≤ hi) (w₂ : hi < w * vs.length) (h : hi / w = lo / w) :
+    extractLsb hi lo (BitVec.flattenList vs) =
+      BitVec.cast (by simp [Nat.mod_eq_sub_mul_div, h]; have := Nat.mul_div_le lo w; omega)
+        (extractLsb (hi % w) (lo % w) (vs[vs.length - 1 - hi / w]?.getD 0#w)) := by
+  by_cases p : 0 < w
+  · have t : w * (lo / w) ≤ lo := Nat.mul_div_le lo w
+    have t' : hi < w * (lo / w) + w := by rw [← h]; apply Nat.lt_mul_div_self_add p
+    apply eq_of_getLsbD_eq
+    intro i hi_lt
+    have q : i < hi - lo + 1 := hi_lt
+    have q' : i < hi % w - lo % w + 1 := by simp [Nat.mod_eq_sub_mul_div, h]; omega
+    have q'' : lo + i < w * vs.length := by omega
+    have h₁ : (lo + i) / w = lo / w := by
+      apply Nat.div_eq_of_lt_le <;> rw [Nat.mul_comm]
+      · omega
+      · simp only [Nat.mul_add]; omega
+    have h₂ : (lo + i) % w = lo % w + i := by simp [Nat.mod_eq_sub_mul_div, h₁]; omega
+    simp only [getLsbD_extractLsb, getLsbD_flattenList, q, q', q'', h₁, h₂, h, getLsbD_cast,
+      decide_true, Bool.true_and]
+  · simp at p
+    subst p
+    omega
+
+/-! ### rev -/
 theorem shiftRight_add {w : Nat} (x : BitVec w) (n m : Nat) :
     x >>> (n + m) = (x >>> n) >>> m:= by
   ext i

src/Init/Data/Nat/Lemmas.lean

@@ -1750,6 +1750,19 @@ grind_pattern shiftLeft_add => (m <<< n) <<< k where
 @[simp] theorem shiftLeft_shiftRight (x n : Nat) : x <<< n >>> n = x := by
   rw [Nat.shiftLeft_eq, Nat.shiftRight_eq_div_pow, Nat.mul_div_cancel _ (Nat.two_pow_pos _)]
 
+theorem lt_mul_div_self_add {x k : Nat} (h : 0 < k) : x < k * (x / k) + k := by
+  rw [← Int.ofNat_lt]
+  simpa using Int.lt_mul_ediv_self_add (by omega)
+
+theorem mul_sub_mod {x n p : Nat} (h : x < n * p) : (n * p - (x + 1)) % n = n - (x % n + 1) := by
+  rw [mod_eq_sub_div_mul, mul_sub_div _ _ _ h, Nat.sub_sub, Nat.add_comm, ← Nat.sub_sub,
+    Nat.mul_comm _ n, ← Nat.mul_sub_left_distrib, Nat.sub_sub_self, mul_succ, Nat.add_comm]
+  · conv in (_ + 1) => rw [← mod_add_div x n, Nat.add_right_comm]
+    apply Nat.add_sub_add_right
+  rw [Nat.succ_le_iff, Nat.div_lt_iff_lt_mul]
+  · rwa [Nat.mul_comm]
+  · refine Nat.pos_of_mul_pos_right (by omega : 0 < n * p)
+
 /-! ### Decidability of predicates -/
 
 -- `noncomputable` so the non-tail-recursive code is never compiled; the tail-recursive