diff --git a/src/include/lean/lean.h b/src/include/lean/lean.h index fceaee2d38..3c2495ce04 100644 --- a/src/include/lean/lean.h +++ b/src/include/lean/lean.h @@ -1424,11 +1424,22 @@ 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_shiftr(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 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); +static inline lean_obj_res lean_nat_shiftr(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); + size_t r = (s2 < sizeof(size_t)*8) ? s1 >> s2 : 0; + return lean_box(r); + } else { + return lean_nat_big_shiftr(a1, a2); + } +} + /* Integers */ #define LEAN_MAX_SMALL_INT (sizeof(void*) == 8 ? INT_MAX : (INT_MAX >> 1)) diff --git a/src/runtime/object.cpp b/src/runtime/object.cpp index 7c390dd679..ab12a58c31 100644 --- a/src/runtime/object.cpp +++ b/src/runtime/object.cpp @@ -1462,7 +1462,7 @@ extern "C" LEAN_EXPORT lean_obj_res lean_nat_shiftl(b_lean_obj_arg a1, b_lean_ob return mpz_to_nat(r); } -extern "C" LEAN_EXPORT lean_obj_res lean_nat_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2) { +extern "C" LEAN_EXPORT lean_obj_res lean_nat_big_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2) { if (!lean_is_scalar(a2)) { return lean_box(0); // This large of an exponent must be 0. } diff --git a/tests/compiler/nat_shiftr.lean b/tests/compiler/nat_shiftr.lean new file mode 100644 index 0000000000..dd6d22d6ac --- /dev/null +++ b/tests/compiler/nat_shiftr.lean @@ -0,0 +1,19 @@ +def test (a : Nat) : IO Unit := + for b in #[0, 1, 14, 15, 16, 17, 31, 32, 33, 63, 64, 65] do + IO.println f!"{a >>> b}" + +def main : IO Unit := do + test 0 + test 1 + test 0xff + test 0x100 + test 0x101 + test 0xffff + test 0x1000_0 + test 0x1000_1 + test 0xffff_ffff + test 0x1_0000_0000 + test 0x1_0000_0001 + test 0xffff_ffff_ffff_ffff + test 0x1_0000_0000_0000_0000 + test 0x1_0000_0000_0000_0001 diff --git a/tests/compiler/nat_shiftr.lean.expected.out b/tests/compiler/nat_shiftr.lean.expected.out new file mode 100644 index 0000000000..01fed06ccc --- /dev/null +++ b/tests/compiler/nat_shiftr.lean.expected.out @@ -0,0 +1,168 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +1 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +255 +127 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +256 +128 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +257 +128 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +65535 +32767 +3 +1 +0 +0 +0 +0 +0 +0 +0 +0 +65536 +32768 +4 +2 +1 +0 +0 +0 +0 +0 +0 +0 +65537 +32768 +4 +2 +1 +0 +0 +0 +0 +0 +0 +0 +4294967295 +2147483647 +262143 +131071 +65535 +32767 +1 +0 +0 +0 +0 +0 +4294967296 +2147483648 +262144 +131072 +65536 +32768 +2 +1 +0 +0 +0 +0 +4294967297 +2147483648 +262144 +131072 +65536 +32768 +2 +1 +0 +0 +0 +0 +18446744073709551615 +9223372036854775807 +1125899906842623 +562949953421311 +281474976710655 +140737488355327 +8589934591 +4294967295 +2147483647 +1 +0 +0 +18446744073709551616 +9223372036854775808 +1125899906842624 +562949953421312 +281474976710656 +140737488355328 +8589934592 +4294967296 +2147483648 +2 +1 +0 +18446744073709551617 +9223372036854775808 +1125899906842624 +562949953421312 +281474976710656 +140737488355328 +8589934592 +4294967296 +2147483648 +2 +1 +0