chore: upstream Nat.binaryRec

src/Init/Data/BitVec/Lemmas.lean

@@ -166,8 +166,7 @@ theorem msb_eq_getLsb_last (x : BitVec w) :
     x.getLsb (w-1) = decide (2 ^ (w-1) ≤ x.toNat) := by
   rcases w with rfl | w
   · simp
-  · simp only [Nat.zero_lt_succ, decide_True, getLsb, Nat.testBit, Nat.succ_sub_succ_eq_sub,
-    Nat.sub_zero, Nat.and_one_is_mod, Bool.true_and, Nat.shiftRight_eq_div_pow]
+  · simp only [getLsb, Nat.testBit_to_div_mod, Nat.succ_sub_succ_eq_sub, Nat.sub_zero]
     rcases (Nat.lt_or_ge (BitVec.toNat x) (2 ^ w)) with h | h
     · simp [Nat.div_eq_of_lt h, h]
     · simp only [h]
@@ -716,7 +715,6 @@ theorem toNat_cons' {x : BitVec w} :
     have p2 : i ≠ n := by omega
     simp [p1, p2]
   · simp [i_eq_n, testBit_toNat]
-    cases b <;> trivial
   · have p1 : i ≠ n := by omega
     have p2 : i - n ≠ 0 := by omega
     simp [p1, p2, Nat.testBit_bool_to_nat]

src/Init/Data/Bool.lean

@@ -360,6 +360,8 @@ theorem toNat_lt (b : Bool) : b.toNat < 2 :=
   cases b <;> simp
 @[simp] theorem toNat_eq_one  (b : Bool) : b.toNat = 1 ↔ b = true := by
   cases b <;> simp
+@[simp] theorem toNat_mod_two (b : Bool) : b.toNat % 2 = b.toNat := by
+  cases b <;> simp
 
 /-! ### ite -/
 

src/Init/Data/Nat/Bitwise/Basic.lean

@@ -71,13 +71,92 @@ theorem shiftRight_eq_div_pow (m : Nat) : ∀ n, m >>> n = m / 2 ^ n
     rw [shiftRight_add, shiftRight_eq_div_pow m k]
     simp [Nat.div_div_eq_div_mul, ← Nat.pow_succ, shiftRight_succ]
 
+/-- `bit b` appends the digit `b` to the little end of the binary representation of
+  its natural number input. -/
+def bit (b : Bool) (n : Nat) : Nat :=
+  cond b (n + n + 1) (n + n)
+
 /-!
 ### testBit
 We define an operation for testing individual bits in the binary representation
 of a number.
 -/
 
 /-- `testBit m n` returns whether the `(n+1)` least significant bit is `1` or `0`-/
