fast implementation of digits with base a power of 2 (#31722)

base/intfuncs.jl

@@ -460,6 +460,10 @@ function ndigits0zpb(x::Integer, b::Integer)
         b == 8  && return (sizeof(x)<<3 - leading_zeros(x) + 2) ÷ 3
         b == 16 && return sizeof(x)<<1 - leading_zeros(x)>>2
         b == 10 && return bit_ndigits0z(x)
+        if ispow2(b)
+            dv, rm = divrem(sizeof(x)<<3 - leading_zeros(x), trailing_zeros(b))
+            return iszero(rm) ? dv : dv + 1
+        end
     end
 
     d = 0
@@ -755,9 +759,18 @@ function digits!(a::AbstractVector{T}, n::Integer; base::Integer = 10) where T<:
     isempty(a) && return a
 
     if base > 0
-        for i in eachindex(a)
-            n, d = divrem(n, base)
-            a[i] = d
+        if ispow2(base) && n >= 0 && n isa Base.BitInteger && base <= typemax(Int)
+            base = Int(base)
+            k = trailing_zeros(base)
+            c = base - 1
+            for i in eachindex(a)
+                a[i] = (n >> (k * (i - firstindex(a)))) & c
+            end
+        else
+            for i in eachindex(a)
+                n, d = divrem(n, base)
+                a[i] = d
+            end
         end
     else
         # manually peel one loop iteration for type stability

test/intfuncs.jl

@@ -117,6 +117,14 @@ end
 
         @test ndigits(146, base=-3) == 5
     end
+    @testset "ndigits with base power of 2" begin
+        @test ndigits(17, base = 2) == 5
+        @test ndigits(123, base = 4) == 4
+        @test ndigits(64, base = 8) == 3
+        @test ndigits(8436, base = 16) == 4
+        @test ndigits(159753, base = 32) == 4
+        @test ndigits(3578951, base = 64) == 4
+    end
     let (n, b) = rand(Int, 2)
         -1 <= b <= 1 && (b = 2) # invalid bases
         @test ndigits(n) == ndigits(big(n)) == ndigits(n, base=10)
@@ -181,11 +189,22 @@ end
     @test bitstring(Int128(3)) == "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000011"
 end
 @testset "digits/base" begin
-    @test digits(4, base = 2) == [0, 0, 1]
     @test digits(5, base = 3) == [2, 1]
-    @test digits(5, base = Int32(2), pad=Int32(3)) == [1, 0, 1]
     @test digits(5, pad = 3) == [5, 0, 0]
     @test digits(5, pad = Int32(3)) == [5, 0, 0]
+    # The following have bases powers of 2, but don't enter the fast path
+    @test digits(-3, base = 2) == -[1, 1]
+    @test digits(-42, base = 4) == -[2, 2, 2]
+
+    @testset "digits/base with bases powers of 2" begin
+        @test digits(4, base = 2) == [0, 0, 1]
+        @test digits(5, base = Int32(2), pad=Int32(3)) == [1, 0, 1]
+        @test digits(42, base = 4) == [2, 2, 2]
+        @test digits(321, base = 8) == [1, 0, 5]
+        @test digits(0x123456789abcdef, base = 16) == 15:-1:1
+        @test digits(0x2b1a210a750, base = 64) == [16, 29, 10, 4, 34, 6, 43]
+        @test digits(0x02a01407, base = Int128(1024)) == [7, 5, 42]
+    end
 
     @testset "digits/base with negative bases" begin
         @testset "digits(n::$T, base = b)" for T in (Int, UInt, BigInt, Int32, UInt32)