From 5b23a1297aba9683f231c4b1a7ab4076af4ad53d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20G=2E=20Dorais?= <fgdorais@gmail.com>
Date: Tue, 14 Jan 2025 18:42:30 -0500
Subject: [PATCH] feat: add `Nat.ofBits` and `Fin.ofBits` (#1089)

---
 Batteries/Data/Fin.lean        |  1 +
 Batteries/Data/Fin/OfBits.lean | 16 +++++++++++++
 Batteries/Data/Nat/Basic.lean  |  7 +++++-
 Batteries/Data/Nat/Lemmas.lean | 42 ++++++++++++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 Batteries/Data/Fin/OfBits.lean

diff --git a/Batteries/Data/Fin.lean b/Batteries/Data/Fin.lean
index 7a5b9c16e8..c40ed38466 100644
--- a/Batteries/Data/Fin.lean
+++ b/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
diff --git a/Batteries/Data/Fin/OfBits.lean b/Batteries/Data/Fin/OfBits.lean
new file mode 100644
index 0000000000..55c081eb8b
--- /dev/null
+++ b/Batteries/Data/Fin/OfBits.lean
@@ -0,0 +1,16 @@
+/-
+Copyright (c) 2025 François G. Dorais. All rights reserved.
+Released under Apache 2. license as described in the file LICENSE.
+Authors: François G. Dorais
+-/
+
+import Batteries.Data.Nat.Lemmas
+
+namespace Fin
+
+/--
+Construct an element of `Fin (2 ^ n)` from a sequence of bits (little endian).
+-/
+abbrev ofBits (f : Fin n → Bool) : Fin (2 ^ n) := ⟨Nat.ofBits f, Nat.ofBits_lt_two_pow f⟩
+
+@[simp] theorem val_ofBits (f : Fin n → Bool) : (ofBits f).val = Nat.ofBits f := rfl
diff --git a/Batteries/Data/Nat/Basic.lean b/Batteries/Data/Nat/Basic.lean
index 21b4a4dd43..991791687f 100644
--- a/Batteries/Data/Nat/Basic.lean
+++ b/Batteries/Data/Nat/Basic.lean
@@ -6,7 +6,6 @@ Authors: Leonardo de Moura, Jeremy Avigad, Mario Carneiro
 
 namespace Nat
 
-
 /--
   Recursor identical to `Nat.recOn` but uses notations `0` for `Nat.zero` and `·+1` for `Nat.succ`
 -/
@@ -103,3 +102,9 @@ where
     else
       guess
   termination_by guess
+
+/--
+Construct a natural number from a sequence of bits using little endian convention.
+-/
+@[inline] def ofBits (f : Fin n → Bool) : Nat :=
+  Fin.foldr n (fun i v => 2 * v + (f i).toNat) 0
diff --git a/Batteries/Data/Nat/Lemmas.lean b/Batteries/Data/Nat/Lemmas.lean
index a97c558c44..944c473d33 100644
--- a/Batteries/Data/Nat/Lemmas.lean
+++ b/Batteries/Data/Nat/Lemmas.lean
@@ -162,3 +162,45 @@ protected def sum_trichotomy (a b : Nat) : a < b ⊕' a = b ⊕' b < a :=
 
 @[simp] theorem sum_append {l₁ l₂ : List Nat}: (l₁ ++ l₂).sum = l₁.sum + l₂.sum := by
   induction l₁ <;> simp [*, Nat.add_assoc]
+
+/-! ### ofBits -/
+
+@[simp] theorem ofBits_zero (f : Fin 0 → Bool) : ofBits f = 0 := rfl
+
+theorem ofBits_succ (f : Fin (n+1) → Bool) : ofBits f = 2 * ofBits (f ∘ Fin.succ) + (f 0).toNat :=
+  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_lt_add_left (Bool.toNat_lt _) ..
+      _ ≤ 2 * 2 ^ n := Nat.mul_le_mul_left 2 (ih ..)
+      _ = 2 ^ (n + 1) := Nat.pow_add_one' .. |>.symm
+
+@[simp] theorem testBit_ofBits_lt (f : Fin n → Bool) (i : Nat) (h : i < n) :
+    (ofBits f).testBit i = f ⟨i, h⟩ := by
+  induction n generalizing i with
+  | zero => contradiction
+  | succ n ih =>
+    simp only [ofBits_succ]
+    match i with
+    | 0 => simp [mod_eq_of_lt (Bool.toNat_lt _)]
+    | i+1 =>
+      rw [testBit_add_one, mul_add_div Nat.zero_lt_two, Nat.div_eq_of_lt (Bool.toNat_lt _)]
+      exact ih (f ∘ Fin.succ) i (Nat.lt_of_succ_lt_succ h)
+
+@[simp] theorem testBit_ofBits_ge (f : Fin n → Bool) (i : Nat) (h : n ≤ i) :
+    (ofBits f).testBit i = false := by
+  apply testBit_lt_two_pow
+  apply Nat.lt_of_lt_of_le
+  · exact ofBits_lt_two_pow f
+  · exact pow_le_pow_of_le_right Nat.zero_lt_two h
+
+theorem testBit_ofBits (f : Fin n → Bool) :
+    (ofBits f).testBit i = if h : i < n then f ⟨i, h⟩ else false := by
+  cases Nat.lt_or_ge i n with
+  | inl h => simp [h]
+  | inr h => simp [h, Nat.not_lt_of_ge h]