-def testBit (m n : Nat) : Bool := (m >>> n) &&& 1 != 0
+def testBit (m n : Nat) : Bool :=
+  -- `1 &&& n` is faster than `n &&& 1` for big `n`. This may change in the future.
+  1 &&& (m >>> n) != 0
+
+theorem shiftRight_one (n) : n >>> 1 = n / 2 := rfl
+
+theorem mod_two_eq_zero_or_one (n : Nat) : n % 2 = 0 ∨ n % 2 = 1 :=
+  match n % 2, @Nat.mod_lt n 2 (by decide) with
+  | 0, _ => .inl rfl
+  | 1, _ => .inr rfl
+
+@[simp] theorem one_land_eq_mod_two (n : Nat) : 1 &&& n = n % 2 := by
+  match Nat.decEq n 0 with
+  | isTrue n0 => subst n0; decide
+  | isFalse n0 =>
+    simp only [HAnd.hAnd, AndOp.and, land]
+    unfold bitwise
+    cases mod_two_eq_zero_or_one n with | _ h => simp [n0, h]; rfl
+
+@[simp] theorem testBit_zero (n : Nat) : testBit n 0 = decide (n % 2 = 1) := by
+  cases mod_two_eq_zero_or_one n with | _ h => simp [testBit, h]
+
+theorem bit_testBit_zero_shiftRight_one (n : Nat) : bit (n.testBit 0) (n >>> 1) = n := by
+  simp only [bit, testBit_zero]
+  cases mod_two_eq_zero_or_one n with | _ h =>
+    simpa [h, shiftRight_one] using Eq.trans (by simp [h, Nat.two_mul]) (Nat.div_add_mod n 2)
+
+theorem bit_eq_zero_iff {n : Nat} {b : Bool} : bit b n = 0 ↔ n = 0 ∧ b = false := by
+  cases n <;> cases b <;> simp [bit, ← Nat.add_assoc]
+
+/-- For a predicate `C : Nat → Sort u`, if instances can be
+  constructed for natural numbers of the form `bit b n`,
+  they can be constructed for any given natural number. -/
+@[inline]
+def bitCasesOn {C : Nat → Sort u} (n) (h : ∀ b n, C (bit b n)) : C n :=
+  -- `1 &&& n != 0` is faster than `n.testBit 0`. This may change when we have faster `testBit`.
+  let x := h (1 &&& n != 0) (n >>> 1)
+  -- `congrArg C _` is `rfl` in non-dependent case
+  congrArg C n.bit_testBit_zero_shiftRight_one ▸ x
+
+/-- A recursion principle for `bit` representations of natural numbers.
+  For a predicate `C : Nat → Sort u`, if instances can be
+  constructed for natural numbers of the form `bit b n`,
+  they can be constructed for all natural numbers. -/
+@[elab_as_elim, specialize]
+def binaryRec {C : Nat → Sort u} (z : C 0) (f : ∀ b n, C n → C (bit b n)) (n : Nat) : C n :=
+  if n0 : n = 0 then congrArg C n0 ▸ z
+  else
+    let x := f (1 &&& n != 0) (n >>> 1) (binaryRec z f (n >>> 1))
+    congrArg C n.bit_testBit_zero_shiftRight_one ▸ x
+decreasing_by exact bitwise_rec_lemma n0
+
+/-- The same as `binaryRec`, but the induction step can assume that if `n=0`,
+  the bit being appended is `true`-/
+@[elab_as_elim, specialize]
+def binaryRec' {C : Nat → Sort u} (z : C 0)
+    (f : ∀ b n, (n = 0 → b = true) → C n → C (bit b n)) : ∀ n, C n :=
+  binaryRec z fun b n ih =>
+    if h : n = 0 → b = true then f b n h ih
+    else
+      have : bit b n = 0 := by
+        rw [bit_eq_zero_iff]
+        cases n <;> cases b <;> simp at h <;> simp [h]
+      congrArg C this ▸ z
+
+/-- The same as `binaryRec`, but special casing both 0 and 1 as base cases -/
+@[elab_as_elim, specialize]
+def binaryRecFromOne {C : Nat → Sort u} (z₀ : C 0) (z₁ : C 1)
+    (f : ∀ b n, n ≠ 0 → C n → C (bit b n)) : ∀ n, C n :=
+  binaryRec' z₀ fun b n h ih =>
+    if h' : n = 0 then
+      have : bit b n = bit true 0 := by
+        rw [h', h h']
+      congrArg C this ▸ z₁
+    else f b n h' ih
 
 end Nat

src/Init/Data/Nat/Bitwise/Lemmas.lean

@@ -35,51 +35,61 @@ private theorem two_mul_sub_one {n : Nat} (n_pos : n > 0) : (2*n - 1) % 2 = 1 :=
 
 /-! ### Preliminaries -/
 
-/--
-An induction principal that works on divison by two.
--/
-noncomputable def div2Induction {motive : Nat → Sort u}
-    (n : Nat) (ind : ∀(n : Nat), (n > 0 → motive (n/2)) → motive n) : motive n := by
-  induction n using Nat.strongInductionOn with
-  | ind n hyp =>
-    apply ind
-    intro n_pos
-    if n_eq : n = 0 then
-      simp [n_eq] at n_pos
-    else
-      apply hyp
-      exact Nat.div_lt_self n_pos (Nat.le_refl _)
-
 @[simp] theorem zero_and (x : Nat) : 0 &&& x = 0 := by rfl
 
 @[simp] theorem and_zero (x : Nat) : x &&& 0 = 0 := by
   simp only [HAnd.hAnd, AndOp.and, land]
   unfold bitwise
   simp
 
