feat: add BitVec.ofFn and lemmas

Batteries.lean

@@ -16,13 +16,15 @@ import Batteries.Data.Array
 import Batteries.Data.AssocList
 import Batteries.Data.BinaryHeap
 import Batteries.Data.BinomialHeap
+import Batteries.Data.BitVec
 import Batteries.Data.ByteArray
 import Batteries.Data.ByteSubarray
 import Batteries.Data.Char
 import Batteries.Data.DList
 import Batteries.Data.Fin
 import Batteries.Data.FloatArray
 import Batteries.Data.HashMap
+import Batteries.Data.Int
 import Batteries.Data.LazyList
 import Batteries.Data.List
 import Batteries.Data.MLList

Batteries/Data/BitVec.lean

@@ -0,0 +1,88 @@
+/-
+Copyright (c) 2024 François G. Dorais. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: François G. Dorais
+-/
+
+import Batteries.Data.Fin.OfBits
+
+namespace BitVec
+
+theorem getElem_shiftConcat (v : BitVec n) (b : Bool) (i) (h : i < n) :
+    (v.shiftConcat b)[i] = if i = 0 then b else v[i-1] := by
+  rw [← getLsbD_eq_getElem, getLsbD_shiftConcat, getLsbD_eq_getElem, decide_eq_true h,
+    Bool.true_and]
+
+@[simp] theorem getElem_shiftConcat_zero (v : BitVec n) (b : Bool) (h : 0 < n) :
+    (v.shiftConcat b)[0] = b := by simp [getElem_shiftConcat]
+
+@[simp] theorem getElem_shiftConcat_succ (v : BitVec n) (b : Bool) (h : i + 1 < n) :
+    (v.shiftConcat b)[i+1] = v[i] := by simp [getElem_shiftConcat]
+
+/-- `ofFnLEAux f` returns the `BitVec m` whose `i`th bit is `f i` when `i < m`, little endian. -/
+@[inline] def ofFnLEAux (m : Nat) (f : Fin n → Bool) : BitVec m :=
+  Fin.foldr n (fun i v => v.shiftConcat (f i)) 0
+
+@[simp] theorem toNat_ofFnLEAux (m : Nat) (f : Fin n → Bool) :
+    (ofFnLEAux m f).toNat = Nat.ofBits f % 2 ^ m := by
+  simp only [ofFnLEAux]
+  induction n with
+  | zero => rfl
+  | succ n ih =>
+    rw [Fin.foldr_succ, toNat_shiftConcat, Nat.shiftLeft_eq, Nat.pow_one, Nat.ofBits_succ, ih,
+      ← Nat.mod_add_div (Nat.ofBits (f ∘ Fin.succ)) (2 ^ m), Nat.mul_add 2, Nat.add_right_comm,
+      Nat.mul_left_comm, Nat.add_mul_mod_self_left, Nat.mul_comm 2]
+    rfl
+
+@[simp] theorem toFin_ofFnLEAux (m : Nat) (f : Fin n → Bool) :
+    (ofFnLEAux m f).toFin = Fin.ofNat' (2 ^ m) (Nat.ofBits f) := by
+  ext; simp
+
+/-- `ofFnLE f` returns the `BitVec n` whose `i`th bit is `f i` with little endian ordering. -/
+@[inline] def ofFnLE (f : Fin n → Bool) : BitVec n := ofFnLEAux n f
+
+@[simp] theorem toNat_ofFnLE (f : Fin n → Bool) : (ofFnLE f).toNat = Nat.ofBits f := by
+  rw [ofFnLE, toNat_ofFnLEAux, Nat.mod_eq_of_lt (Nat.ofBits_lt_two_pow f)]
+
+@[simp] theorem toFin_ofFnLE (f : Fin n → Bool) : (ofFnLE f).toFin = Fin.ofBits f := by
+  ext; simp
+
+theorem getElem_ofFnLEAux (f : Fin n → Bool) (i) (h : i < n) (h' : i < m) :
+    (ofFnLEAux m f)[i] = f ⟨i, h⟩ := by
+  simp only [ofFnLEAux]
+  induction n generalizing i m with
+  | zero => contradiction
+  | succ n ih =>
+    simp only [Fin.foldr_succ, getElem_shiftConcat]
+    cases i with
+    | zero =>
+      simp
+    | succ i =>
+      rw [ih (fun i => f i.succ)] <;> try omega
+      simp
+
+@[simp] theorem getElem_ofFnLE (f : Fin n → Bool) (i) (h : i < n) : (ofFnLE f)[i] = f ⟨i, h⟩ :=
+  getElem_ofFnLEAux ..
+
+theorem getLsb'_ofFnLE (f : Fin n → Bool) (i) : (ofFnLE f).getLsb' i = f i := by simp
+
+@[simp] theorem getMsb'_ofFnLE (f : Fin n → Bool) (i) : (ofFnLE f).getMsb' i = f i.rev := by
+  simp [getMsb'_eq_getLsb', Fin.rev]; congr 2; omega
+
+/-- `ofFnBE f` returns the `BitVec n` whose `i`th bit is `f i` with big endian ordering. -/
+@[inline] def ofFnBE (f : Fin n → Bool) : BitVec n := ofFnLE fun i => f i.rev
+
+@[simp] theorem toNat_ofFnBE (f : Fin n → Bool) : (ofFnBE f).toNat = Nat.ofBits (f ∘ Fin.rev) := by
+  simp [ofFnBE]; rfl
+
+@[simp] theorem toFin_ofFnBE (f : Fin n → Bool) : (ofFnBE f).toFin = Fin.ofBits (f ∘ Fin.rev) := by
+  ext; simp
+
+@[simp] theorem getElem_ofFnBE (f : Fin n → Bool) (i) (h : i < n) :
+    (ofFnBE f)[i] = f (Fin.rev ⟨i, h⟩) := by simp [ofFnBE]
+
+theorem getLsb'_ofFnBE (f : Fin n → Bool) (i) : (ofFnBE f).getLsb' i = f i.rev := by
+  simp
+
+@[simp] theorem getMsb'_ofFnBE (f : Fin n → Bool) (i) : (ofFnBE f).getMsb' i = f i := by
+  simp [ofFnBE]

Batteries/Data/Fin.lean

@@ -1,3 +1,4 @@
 import Batteries.Data.Fin.Basic
 import Batteries.Data.Fin.Fold
 import Batteries.Data.Fin.Lemmas
+import Batteries.Data.Fin.OfBits

Batteries/Data/Fin/OfBits.lean

@@ -0,0 +1,14 @@
+/-
+Copyright (c) 2025 François G. Dorais. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: François G. Dorais
+-/
+
+import Batteries.Data.Nat.OfBits
+
+namespace Fin
+
+/-- Construct a `Fin (2 ^ n)` element from bit values (little endian). -/
+abbrev ofBits (f : Fin n → Bool) : Fin (2 ^ n) := ⟨Nat.ofBits f, Nat.ofBits_lt_two_pow f⟩
+
+theorem val_ofBits (f : Fin n → Bool) : (ofBits f).val = Nat.ofBits f := rfl

Batteries/Data/Int.lean

@@ -0,0 +1 @@
+import Batteries.Data.Int.OfBits

Batteries/Data/Int/OfBits.lean

@@ -0,0 +1,46 @@
+/-
+Copyright (c) 2025 François G. Dorais. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: François G. Dorais
+-/
+
+import Batteries.Data.Nat.OfBits
+
+namespace Int
+
+/-- `testBit m n` returns whether the `n+1` least significant bit is `1` or `0` -/
+@[inline] def testBit : Int → Nat → Bool
+  | .ofNat m, n => Nat.testBit m n
+  | .negSucc m, n => ! Nat.testBit m n
+
+
+/--
+Construct an integer from bit values (little endian).
+
+If `fill` is false, the result will be a nonnegative integer with all higher unspecified bits zero.
+If `fill` is true, the result will be a negative integer with all higher unspecified bits one.
+-/
+@[inline] def ofBitsFill (f : Fin n → Bool) : (fill : Bool := false) → Int
+  | false => Int.ofNat (Nat.ofBits f)
+  | true => Int.negSucc (Nat.ofBits fun i => ! f i)
+
+@[simp] theorem testBit_ofBitsFill_lt (f : Fin n → Bool) (fill : Bool) (i) (h : i < n) :
+    (ofBitsFill f fill).testBit i = f ⟨i, h⟩ := by
+  cases fill <;> simp [ofBitsFill, testBit, h]
+
+@[simp] theorem testBit_ofBitsFill_ge (f : Fin n → Bool) (fill : Bool) (i) (h : n ≤ i) :
+    (ofBitsFill f fill).testBit i = fill := by
+  cases fill <;> simp [ofBitsFill, testBit, h]
+
+/--
+Construct an integer from bit values (little endian). The most significant specified bit is used to
+determine the sign:
+* If the most significant specified bit is unset then the result will be nonnegative and all
+  higher order bits will read as unset.
+* If the most significant specified bit is set then the result will be negative and all higher
+  order bits will read as set.
+-/
+@[inline] def ofBits (f : Fin n → Bool) : Int :=
+  match n with
+  | 0 => 0
+  | n+1 => ofBitsFill f (f ⟨n, Nat.lt_succ_self n⟩)

Batteries/Data/Nat.lean

@@ -2,3 +2,5 @@ import Batteries.Data.Nat.Basic
 import Batteries.Data.Nat.Bisect
 import Batteries.Data.Nat.Gcd
 import Batteries.Data.Nat.Lemmas
+import Batteries.Data.Nat.OfBits
+

Batteries/Data/Nat/OfBits.lean

@@ -0,0 +1,55 @@
+/-
+Copyright (c) 2025 François G. Dorais. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: François G. Dorais
+-/
+
+namespace Nat
+
+/-- Construct a natural number from bit values (little endian). -/
+@[inline] def ofBits (f : Fin n → Bool) : Nat :=
+  Fin.foldr n (fun i v => 2 * v + (f i).toNat) 0
+
+@[simp] theorem ofBits_zero (f : Fin 0 → Bool) : ofBits f = 0 := rfl
+
+@[simp] theorem ofBits_one (f : Fin 1 → Bool) : ofBits f = (f 0).toNat := Nat.zero_add ..
+
+theorem ofBits_succ (f : Fin (n+1) → Bool) :
+    ofBits f = 2 * ofBits (f ∘ Fin.succ) + (f 0).toNat := by
+  simp [ofBits, Fin.foldr_succ]
+
+theorem ofBits_lt_two_pow (f : Fin n → Bool) : ofBits f < 2 ^ n := by
+  induction n with
+  | zero => simp
+  | succ n ih =>
+    calc ofBits f
+        = 2 * ofBits (f ∘ Fin.succ) + (f 0).toNat := ofBits_succ ..
+      _ ≤ 2 * ofBits (f ∘ Fin.succ) + 1 := Nat.add_le_add_left (Bool.toNat_le _) ..
+      _ < 2 * (ofBits (f ∘ Fin.succ) + 1) := Nat.lt_add_one ..
+      _ ≤ 2 * 2 ^ n := Nat.mul_le_mul_left 2 (Nat.succ_le_of_lt (ih _))
+      _ = 2 ^ (n + 1) := (Nat.pow_add_one' ..).symm
+
+@[simp] theorem testBit_ofBits_lt (f : Fin n → Bool) (i) (h : i < n) :
+    (ofBits f).testBit i = f ⟨i, h⟩ := by
+  induction n generalizing i with
+  | zero => contradiction
+  | succ n ih =>
+    have h0 : (f 0).toNat < 2 := Bool.toNat_lt ..
+    match i with
+    | 0 =>
+      rw [ofBits_succ, testBit_zero, Nat.mul_add_mod, Nat.mod_eq_of_lt h0, Fin.zero_eta]
+      cases f 0 <;> rfl
+    | i+1 =>
+      rw [ofBits_succ, testBit_add_one, Nat.mul_add_div Nat.zero_lt_two, Nat.div_eq_of_lt h0]
+      exact ih (f ∘ Fin.succ) i (Nat.lt_of_succ_lt_succ h)
+
+@[simp] theorem testBit_ofBits_ge (f : Fin n → Bool) (i) (h : n ≤ i) :
+    (ofBits f).testBit i = false := by
+  induction n generalizing i with
+  | zero => simp
+  | succ n ih =>
+    have h0 : (f 0).toNat < 2 := Bool.toNat_lt ..
+    match i with
+    | i+1 =>
+      rw [ofBits_succ, testBit_succ, Nat.mul_add_div Nat.zero_lt_two, Nat.div_eq_of_lt h0]
+      exact ih (f ∘ Fin.succ) i (Nat.le_of_succ_le_succ h)