feat: Decidable instance for Nat.isPowerOfTwo (#11905)

This PR provides a `Decidable` instance for `Nat.isPowerOfTwo` based on
the formula `(n ≠ 0) ∧ (n &&& (n - 1)) = 0`.

To do this it includes theorems about `Nat.testBit` to show that the
`n.log2`th bit is set in `n` and `n - 1` for non powers of two.

Bitwise lemmas are needed to reason about the `&&&` so the file
`Init.Data.Nat.Power2` is renamed to `Init.Data.Nat.Power2.Basic` and
`Init.Data.Nat.Power2.Lemmas` introduced that depends on
`Init.Data.Nat.Bitwise.Lemmas` to prevent circular includes.

---------

Co-authored-by: Kim Morrison <477956+kim-em@users.noreply.github.com>

Author: georgerennie

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

@@ -223,6 +223,16 @@ theorem testBit_lt_two_pow {x i : Nat} (lt : x < 2^i) : x.testBit i = false := b
     exfalso
     exact Nat.not_le_of_gt lt (ge_two_pow_of_testBit p)
 
+theorem testBit_of_two_pow_le_and_two_pow_add_one_gt {n i : Nat}
+    (hle : 2^i ≤ n) (hgt : n < 2^(i + 1)) : n.testBit i = true := by
+  rcases exists_ge_and_testBit_of_ge_two_pow hle with ⟨i', ⟨_, _⟩⟩
+  have : i = i' := by
+    false_or_by_contra
+    have : 2 ^ (i + 1) ≤ 2 ^ i' := Nat.pow_le_pow_of_le (by decide) (by omega)
+    have : n.testBit i' = false := testBit_lt_two_pow (by omega)
+    simp_all only [Bool.false_eq_true]
+  rwa [this]
+
 theorem lt_pow_two_of_testBit (x : Nat) (p : ∀i, i ≥ n → testBit x i = false) : x < 2^n := by
   apply Decidable.by_contra
   intro not_lt
@@ -231,6 +241,10 @@ theorem lt_pow_two_of_testBit (x : Nat) (p : ∀i, i ≥ n → testBit x i = fal
   have test_false := p _ i_ge_n
   simp [test_true] at test_false
 
+theorem testBit_log2 {n : Nat} (h : n ≠ 0) : n.testBit n.log2 = true := by
+  have := log2_eq_iff (n := n) (k := n.log2) (by omega)
+  apply testBit_of_two_pow_le_and_two_pow_add_one_gt <;> omega
+
 private theorem succ_mod_two : succ x % 2 = 1 - x % 2 := by
   induction x with
   | zero =>

src/Init/Data/Nat/Lemmas.lean

@@ -10,7 +10,7 @@ import all Init.Data.Nat.Bitwise.Basic
 public import Init.Data.Nat.MinMax
 public import Init.Data.Nat.Log2
 import all Init.Data.Nat.Log2
-public import Init.Data.Nat.Power2
+public import Init.Data.Nat.Power2.Basic
 public import Init.Data.Nat.Mod
 import Init.TacticsExtra
 import Init.BinderPredicates

src/Init/Data/Nat/Power2.lean

@@ -6,66 +6,5 @@ Authors: Leonardo de Moura
 module
 
 prelude
-public import Init.Data.Nat.Linear
-
-public section
-
-namespace Nat
-
-theorem nextPowerOfTwo_dec {n power : Nat} (h₁ : power > 0) (h₂ : power < n) : n - power * 2 < n - power := by
-  have : power * 2 = power + power := by simp +arith
-  rw [this, Nat.sub_add_eq]
-  exact Nat.sub_lt (Nat.zero_lt_sub_of_lt h₂) h₁
-
-/--
-Returns the least power of two that's greater than or equal to `n`.
-
-Examples:
- * `Nat.nextPowerOfTwo 0 = 1`
- * `Nat.nextPowerOfTwo 1 = 1`
- * `Nat.nextPowerOfTwo 2 = 2`
- * `Nat.nextPowerOfTwo 3 = 4`
- * `Nat.nextPowerOfTwo 5 = 8`
--/
-def nextPowerOfTwo (n : Nat) : Nat :=
-  go 1 (by decide)
-where
-  go (power : Nat) (h : power > 0) : Nat :=
-    if power < n then
-      go (power * 2) (Nat.mul_pos h (by decide))
-    else
-      power
-  termination_by n - power
-  decreasing_by simp_wf; apply nextPowerOfTwo_dec <;> assumption
-
-/--
-A natural number `n` is a power of two if there exists some `k : Nat` such that `n = 2 ^ k`.
--/
-def isPowerOfTwo (n : Nat) := ∃ k, n = 2 ^ k
-
-theorem isPowerOfTwo_one : isPowerOfTwo 1 :=
-  ⟨0, by decide⟩
-
-theorem isPowerOfTwo_mul_two_of_isPowerOfTwo (h : isPowerOfTwo n) : isPowerOfTwo (n * 2) :=
-  have ⟨k, h⟩ := h
-  ⟨k+1, by simp [h, Nat.pow_succ]⟩
-
-theorem pos_of_isPowerOfTwo (h : isPowerOfTwo n) : n > 0 := by
-  have ⟨k, h⟩ := h
-  rw [h]
-  apply Nat.pow_pos
-  decide
-
-theorem isPowerOfTwo_nextPowerOfTwo (n : Nat) : n.nextPowerOfTwo.isPowerOfTwo := by
-  apply isPowerOfTwo_go
-  apply isPowerOfTwo_one
-where
-  isPowerOfTwo_go (power : Nat) (h₁ : power > 0) (h₂ : power.isPowerOfTwo) : (nextPowerOfTwo.go n power h₁).isPowerOfTwo := by
-    unfold nextPowerOfTwo.go
-    split
-    . exact isPowerOfTwo_go (power*2) (Nat.mul_pos h₁ (by decide)) (Nat.isPowerOfTwo_mul_two_of_isPowerOfTwo h₂)
-    . assumption
-  termination_by n - power
-  decreasing_by simp_wf; apply nextPowerOfTwo_dec <;> assumption
-
-end Nat
+public import Init.Data.Nat.Power2.Basic
+public import Init.Data.Nat.Power2.Lemmas

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

@@ -0,0 +1,71 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+
+prelude
+public import Init.Data.Nat.Linear
+
+public section
+
+namespace Nat
+
+theorem nextPowerOfTwo_dec {n power : Nat} (h₁ : power > 0) (h₂ : power < n) : n - power * 2 < n - power := by
+  have : power * 2 = power + power := by simp +arith
+  rw [this, Nat.sub_add_eq]
+  exact Nat.sub_lt (Nat.zero_lt_sub_of_lt h₂) h₁
+
+/--
+Returns the least power of two that's greater than or equal to `n`.
+
+Examples:
+ * `Nat.nextPowerOfTwo 0 = 1`
+ * `Nat.nextPowerOfTwo 1 = 1`
+ * `Nat.nextPowerOfTwo 2 = 2`
+ * `Nat.nextPowerOfTwo 3 = 4`
+ * `Nat.nextPowerOfTwo 5 = 8`
+-/
+def nextPowerOfTwo (n : Nat) : Nat :=
+  go 1 (by decide)
+where
+  go (power : Nat) (h : power > 0) : Nat :=
+    if power < n then
+      go (power * 2) (Nat.mul_pos h (by decide))
+    else
+      power
+  termination_by n - power
+  decreasing_by simp_wf; apply nextPowerOfTwo_dec <;> assumption
+
+/--
+A natural number `n` is a power of two if there exists some `k : Nat` such that `n = 2 ^ k`.
+-/
+def isPowerOfTwo (n : Nat) := ∃ k, n = 2 ^ k
+
+theorem isPowerOfTwo_one : isPowerOfTwo 1 :=
+  ⟨0, by decide⟩
+
+theorem isPowerOfTwo_mul_two_of_isPowerOfTwo (h : isPowerOfTwo n) : isPowerOfTwo (n * 2) :=
+  have ⟨k, h⟩ := h
+  ⟨k+1, by simp [h, Nat.pow_succ]⟩
+
+theorem pos_of_isPowerOfTwo (h : isPowerOfTwo n) : n > 0 := by
+  have ⟨k, h⟩ := h
+  rw [h]
+  apply Nat.pow_pos
+  decide
+
+theorem isPowerOfTwo_nextPowerOfTwo (n : Nat) : n.nextPowerOfTwo.isPowerOfTwo := by
+  apply isPowerOfTwo_go
+  apply isPowerOfTwo_one
+where
+  isPowerOfTwo_go (power : Nat) (h₁ : power > 0) (h₂ : power.isPowerOfTwo) : (nextPowerOfTwo.go n power h₁).isPowerOfTwo := by
+    unfold nextPowerOfTwo.go
+    split
+    . exact isPowerOfTwo_go (power*2) (Nat.mul_pos h₁ (by decide)) (Nat.isPowerOfTwo_mul_two_of_isPowerOfTwo h₂)
+    . assumption
+  termination_by n - power
+  decreasing_by simp_wf; apply nextPowerOfTwo_dec <;> assumption
+
+end Nat

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

@@ -0,0 +1,62 @@
+/-
+Copyright (c) George Rennie. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: George Rennie
+-/
+module
+
+prelude
+import all Init.Data.Nat.Power2.Basic
+public import Init.Data.Nat.Bitwise.Lemmas
+
+public section
+
+/-!
+# Further lemmas about `Nat.isPowerOfTwo`, with the convenience of having bitwise lemmas available.
+-/
+
+namespace Nat
+
+theorem not_isPowerOfTwo_zero : ¬isPowerOfTwo 0 := by
+  rw [isPowerOfTwo, not_exists]
+  intro x
+  have := one_le_pow x 2 (by decide)
+  omega
+
+theorem and_sub_one_testBit_log2 {n : Nat} (h : n ≠ 0) (hpow2 : ¬n.isPowerOfTwo) :
+    (n &&& (n - 1)).testBit n.log2 := by
+  rw [testBit_and, Bool.and_eq_true]
+  constructor
+  · exact testBit_log2 (by omega)
+  · by_cases n = 2^n.log2
+    · rw [isPowerOfTwo, not_exists] at hpow2
+      have := hpow2 n.log2
+      trivial
+    · have := log2_eq_iff (n := n) (k := n.log2) (by omega)
+      have : (n - 1).log2 = n.log2 := by rw [log2_eq_iff] <;> omega
+      rw [←this]
+      exact testBit_log2 (by omega)
+
+theorem and_sub_one_eq_zero_iff_isPowerOfTwo {n : Nat} (h : n ≠ 0) :
+    (n &&& (n - 1)) = 0 ↔ n.isPowerOfTwo := by
+  constructor
+  · intro hbitwise
+    false_or_by_contra
+    rename_i hpow2
+    have := and_sub_one_testBit_log2 h hpow2
+    rwa [hbitwise, zero_testBit n.log2, Bool.false_eq_true] at this
+  · intro hpow2
+    rcases hpow2 with ⟨_, hpow2⟩
+    rw [hpow2, and_two_pow_sub_one_eq_mod, mod_self]
+
+theorem ne_zero_and_sub_one_eq_zero_iff_isPowerOfTwo {n : Nat} :
+    ((n ≠ 0) ∧ (n &&& (n - 1)) = 0) ↔ n.isPowerOfTwo := by
+  match h : n with
+  | 0 => simp [not_isPowerOfTwo_zero]
+  | n + 1 => simp; exact and_sub_one_eq_zero_iff_isPowerOfTwo (by omega)
+
+@[inline]
+instance {n : Nat} : Decidable n.isPowerOfTwo :=
+  decidable_of_iff _ ne_zero_and_sub_one_eq_zero_iff_isPowerOfTwo
+
+end Nat