-@[simp] theorem and_one_is_mod (x : Nat) : x &&& 1 = x % 2 := by
-  if xz : x = 0 then
-    simp [xz, zero_and]
-  else
-    have andz := and_zero (x/2)
-    simp only [HAnd.hAnd, AndOp.and, land] at andz
-    simp only [HAnd.hAnd, AndOp.and, land]
-    unfold bitwise
-    cases mod_two_eq_zero_or_one x with | _ p =>
-      simp [xz, p, andz, one_div_two, mod_eq_of_lt]
+/-! ### bitwise -/
+
+@[simp]
+theorem bitwise_zero_left (m : Nat) : bitwise f 0 m = if f false true then m else 0 :=
+  rfl
+
+@[simp]
+theorem bitwise_zero_right (n : Nat) : bitwise f n 0 = if f true false then n else 0 := by
+  unfold bitwise
+  simp only [ite_self, decide_False, Nat.zero_div, ite_true]
+  cases n <;> simp
+
+theorem bitwise_zero : bitwise f 0 0 = 0 := by
+  simp only [bitwise_zero_right, ite_self]
+
+/-! ### bit -/
+
+theorem bit_val (b n) : bit b n = 2 * n + b.toNat := by
+  rw [Nat.mul_comm]
+  induction b with
+  | false => exact congrArg (· + n) n.zero_add.symm
+  | true => exact congrArg (· + n + 1) n.zero_add.symm
+
+@[simp]
+theorem bit_div_two (b n) : bit b n / 2 = n := by
+  rw [bit_val, Nat.add_comm, add_mul_div_left, div_eq_of_lt, Nat.zero_add]
+  · cases b <;> decide
+  · decide
+
+@[simp]
+theorem bit_mod_two (b n) : bit b n % 2 = b.toNat := by
+  cases b <;> simp [bit_val, mul_add_mod]
+
+theorem mul_two_le_bit {x b n} : x * 2 ≤ bit b n ↔ x ≤ n := by
+  rw [← le_div_iff_mul_le Nat.two_pos, bit_div_two]
 
 /-! ### testBit -/
 
 @[simp] theorem zero_testBit (i : Nat) : testBit 0 i = false := by
-  simp only [testBit, zero_shiftRight, zero_and, bne_self_eq_false]
+  simp only [testBit, zero_shiftRight, and_zero, bne_self_eq_false]
 
-@[simp] theorem testBit_zero (x : Nat) : testBit x 0 = decide (x % 2 = 1) := by
-  cases mod_two_eq_zero_or_one x with | _ p => simp [testBit, p]
-
-@[simp] theorem testBit_succ (x i : Nat) : testBit x (succ i) = testBit (x/2) i := by
+@[simp] theorem testBit_succ (x i : Nat) : testBit x (i + 1) = testBit (x / 2) i := by
   unfold testBit
   simp [shiftRight_succ_inside]
 
+theorem bit_testBit_zero (b n) : (bit b n).testBit 0 = b := by
+  simp
+
 theorem testBit_to_div_mod {x : Nat} : testBit x i = decide (x / 2^i % 2 = 1) := by
   induction i generalizing x with
   | zero =>
