Merge pull request #115 from JuliaMath/oscardssmith-faster-Factorization

WIP: increase speed of `Factorization` objects

Author: oscardssmith

src/Primes.jl

@@ -254,12 +254,15 @@ function factor!(n::T, h::AbstractDict{K,Int}) where {T<:Integer,K<:Integer}
 
     local p::T
     for p in PRIMES
+        num_p = 0
         while true
             q, r = divrem(n, T(p)) # T(p) so julia <1.9 uses fast divrem for `BigInt`
             r == 0 || break
-            h[p] = get(h, p, 0) + 1
+            num_p += 1
             n = q
         end
+        # h[p] += num_p (about 2x faster, but the speed only matters for small numbers)
+        num_p > 0 && increment!(h, num_p, p)
         p*p > n && break
     end
     n == 1 && return h

src/factorization.jl

@@ -4,7 +4,12 @@
 struct Factorization{T<:Integer} <: AbstractDict{T, Int}
     pe::Vector{Pair{T, Int}} # Prime-Exponent
 
-    Factorization{T}() where {T<:Integer} = new{T}(Vector{Pair{T, Int}}())
+    function Factorization{T}() where {T<:Integer}
+        # preallocates enough space that numbers smaller than 2310 won't need to resize
+        v = Vector{Pair{T, Int}}(undef, 4)
+        empty!(v)
+        new{T}(v)
+    end
 end
 
 function Factorization{T}(d::AbstractDict) where T<:Integer
@@ -19,23 +24,39 @@ Base.convert(::Type{Factorization}, d::AbstractDict) = Factorization(d)
 Base.iterate(f::Factorization, state...) = iterate(f.pe, state...)
 
 function Base.get(f::Factorization, p, default)
-    found = searchsorted(f.pe, p, by=first)
-    isempty(found) ?
-        default :
-        last(f.pe[first(found)])
+    found = searchsortedfirst(f.pe, p, by=first)
+    (found > length(f.pe) || first(f.pe[found])) != p  ? default : last(f.pe[found])
 end
 
 Base.getindex(f::Factorization, p::Integer) = get(f, p, 0)
 
 function Base.setindex!(f::Factorization{T}, e::Int, p::Integer) where T
-    found = searchsorted(f.pe, p, by=first)
-    if isempty(found)
-        insert!(f.pe, first(found), T(p)=>e)
+    found = searchsortedfirst(f.pe, p, by=first)
+    if found > length(f.pe)
+        push!(f.pe, T(p)=>e)
+    elseif first(f.pe[found]) != p
+        insert!(f.pe, found, T(p)=>e)
+    else
+        f.pe[found] = T(p)=>e
+    end
+    f
+end
+
+"""
+    impliments f[p] += e faster
+"""
+function increment!(f::Factorization{T}, e::Int, p::Integer) where T
+    found = searchsortedfirst(f.pe, p, by=first)
+    if found > length(f.pe)
+        push!(f.pe, T(p)=>e)
+    elseif first(f.pe[found]) != p
+        insert!(f.pe, found, T(p)=>e)
     else
-        f.pe[first(found)] = T(p)=>e
+        f.pe[found] = T(p)=>(last(f.pe[found])+e)
     end
     f
 end
+increment!(f::AbstractDict, e::Int, p::Integer) = (f[p] = get(f, p, 0) + e)
 
 Base.length(f::Factorization) = length(f.pe)