Skip to content

Commit

Permalink
feat: add Nat.ofBits and Fin.ofBits (#1089)
Browse files Browse the repository at this point in the history
  • Loading branch information
fgdorais authored Jan 14, 2025
1 parent c104265 commit 5b23a12
Show file tree
Hide file tree
Showing 4 changed files with 65 additions and 1 deletion.
1 change: 1 addition & 0 deletions Batteries/Data/Fin.lean
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import Batteries.Data.Fin.Basic
import Batteries.Data.Fin.Fold
import Batteries.Data.Fin.Lemmas
import Batteries.Data.Fin.OfBits
16 changes: 16 additions & 0 deletions Batteries/Data/Fin/OfBits.lean
Original file line number Diff line number Diff line change
@@ -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
7 changes: 6 additions & 1 deletion Batteries/Data/Nat/Basic.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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`
-/
Expand Down Expand Up @@ -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
42 changes: 42 additions & 0 deletions Batteries/Data/Nat/Lemmas.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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]

0 comments on commit 5b23a12

Please sign in to comment.