Improve factor some more (#1957)

Now using Nemo.flintify to handle different integer types

src/flint/fmpz_factor.jl

@@ -18,30 +18,47 @@ function _factor(a::ZZRingElem)
 end
 
 # Just to give a helpful error message if someone tries to factor a boolean
-function factor(b::Bool)
-  throw(DomainError("Cannot factorize a boolean"));
+function factor(::Bool)
+  throw(DomainError("Cannot factorize a boolean"))
 end
 
-# This function handles machine integer types (up to 64 bits)
-function factor(a::T) where T <: Union{Int, UInt}
-  iszero(a) && throw(ArgumentError("Argument must be non-zero"))
-  u = sign(a)
-  a = u < 0 ? -a : a
+
+# The main dispatch function: flintify gives either Int or ZZRingElem,
+# then use Julia's dispatcher to select the worker function.
+function factor(a::T)  where T <: Integer
+  @req  !iszero(a)  "Argument must be non-zero"
+  return _factor(typeof(a), Nemo.flintify(a))
+end
+
+
+# Three internal worker functions: one for Int, one for UInt, one for ZZRingElem
+function _factor(::Type{T}, a::Int) where T <: Integer
+  abs_a = reinterpret(UInt, a < 0 ? -a : a) # like abs(a), but correct also when a == typemin(Int)
+  fac = _factor(T, abs_a)
+  if a < 0
+    fac.unit = T(-1)  # OK since fac is mutable; gives error if T is Unsigned
+  end
+  return fac
+end
+
+function _factor(::Type{T}, a::UInt)  where T <: Integer
+  @req  (T != Bool)  "Cannot have a factorization into booleans"
+  @req  !iszero(a)  "Argument must be non-zero"
   F = n_factor()
   @ccall libflint.n_factor(F::Ref{n_factor}, a::UInt)::Nothing
   res = Dict{T, Int}()  # factor-multiplicity pairs
   for i in 1:F.num
-    z = F.p[i]
+    z = T(F.p[i])
     res[z] = F.exp[i]
   end
-  return Fac(u, res)
+  return Fac(T(1), res)
 end
 
-# This is supposed to be called only for T in [Int128, UInt128, BigInt]
-function factor(a::T) where T <: Integer
-  iszero(a) && throw(ArgumentError("Argument must be non-zero"))
-  u = sign(a)
-  F = factor(ZZ(abs(a)))
+function _factor(::Type{T}, a::ZZRingElem)  where T <: Integer
+  @req  (T != Bool)  "Cannot have a factorization into booleans"
+  @req  !iszero(a)  "Argument must be non-zero"
+  u = T(sign(a))
+  F = factor(abs(a))
   res = Dict{T, Int}()  # factor-multiplicity pairs
   for (fac,exp) in F.fac
     res[T(fac)] = exp

test/flint/fmpz_factor-test.jl

@@ -15,8 +15,8 @@
   # trivial case: input is 1
   F1 = factor(1)
   F1_ZZ = factor(ZZ(1))
-  @test length(F1.fac) == 0
-  @test length(F1_ZZ.fac) == 0
+  @test length(F1) == 0
+  @test length(F1_ZZ) == 0
   @test unit(F1) == 1
   @test unit(F1_ZZ) == 1
 
@@ -56,4 +56,17 @@
   @test unit(F_minus1) == -1
   @test unit(F_minus1_ZZ) == -1
 
+  for T in [ Int8,  Int16,  Int32,  Int64,  Int128,  UInt8,  UInt16,  UInt32,  UInt64,  UInt128,  BigInt,  ZZRingElem ]
+    @test_throws ArgumentError factor(T(0))
+    fac99 = factor(T(99))
+    @test  unit(fac99) == 1
+    @test  typeof(unit(fac99)) == T
+    @test  length(fac99) == 2
+    @test  issetequal([a  for (a,_) in fac99], [3, 11])
+    @test  issetequal([e  for (_,e) in fac99], [2, 1])
+    for (a,_) in fac99
+      @test  typeof(a) == T
+    end
+  end
+
 end