@@ -88,51 +98,37 @@ theorem testBit_to_div_mod {x : Nat} : testBit x i = decide (x / 2^i % 2 = 1) :=
   | succ i hyp =>
     simp [hyp, Nat.div_div_eq_div_mul, Nat.pow_succ']
 
+@[simp] theorem _root_.Bool.toNat_testBit_zero (b : Bool) : b.toNat.testBit 0 = b := by
+  cases b <;> rfl
+
 theorem toNat_testBit (x i : Nat) :
     (x.testBit i).toNat = x / 2 ^ i % 2 := by
   rw [Nat.testBit_to_div_mod]
   rcases Nat.mod_two_eq_zero_or_one (x / 2^i) <;> simp_all
 
 theorem ne_zero_implies_bit_true {x : Nat} (xnz : x ≠ 0) : ∃ i, testBit x i := by
-  induction x using div2Induction with
-  | ind x hyp =>
-    have x_pos : x > 0 := Nat.pos_of_ne_zero xnz
-    match mod_two_eq_zero_or_one x with
-    | Or.inl mod2_eq =>
-      rw [←div_add_mod x 2] at xnz
-      simp only [mod2_eq, ne_eq, Nat.mul_eq_zero, Nat.add_zero, false_or] at xnz
-      have ⟨d, dif⟩   := hyp x_pos xnz
-      apply Exists.intro (d+1)
-      simp_all
-    | Or.inr mod2_eq =>
-      apply Exists.intro 0
-      simp_all
+  induction x using binaryRecFromOne with
+  | z₀ => exact absurd rfl xnz
+  | z₁ => exact ⟨0, rfl⟩
+  | f b n n0 hyp =>
+    obtain ⟨i, h⟩ := hyp n0
+    refine ⟨i + 1, ?_⟩
+    rwa [testBit_succ, bit_div_two]
 
 theorem ne_implies_bit_diff {x y : Nat} (p : x ≠ y) : ∃ i, testBit x i ≠ testBit y i := by
-  induction y using Nat.div2Induction generalizing x with
-  | ind y hyp =>
-    cases Nat.eq_zero_or_pos y with
-    | inl yz =>
-      simp only [yz, Nat.zero_testBit, Bool.eq_false_iff]
-      simp only [yz] at p
-      have ⟨i,ip⟩  := ne_zero_implies_bit_true p
-      apply Exists.intro i
-      simp [ip]
-    | inr ypos =>
-      if lsb_diff : x % 2 = y % 2 then
-        rw [←Nat.div_add_mod x 2, ←Nat.div_add_mod y 2] at p
-        simp only [ne_eq, lsb_diff, Nat.add_right_cancel_iff,
-          Nat.zero_lt_succ, Nat.mul_left_cancel_iff] at p
-        have ⟨i, ieq⟩  := hyp ypos p
-        apply Exists.intro (i+1)
-        simpa
-      else
-        apply Exists.intro 0
-        simp only [testBit_zero]
-        revert lsb_diff
-        cases mod_two_eq_zero_or_one x with | _ p =>
-          cases mod_two_eq_zero_or_one y with | _ q =>
-            simp [p,q]
+  induction y using binaryRec' generalizing x with
+  | z =>
+    simp only [zero_testBit, Bool.ne_false_iff]
+    exact ne_zero_implies_bit_true p
+  | f yb y h hyp =>
+    rw [← x.bit_testBit_zero_shiftRight_one] at *
+    by_cases hb : testBit x 0 = yb
+    · subst hb
+      obtain ⟨i, h⟩ := hyp (mt (congrArg _) p)
+      refine ⟨i + 1, ?_⟩
+      rwa [testBit_succ, bit_div_two, testBit_succ, bit_div_two]
+    · refine ⟨0, ?_⟩
+      simpa using hb
 
 /--
 `eq_of_testBit_eq` allows proving two natural numbers are equal
@@ -146,25 +142,18 @@ theorem eq_of_testBit_eq {x y : Nat} (pred : ∀i, testBit x i = testBit y i) :
     have p := pred i
     contradiction
 
-theorem ge_two_pow_implies_high_bit_true {x : Nat} (p : x ≥ 2^n) : ∃ i, i ≥ n ∧ testBit x i := by
-  induction x using div2Induction generalizing n with
-  | ind x hyp =>
-    have x_pos : x > 0 := Nat.lt_of_lt_of_le (Nat.two_pow_pos n) p
-    have x_ne_zero : x ≠ 0 := Nat.ne_of_gt x_pos
+theorem ge_two_pow_implies_high_bit_true {x : Nat} (p : 2^n ≤ x) : ∃ i, n ≤ i ∧ testBit x i := by
+  induction x using binaryRec generalizing n with
+  | z => exact absurd p (Nat.not_le_of_lt n.two_pow_pos)
+  | f xb x hyp =>
     match n with
-    | zero =>
-      let ⟨j, jp⟩ := ne_zero_implies_bit_true x_ne_zero
-      exact Exists.intro j (And.intro (Nat.zero_le _) jp)
-    | succ n =>
-      have x_ge_n : x / 2 ≥ 2 ^ n := by
-          simpa [le_div_iff_mul_le, ← Nat.pow_succ'] using p
-      have ⟨j, jp⟩ := @hyp x_pos n x_ge_n
-      apply Exists.intro (j+1)
-      apply And.intro
-      case left =>
-        exact (Nat.succ_le_succ jp.left)
-      case right =>
-        simpa using jp.right
+    | 0 =>
+      simp only [zero_le, true_and]
+      exact ne_zero_implies_bit_true (ne_of_gt (Nat.lt_of_succ_le p))
+    | n+1 =>
+      obtain ⟨i, h⟩ := hyp (mul_two_le_bit.mp p)
+      refine ⟨i + 1, ?_⟩
+      rwa [Nat.add_le_add_iff_right, testBit_succ, bit_div_two]
 
 theorem testBit_implies_ge {x : Nat} (p : testBit x i = true) : x ≥ 2^i := by
   simp only [testBit_to_div_mod] at p
@@ -188,8 +177,6 @@ theorem lt_pow_two_of_testBit (x : Nat) (p : ∀i, i ≥ n → testBit x i = fal
   have test_false := p _ i_ge_n
   simp only [test_true] at test_false
 
-/-! ### testBit -/
-
 private theorem succ_mod_two : succ x % 2 = 1 - x % 2 := by
   induction x with
   | zero =>
@@ -233,7 +220,7 @@ theorem testBit_two_pow_add_gt {i j : Nat} (j_lt_i : j < i) (x : Nat) :
     rw [Nat.sub_eq_zero_iff_le] at i_sub_j_eq
     exact Nat.not_le_of_gt j_lt_i i_sub_j_eq
   | d+1 =>
-    simp [Nat.pow_succ, Nat.mul_comm _ 2,  Nat.mul_add_mod]
+    simp [Nat.pow_succ, Nat.mul_comm _ 2, Nat.mul_add_mod]
 
 @[simp] theorem testBit_mod_two_pow (x j i : Nat) :
     testBit (x % 2^j) i = (decide (i < j) && testBit x i) := by
@@ -257,7 +244,7 @@ theorem testBit_two_pow_add_gt {i j : Nat} (j_lt_i : j < i) (x : Nat) :
             exact Nat.lt_add_of_pos_right (Nat.two_pow_pos j)
       simp only [hyp y y_lt_x]
       if i_lt_j : i < j then
-        rw [ Nat.add_comm _ (2^_), testBit_two_pow_add_gt i_lt_j]
+        rw [Nat.add_comm _ (2^_), testBit_two_pow_add_gt i_lt_j]
       else
         simp [i_lt_j]
 
@@ -274,7 +261,7 @@ theorem testBit_two_pow_sub_succ (h₂ : x < 2 ^ n) (i : Nat) :
     match n with
     | 0 => simp [succ_sub_succ_eq_sub]
     | n+1 =>
-      simp [not_decide_mod_two_eq_one]
+      simp [← decide_not]
       omega
   | succ i ih =>
     simp only [testBit_succ]
@@ -487,3 +474,38 @@ theorem mul_add_lt_is_or {b : Nat} (b_lt : b < 2^i) (a : Nat) : 2^i * a + b = 2^
 
 @[simp] theorem testBit_shiftRight (x : Nat) : testBit (x >>> i) j = testBit x (i+j) := by
   simp [testBit, ←shiftRight_add]
+
+@[simp] theorem bit_shiftRight_one (b n) : bit b n >>> 1 = n :=
+  bit_div_two b n
+
+/-! ### binaryRec -/
+
+@[simp]
+theorem bitCasesOn_bit {C : Nat → Sort u} (h : ∀ b n, C (bit b n)) (b : Bool) (n : Nat) :
+    bitCasesOn (bit b n) h = h b n := by
+  change congrArg C (bit b n).bit_testBit_zero_shiftRight_one ▸ h _ _ = h b n
+  generalize congrArg C (bit b n).bit_testBit_zero_shiftRight_one = e; revert e
+  rw [bit_testBit_zero, bit_shiftRight_one]
+  intros; rfl
+
+@[simp]
+theorem binaryRec_zero {C : Nat → Sort u} (z : C 0) (f : ∀ b n, C n → C (bit b n)) :
+    binaryRec z f 0 = z :=
+  rfl
+
+theorem binaryRec_eq {C : Nat → Sort u} {z : C 0} {f : ∀ b n, C n → C (bit b n)} (b n)
+    (h : f false 0 z = z ∨ (n = 0 → b = true)) :
+    binaryRec z f (bit b n) = f b n (binaryRec z f n) := by
+  by_cases h' : bit b n = 0
+  case pos =>
+    obtain ⟨rfl, rfl⟩ := bit_eq_zero_iff.mp h'
+    simp only [forall_const, or_false] at h
+    exact h.symm
+  case neg =>
+    rw [binaryRec, dif_neg h']
+    change congrArg C (bit b n).bit_testBit_zero_shiftRight_one ▸ f _ _ _ = _
+    generalize congrArg C (bit b n).bit_testBit_zero_shiftRight_one = e; revert e
+    rw [bit_testBit_zero, bit_shiftRight_one]
+    intros; rfl
+
+end Nat