Quadratic Sieve

-- A record representing the immutable data used during factoring
data FactorData = FD { smoothNum :: Int,
                       segSize :: Integer,
                       factorBase :: [Integer],
                       ceilSqrt :: Integer,
                       quadF :: Integer -> Integer,
                       toSieve :: [(Integer, [Integer])]
                     }

-- Constructs the FactorData record given the number to factor
-- TODO: smoothness bound or smoothness number?
mkFactorData :: Integer -> FactorData
mkFactorData n = FD {
  smoothNum = sn,
  segSize = 100000, -- temp
  factorBase = fb,
  ceilSqrt = cs,
  quadF = \x -> (x + cs) * (x + cs) - n,
  toSieve = map (\p -> (p, yRoots p)) fb
  }
  where
    sn = 150 -- temp
    fb = take sn $ filter (\p -> legendre n p == 1) primes
    cs = ceiling $ (sqrt :: Float -> Float) $ fromIntegral n
    yRoots p = map (\r -> (r - cs) `mod` p) (psqrt n p)

-- Do the sieving with the given prime and roots
sieve' :: STArray s Integer Integer -> (Integer, [Integer]) -> ST s ()
sieve' _ (_,[]) = error "polynomial should have at least one solution"
sieve' a (p,rs) = mapM_ sieve'' rs
  where
    sieve'' r = do
      (lo, hi) <- getBounds a
      -- Get the first index that equals (r mod p)
      let s = lo + mod (r - lo) p in
        -- Sieve with p
        forM_ [s, s + p .. hi]
        (\i -> do
            n <- readArray a i
            writeArray a i (fst $ powsN n p))

-- Sieve the array with each of the (prime,roots) pairs
sieve :: STArray s Integer Integer -> FactorData -> ST s ()
sieve a = mapM_ (sieve' a) . toSieve

-- Fill the array with the initial values of the polynomial
fill :: (Integer -> Integer) -> STArray s Integer Integer -> ST s ()
fill f a = do
  (lo, hi) <- getBounds a
  forM_ [lo .. hi] (\i -> writeArray a i (f i))

-- Loop fusion may handle this, but explicitly writing it out will help
-- me sleep at night
filterMap :: (a -> Bool) -> (a -> b) -> [a] -> [b]
filterMap _ _ [] = []
filterMap p f (x:xs)
  | p x       = f x : filterMap p f xs
  | otherwise = filterMap p f xs

-- Rips the smooth numbers out of the sieved array
rip :: FactorData -> Array Integer Integer -> [(Integer, Integer)]
rip fd = (filterMap ((== 1) . snd) pair) . assocs
  where pair (i,_) = (i + ceilSqrt fd, quadF fd i)

-- Creates an infinite list of smooth numbers
doSieve :: FactorData -> [(Integer, Integer)]
doSieve fd = concatMap (\i -> (rip fd) $ runSTArray $ do
                           a <- newArray_ (i * s, (i + 1) * s)
                           fill (quadF fd) a
                           sieve a fd
                           return a
                       ) [0..]
  where s = segSize fd

-- Factors smooth numbers over the factor base, with exponents mod 2
smoothFactor :: [Integer] -> Integer -> [(Integer, Integer)]
smoothFactor [] _ = []
smoothFactor (p:ps) y = (p, e `mod` 2) : smoothFactor ps d
  where (d,e) = powsN y p

-- Converts a smooth number to its exponent vector over the factor base
toBitVector :: [Integer] -> Integer -> [Int]
toBitVector ps = (map (fromIntegral . snd)) . (smoothFactor ps)

-- External module that has been thoroughly tested
nullSpace2 :: [[Int]] -> [[Int]]
nullSpace2 = allNullVectorsF2 . transpose

-- Finds all (a,b) such that a^2 = b^2 (mod n)
findAB :: FactorData -> [(Integer, Integer)]
findAB fd = map (\vect ->
                  (squareRoot $ product $ zipWith (Prelude.^) smooth vect,
                   product $ zipWith (Prelude.^) xs vect)) choices
  where sieved  = take (1 + smoothNum fd) (doSieve fd)
        fb      = factorBase fd
        xs      = map fst sieved
        smooth  = map snd sieved
        vectors = map (toBitVector fb) smooth
        choices = nullSpace2 vectors

-- Checks to see whether (a-b) or (a+b) share any factors with n
checkAB :: Integer -> (Integer, Integer) -> Maybe Integer
checkAB n (a,b) = if g1 == 1 || g1 == n then
                    if g2 == 1 || g2 == n then Nothing
                    else Just g2
                  else Just g1
  where
    g1 = gcd (a-b) n
    g2 = gcd (a+b) n

-- Merges two sorted lists of factors
merge :: (Num a, Ord a, Num b) => [(a, b)] -> [(a, b)] -> [(a, b)]
merge [] [] = []
merge l1 [] = l1
merge [] l2 = l2
merge l1@((p1,e1):fs1) l2@((p2,e2):fs2)
  | p1 == 1   = merge fs1 l2
  | p2 == 1   = merge l1 fs2
  | p1 == p2  = (p1, e1 + e2) : merge fs1 fs2
  | p1 < p2   = (p1, e1) : merge fs1 l2
  | otherwise = (p2, e2) : merge l1 fs2

unSquare :: Integer -> (Integer, Int)
unSquare n
  | isSquare n = fmap (+1) (unSquare (squareRoot n))
  | otherwise  = (n, 0)

-- The normal quadratic sieve has a serious issue with squares
qsieveSquare :: Integer -> [(Integer, Integer)]
qsieveSquare n = let (q, s) = unSquare n in
  map (fmap (* 2^s)) (qsieve q)

-- | Factors an integer using only the quadratic sieve. You are better off
--   calling factor, which will try simpler methods before resorting to the
--   quadratic sieve
qsieve :: Integer -> [(Integer, Integer)]
qsieve n
  | n <= 1     = [(n,1)]
  | isPrime n  = [(n,1)]
  | isSquare n = qsieveSquare n
  | otherwise  = case mapMaybe (checkAB n) aBs of
    [] -> error "I seriously don't know what to do here"
    (g:_) -> qsieve' g
  where
    aBs = findAB (mkFactorData n)
    qsieve' g = let (q,s) = powsN n g in
      merge (map (fmap (*s)) (qsieve g)) (qsieve q)