From 504336822fa080b6d214c3dc681b925067000f85 Mon Sep 17 00:00:00 2001
From: Joachim Breitner <mail@joachim-breitner.de>
Date: Wed, 17 Apr 2024 20:11:05 +0200
Subject: [PATCH] perf: faster Nat.repr implementation in C (#3876)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`Nat.repr` was implemented by generating a list of `Chars`, each created
by a 10-way if-then-else. This can cause significant slow down in some
particular use cases.

Now `Nat.repr` is `implemented_by` a faster implementation that uses
C++’s `std::to_string` on small numbers (< USize.size) and maintains an
array of pre-allocated strings for the first 128 numbers.

The handling of big numbers (≥ USize.size) remains as before.
---
 src/Init/Data/Repr.lean   | 19 +++++++++++++++++++
 src/include/lean/lean.h   |  1 +
 src/runtime/object.cpp    |  8 ++++++++
 stage0/src/stdlib_flags.h |  2 +-
 4 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/Init/Data/Repr.lean b/src/Init/Data/Repr.lean
index 42ec4ccecc0b..bfcfbac8e808 100644
--- a/src/Init/Data/Repr.lean
+++ b/src/Init/Data/Repr.lean
@@ -116,6 +116,11 @@ instance {p : α → Prop} [Repr α] : Repr (Subtype p) where
 
 namespace Nat
 
+/-
+We have pure functions for calculating the decimal representation of a `Nat` (`toDigits`), but also
+a fast variant that handles small numbers (`USize`) via C code (`lean_string_of_usize`).
+-/
+
 def digitChar (n : Nat) : Char :=
   if n = 0 then '0' else
   if n = 1 then '1' else
@@ -146,6 +151,20 @@ def toDigitsCore (base : Nat) : Nat → Nat → List Char → List Char
 def toDigits (base : Nat) (n : Nat) : List Char :=
   toDigitsCore base (n+1) n []
 
+@[extern "lean_string_of_usize"]
+protected def _root_.USize.repr (n : @& USize) : String :=
+  (toDigits 10 n.toNat).asString
+
+/-- We statically allocate and memoize reprs for small natural numbers. -/
+private def reprArray : Array String := Id.run do
+  List.range 128 |>.map (·.toUSize.repr) |> Array.mk
+
+private def reprFast (n : Nat) : String :=
+  if h : n < 128 then Nat.reprArray.get ⟨n, h⟩ else
+  if h : n < USize.size then (USize.ofNatCore n h).repr
+  else (toDigits 10 n).asString
+
+@[implemented_by reprFast]
 protected def repr (n : Nat) : String :=
   (toDigits 10 n).asString
 
diff --git a/src/include/lean/lean.h b/src/include/lean/lean.h
index 7bd035b57940..80120a38b89d 100644
--- a/src/include/lean/lean.h
+++ b/src/include/lean/lean.h
@@ -1044,6 +1044,7 @@ LEAN_EXPORT bool lean_string_lt(b_lean_obj_arg s1, b_lean_obj_arg s2);
 static inline uint8_t lean_string_dec_eq(b_lean_obj_arg s1, b_lean_obj_arg s2) { return lean_string_eq(s1, s2); }
 static inline uint8_t lean_string_dec_lt(b_lean_obj_arg s1, b_lean_obj_arg s2) { return lean_string_lt(s1, s2); }
 LEAN_EXPORT uint64_t lean_string_hash(b_lean_obj_arg);
+LEAN_EXPORT lean_obj_res lean_string_of_usize(size_t);
 
 /* Thunks */
 
diff --git a/src/runtime/object.cpp b/src/runtime/object.cpp
index 95dd2758daa8..9fce155fc890 100644
--- a/src/runtime/object.cpp
+++ b/src/runtime/object.cpp
@@ -1629,6 +1629,10 @@ object * mk_string(std::string const & s) {
     return lean_mk_string_from_bytes(s.data(), s.size());
 }
 
+object * mk_ascii_string(std::string const & s) {
+    return lean_mk_string_core(s.data(), s.size(), s.size());
+}
+
 std::string string_to_std(b_obj_arg o) {
     lean_assert(string_size(o) > 0);
     return std::string(w_string_cstr(o), lean_string_size(o) - 1);
@@ -1999,6 +2003,10 @@ extern "C" LEAN_EXPORT uint64 lean_string_hash(b_obj_arg s) {
     return hash_str(sz, (unsigned char const *) str, 11);
 }
 
+extern "C" LEAN_EXPORT obj_res lean_string_of_usize(size_t n) {
+    return mk_ascii_string(std::to_string(n));
+}
+
 // =======================================
 // ByteArray & FloatArray
 
diff --git a/stage0/src/stdlib_flags.h b/stage0/src/stdlib_flags.h
index 0699845ba452..cb8a954a4d16 100644
--- a/stage0/src/stdlib_flags.h
+++ b/stage0/src/stdlib_flags.h
@@ -6,7 +6,7 @@ options get_default_options() {
     // see https://lean-lang.org/lean4/doc/dev/bootstrap.html#further-bootstrapping-complications
 #if LEAN_IS_STAGE0 == 1
     // switch to `true` for ABI-breaking changes affecting meta code
-    opts = opts.update({"interpreter", "prefer_native"}, false);
+    opts = opts.update({"interpreter", "prefer_native"}, true);
     // switch to `true` for changing built-in parsers used in quotations
     opts = opts.update({"internal", "parseQuotWithCurrentStage"}, false);
     // toggling `parseQuotWithCurrentStage` may also require toggling the following option if macros/syntax