Skip to content

Commit

Permalink
perf: faster Nat.repr implementation in C (#3876)
Browse files Browse the repository at this point in the history
`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.
  • Loading branch information
nomeata authored Apr 17, 2024
1 parent 4f50544 commit 5043368
Show file tree
Hide file tree
Showing 4 changed files with 29 additions and 1 deletion.
19 changes: 19 additions & 0 deletions src/Init/Data/Repr.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
1 change: 1 addition & 0 deletions src/include/lean/lean.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 */

Expand Down
8 changes: 8 additions & 0 deletions src/runtime/object.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion stage0/src/stdlib_flags.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 5043368

Please sign in to comment.