From 9c6b6982276e469da236d34e2231ffd3cabc2018 Mon Sep 17 00:00:00 2001
From: Rob23oba <152706811+Rob23oba@users.noreply.github.com>
Date: Fri, 22 Aug 2025 13:32:00 +0200
Subject: [PATCH] perf: fast version of `Nat.log2` (#10046)

This PR replaces the implementation of `Nat.log2` with a version that
reduces faster.
The new version can handle:
```lean-4
example : Nat.log2 (1 <<< 500) = 500 := rfl
```
---
 src/Init/Data/Nat/Lemmas.lean |  4 ++--
 src/Init/Data/Nat/Log2.lean   | 35 +++++++++++++++++++++++++++++++----
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/src/Init/Data/Nat/Lemmas.lean b/src/Init/Data/Nat/Lemmas.lean
index 3a3b7519bd..5d2d0f8122 100644
--- a/src/Init/Data/Nat/Lemmas.lean
+++ b/src/Init/Data/Nat/Lemmas.lean
@@ -1313,13 +1313,13 @@ theorem pow_eq_self_iff {a b : Nat} (ha : 1 < a) : a ^ b = a ↔ b = 1 := by
 
 @[simp]
 theorem log2_zero : Nat.log2 0 = 0 := by
-  simp [Nat.log2]
+  simp [Nat.log2_def]
 
 theorem le_log2 (h : n ≠ 0) : k ≤ n.log2 ↔ 2 ^ k ≤ n := by
   match k with
   | 0 => simp [show 1 ≤ n from Nat.pos_of_ne_zero h]
   | k+1 =>
-    rw [log2]; split
+    rw [log2_def]; split
     · have n0 : 0 < n / 2 := (Nat.le_div_iff_mul_le (by decide)).2 ‹_›
       simp only [Nat.add_le_add_iff_right, le_log2 (Nat.ne_of_gt n0),
         Nat.pow_succ]
diff --git a/src/Init/Data/Nat/Log2.lean b/src/Init/Data/Nat/Log2.lean
index a0038099c0..cadf4db9d3 100644
--- a/src/Init/Data/Nat/Log2.lean
+++ b/src/Init/Data/Nat/Log2.lean
@@ -1,7 +1,7 @@
 /-
 Copyright (c) 2021 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Gabriel Ebner
+Authors: Gabriel Ebner, Robin Arnez
 -/
 module
 
@@ -37,11 +37,38 @@ Examples:
 -/
 @[extern "lean_nat_log2"]
 def log2 (n : @& Nat) : Nat :=
-  if n ≥ 2 then log2 (n / 2) + 1 else 0
-decreasing_by exact log2_terminates _ ‹_›
+  -- Lean "assembly"
+  n.rec (fun _ => nat_lit 0) (fun _ ih n =>
+    ((nat_lit 2).ble n).rec (nat_lit 0) ((ih (n.div (nat_lit 2))).succ)) n
+
+private theorem log2_rec_irrel {n k k' : Nat} (hk : n ≤ k) (hk' : n ≤ k') :
+    (k.rec (fun _ => 0) (fun _ ih n => ((2).ble n).rec 0 ((ih (n / 2)).succ)) n : Nat) =
+      k'.rec (fun _ => 0) (fun _ ih n => ((2).ble n).rec 0 ((ih (n / 2)).succ)) n := by
+  induction k generalizing n k' with
+  | zero => cases hk; cases k' <;> rfl
+  | succ k ih =>
+    cases k'
+    · cases hk'; rfl
+    · dsimp only
+      cases h : ble 2 n
+      · rfl
+      · have hn : n / 2 < n := log2_terminates n (Nat.le_of_ble_eq_true h)
+        exact congrArg Nat.succ (ih (Nat.le_of_lt_add_one (Nat.lt_of_lt_of_le hn hk))
+          (Nat.le_of_lt_add_one (Nat.lt_of_lt_of_le hn hk')))
+
+theorem log2_def (n : Nat) : n.log2 = if 2 ≤ n then (n / 2).log2 + 1 else 0 := by
+  rw [Nat.log2, Nat.log2]
+  cases n
+  · rfl
+  split
+  · rename_i n h
+    simp only [ble_eq_true_of_le h]
+    exact congrArg Nat.succ
+      (log2_rec_irrel (Nat.le_of_lt_add_one (log2_terminates _ h)) (Nat.le_refl _))
+  · simp only [mt le_of_ble_eq_true ‹_›]
 
 theorem log2_le_self (n : Nat) : Nat.log2 n ≤ n := by
-  unfold Nat.log2; split
+  rw [log2_def]; split
   next h =>
     have := log2_le_self (n / 2)
     exact Nat.lt_of_le_of_lt this (Nat.div_lt_self (Nat.le_of_lt h) (by decide))