feat(data/nat): Add new binary recursors; prove the relationship between nat.bits and other pieces of code (#14990)

This is in connection to #13208, because I was asked to write `to_bits` in terms of `nat.size` and `nat.bits`, but `nat.bits` was hard to use (there were no lemmas about it). This PR proves some statements about `nat.bits`, so that it can finally be used.

Author: prakol16

src/data/list/basic.lean

@@ -4086,6 +4086,9 @@ lemma inth_eq_iget_nth (n : ℕ) :
   l.inth n = (l.nth n).iget :=
 by rw [←nthd_default_eq_inth, nthd_eq_get_or_else_nth, option.get_or_else_default_eq_iget]
 
+lemma inth_zero_eq_head : l.inth 0 = l.head :=
+by { cases l; refl, }
+
 end inth
 
 end list

src/data/nat/bits.lean

@@ -0,0 +1,94 @@
+/-
+Copyright (c) 2022 Praneeth Kolichala. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Praneeth Kolichala
+-/
+import tactic.generalize_proofs
+import tactic.norm_num
+
+/-!
+# Additional properties of binary recursion on `nat`
+
+This file documents additional properties of binary recursion,
+which allows us to more easily work with operations which do depend
+on the number of leading zeros in the binary representation of `n`.
+For example, we can more easily work with `nat.bits` and `nat.size`.
+
+See also: `nat.bitwise`, `nat.pow` (for various lemmas about `size` and `shiftl`/`shiftr`),
+and `nat.digits`.
+-/
+
+namespace nat
+
+lemma bit_eq_zero_iff {n : ℕ} {b : bool} : bit b n = 0 ↔ n = 0 ∧ b = ff :=
+by { split, { cases b; simp [nat.bit], }, rintro ⟨rfl, rfl⟩, refl, }
+
+/-- The same as binary_rec_eq, but that one unfortunately requires `f` to be the identity when
+  appending `ff` to `0`. Here, we allow you to explicitly say that that case is not happening, i.e.
+  supplying `n = 0 → b = tt`. -/
+lemma binary_rec_eq' {C : ℕ → Sort*} {z : C 0} {f : ∀ b n, C n → C (bit b n)}
+  (b n) (h : f ff 0 z = z ∨ (n = 0 → b = tt)) :
+  binary_rec z f (bit b n) = f b n (binary_rec z f n) :=
+begin
+  rw [binary_rec],
+  split_ifs with h',
+  { rcases bit_eq_zero_iff.mp h' with ⟨rfl, rfl⟩,
+    rw binary_rec_zero,
+    simp only [imp_false, or_false, eq_self_iff_true, not_true] at h,
+    exact h.symm },
+  { generalize_proofs e, revert e,
+    rw [bodd_bit, div2_bit],
+    intros, refl, }
+end
+
+/-- The same as `binary_rec`, but the induction step can assume that if `n=0`,
+  the bit being appended is `tt`-/
+@[elab_as_eliminator]
+def binary_rec' {C : ℕ → Sort*} (z : C 0) (f : ∀ b n, (n = 0 → b = tt) → C n → C (bit b n)) :
+  ∀ n, C n :=
+binary_rec z (λ b n ih, if h : n = 0 → b = tt then f b n h ih else
+  by { convert z, rw bit_eq_zero_iff, simpa using h, })
+
+/-- The same as `binary_rec`, but special casing both 0 and 1 as base cases -/
+@[elab_as_eliminator]
+def binary_rec_from_one {C : ℕ → Sort*} (z₀ : C 0) (z₁ : C 1)
+  (f : ∀ b n, n ≠ 0 → C n → C (bit b n)) : ∀ n, C n :=
+binary_rec' z₀ (λ b n h ih, if h' : n = 0 then by { rw [h', h h'], exact z₁ } else f b n h' ih)
+
+@[simp] lemma zero_bits : bits 0 = [] := by simp [nat.bits]
+
+@[simp] lemma bits_append_bit (n : ℕ) (b : bool) (hn : n = 0 → b = tt) :
+  (bit b n).bits = b :: n.bits :=
+by { rw [nat.bits, binary_rec_eq'], simpa, }
+
+@[simp] lemma bit0_bits (n : ℕ) (hn : n ≠ 0) : (bit0 n).bits = ff :: n.bits :=
+bits_append_bit n ff (λ hn', absurd hn' hn)
+
+@[simp] lemma bit1_bits (n : ℕ) : (bit1 n).bits = tt :: n.bits :=
+bits_append_bit n tt (λ _, rfl)
+
+@[simp] lemma one_bits : nat.bits 1 = [tt] := by { convert bit1_bits 0, simp, }
+
+example : bits 3423 = [tt, tt, tt, tt, tt, ff, tt, ff, tt, ff, tt, tt] := by norm_num
+
+lemma bodd_eq_bits_head (n : ℕ) : n.bodd = n.bits.head :=
+begin
+  induction n using nat.binary_rec' with b n h ih, { simp, },
+  simp [bodd_bit, bits_append_bit _ _ h],
+end
+
+lemma div2_bits_eq_tail (n : ℕ) : n.div2.bits = n.bits.tail :=
+begin
+  induction n using nat.binary_rec' with b n h ih, { simp, },
+  simp [div2_bit, bits_append_bit _ _ h],
+end
+
+lemma size_eq_bits_len (n : ℕ) : n.bits.length = n.size :=
+begin
+  induction n using nat.binary_rec' with b n h ih, { simp, },
+  rw [size_bit, bits_append_bit _ _ h],
+  { simp [ih], },
+  { simpa [bit_eq_zero_iff], }
+end
+
+end nat

src/data/nat/bitwise.lean

@@ -3,6 +3,7 @@ Copyright (c) 2020 Markus Himmel. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Markus Himmel
 -/
+import data.nat.bits
 import tactic.linarith
 
 /-!
@@ -51,6 +52,16 @@ end
 @[simp] lemma zero_test_bit (i : ℕ) : test_bit 0 i = ff :=
 by simp [test_bit]
 
+/-- The ith bit is the ith element of `n.bits`. -/
+lemma test_bit_eq_inth (n i : ℕ) : n.test_bit i = n.bits.inth i :=
+begin
+  induction i with i ih generalizing n,
+  { simp [test_bit, shiftr, bodd_eq_bits_head, list.inth_zero_eq_head], },
+  conv_lhs { rw ← bit_decomp n, },
+  rw [test_bit_succ, ih n.div2, div2_bits_eq_tail],
+  cases n.bits; simp,
+end
+
 /-- Bitwise extensionality: Two numbers agree if they agree at every bit position. -/
 lemma eq_of_test_bit_eq {n m : ℕ} (h : ∀ i, test_bit n i = test_bit m i) : n = m :=
 begin

src/data/nat/digits.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison, Shing Tak Lam, Mario Carneiro
 -/
 import data.int.modeq
+import data.nat.bits
 import data.nat.log
 import data.nat.parity
 import data.list.indexes
@@ -471,6 +472,22 @@ begin
   exact base_pow_length_digits_le' b m,
 end
 
+/-! ### Binary -/
+lemma digits_two_eq_bits (n : ℕ) : digits 2 n = n.bits.map (λ b, cond b 1 0) :=
+begin
+  induction n using nat.binary_rec_from_one with b n h ih,
+  { simp, },
+  { simp, },
+  rw bits_append_bit _ _ (λ hn, absurd hn h),
+  cases b,
+  { rw digits_def' (le_refl 2),
+     { simpa [nat.bit, nat.bit0_val n], },
+     { simpa [pos_iff_ne_zero, bit_eq_zero_iff], }, },
+  { simpa [nat.bit, nat.bit1_val n, add_comm, digits_add 2 le_rfl 1 n (by norm_num)
+    (by norm_num)] },
+end
+
+
 /-! ### Modular Arithmetic -/
 
 -- This is really a theorem about polynomials.