perf: Nat.testBit builtin no-alloc implementation

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

@@ -135,7 +135,7 @@ of a number.
 /--
 Returns `true` if the `(n+1)`th least significant bit is `1`, or `false` if it is `0`.
 -/
-@[expose] def testBit (m n : Nat) : Bool :=
+@[expose, extern "lean_nat_test_bit"] def testBit (m n : Nat) : Bool :=
   -- `1 &&& n` is faster than `n &&& 1` for big `n`.
   1 &&& (m >>> n) != 0
 

src/include/lean/lean.h

@@ -1457,6 +1457,7 @@ static inline lean_obj_res lean_nat_lxor(b_lean_obj_arg a1, b_lean_obj_arg a2) {
 
 LEAN_EXPORT lean_obj_res lean_nat_shiftl(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_big_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2);
+LEAN_EXPORT bool lean_nat_big_test_bit(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_pow(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_gcd(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_log2(b_lean_obj_arg a);
@@ -1472,6 +1473,16 @@ static inline lean_obj_res lean_nat_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2)
     }
 }
 
+static inline uint8_t lean_nat_test_bit(b_lean_obj_arg a1, b_lean_obj_arg a2) {
+    if (LEAN_LIKELY(lean_is_scalar(a1) && lean_is_scalar(a2))) {
+        size_t s1 = lean_unbox(a1);
+        size_t s2 = lean_unbox(a2);
+        return ((s2 < sizeof(size_t)*8) ? s1 >> s2 : 0) & 1;
+    } else {
+        return lean_nat_big_test_bit(a1, a2);
+    }
+}
+
 /* Integers */
 
 #define LEAN_MAX_SMALL_INT (sizeof(void*) == 8 ? INT_MAX : (INT_MAX >> 1))

src/runtime/mpz.cpp

@@ -222,6 +222,10 @@ size_t mpz::log2() const {
     return r - 1;
 }
 
+bool mpz::test_bit(size_t n) const {
+    return mpz_tstbit(m_val, n) != 0;
+}
+
 mpz & mpz::operator&=(mpz const & o) {
     mpz_and(m_val, m_val, o.m_val);
     return *this;

src/runtime/mpz.h

@@ -284,6 +284,8 @@ class LEAN_EXPORT mpz {
     */
     size_t log2() const;
 
+    bool test_bit(size_t n) const;
+
     friend void power(mpz & a, mpz const & b, unsigned k);
     friend void _power(mpz & a, mpz const & b, unsigned k) { power(a, b, k); }
     friend mpz pow(mpz a, unsigned k) { power(a, a, k); return a; }

src/runtime/object.cpp

@@ -1509,6 +1509,14 @@ extern "C" LEAN_EXPORT lean_obj_res lean_nat_big_shiftr(b_lean_obj_arg a1, b_lea
     return mpz_to_nat(r);
 }
 
+extern "C" LEAN_EXPORT bool lean_nat_big_test_bit(b_lean_obj_arg a1, b_lean_obj_arg a2) {
+    if (!lean_is_scalar(a2)) {
+        return false; // This large of an exponent must be 0.
+    }
+    // lean_is_scalar(a1) must be false now
+    return mpz_value(a1).test_bit(lean_unbox(a2));
+}
+
 extern "C" LEAN_EXPORT lean_obj_res lean_nat_pow(b_lean_obj_arg a1, b_lean_obj_arg a2) {
     if (!lean_is_scalar(a2) || lean_unbox(a2) > UINT_MAX) {
         lean_internal_panic("Nat.pow exponent is too big");