Jacobi facelifting

Math/NumberTheory/Moduli/Jacobi.hs

@@ -1,33 +1,35 @@
 -- |
 -- Module:      Math.NumberTheory.Moduli.Jacobi
--- Copyright:   (c) 2011 Daniel Fischer, 2017 Andrew Lelechenko
+-- Copyright:   (c) 2011 Daniel Fischer, 2017-2018 Andrew Lelechenko
 -- Licence:     MIT
 -- Maintainer:  Andrew Lelechenko <[email protected]>
 -- Stability:   Provisional
 -- Portability: Non-portable (GHC extensions)
 --
--- Jacobi symbol.
+-- <https://en.wikipedia.org/wiki/Jacobi_symbol Jacobi symbol>
+-- is a generalization of the Legendre symbol, useful for primality
+-- testing and integer factorization.
 --
 
-{-# LANGUAGE CPP        #-}
-{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP          #-}
+{-# LANGUAGE LambdaCase   #-}
 
 module Math.NumberTheory.Moduli.Jacobi
   ( JacobiSymbol(..)
   , jacobi
-  , jacobi'
   ) where
 
-import Data.Array.Unboxed
 import Data.Bits
 #if __GLASGOW_HASKELL__ < 803
 import Data.Semigroup
 #endif
+import Numeric.Natural
 
-import Math.NumberTheory.Unsafe
 import Math.NumberTheory.Utils
 
--- | Type for result of 'jacobi'.
+-- | Represents three possible values of
+-- <https://en.wikipedia.org/wiki/Jacobi_symbol Jacobi symbol>.
 data JacobiSymbol = MinusOne | Zero | One
   deriving (Eq, Ord, Show)
 
@@ -47,103 +49,72 @@ negJS = \case
   Zero     -> Zero
   One      -> MinusOne
 
--- | Jacobi symbol of two numbers.
---   The \"denominator\" must be odd and positive, this condition is checked.
+-- | <https://en.wikipedia.org/wiki/Jacobi_symbol Jacobi symbol> of two arguments.
+-- The lower argument (\"denominator\") must be odd and positive,
+-- this condition is checked.
 --
---   If both numbers have a common prime factor, the result
---   is @0@, otherwise it is &#177;1.
+-- If arguments have a common factor, the result
+-- is 'Zero', otherwise it is 'MinusOne' or 'One'.
+--
+-- > > jacobi 1001 9911
+-- > Zero -- arguments have a common factor 11
+-- > > jacobi 1001 9907
+-- > MinusOne
 {-# SPECIALISE jacobi :: Integer -> Integer -> JacobiSymbol,
+                         Natural -> Natural -> JacobiSymbol,
                          Int -> Int -> JacobiSymbol,
                          Word -> Word -> JacobiSymbol
   #-}
 jacobi :: (Integral a, Bits a) => a -> a -> JacobiSymbol
+jacobi _ 1 = One
 jacobi a b
-  | b < 0       = error "Math.NumberTheory.Moduli.jacobi: negative denominator"
-  | evenI b     = error "Math.NumberTheory.Moduli.jacobi: even denominator"
-  | b == 1      = One
-  | otherwise   = jacobi' a b   -- b odd, > 1
-
--- Invariant: b > 1 and odd
--- | Jacobi symbol of two numbers without validity check of
---   the \"denominator\".
-{-# SPECIALISE jacobi' :: Integer -> Integer -> JacobiSymbol,
-                          Int -> Int -> JacobiSymbol,
-                          Word -> Word -> JacobiSymbol
-  #-}
+  | b < 0     = error "Math.NumberTheory.Moduli.jacobi: negative denominator"
+  | evenI b   = error "Math.NumberTheory.Moduli.jacobi: even denominator"
+  | otherwise = jacobi' a b   -- b odd, > 1
+
 jacobi' :: (Integral a, Bits a) => a -> a -> JacobiSymbol
+jacobi' 0 _ = Zero
+jacobi' 1 _ = One
 jacobi' a b
-  | a == 0      = Zero
-  | a == 1      = One
-  | a < 0       = let n | rem4 b == 1 = One
-                        | otherwise   = MinusOne
-                      -- Blech, minBound may pose problems
-                      (z,o) = shiftToOddCount (abs $ toInteger a)
-                      s | evenI z || unsafeAt jac2 (rem8 b) == 1 = n
-                        | otherwise                              = negJS n
-                  in s <> jacobi' (fromInteger o) b
-  | a >= b      = case a `rem` b of
-                    0 -> Zero
-                    r -> jacPS One r b
-  | evenI a     = case shiftToOddCount a of
-                    (z,o) -> let r | rem4 o .&. rem4 b == 1 = One
-                                   | otherwise              = MinusOne
-                                 s | evenI z || unsafeAt jac2 (rem8 b) == 1 = r
-                                   | otherwise                              = negJS r
-                             in jacOL s b o
-  | otherwise   = case rem4 a .&. rem4 b of
-                    3 -> jacOL MinusOne b a
-                    _ -> jacOL One      b a
+  | a < 0     = let n = if rem4is3 b then MinusOne else One
+                    (z, o) = shiftToOddCount (negate a)
+                    s = if evenI z || rem8is1or7 b then n else negJS n
+                in s <> jacobi' o b
+  | a >= b    = case a `rem` b of
+                  0 -> Zero
+                  r -> jacPS One r b
+  | evenI a   = case shiftToOddCount a of
+                  (z, o) -> let r = if rem4is3 o && rem4is3 b then MinusOne else One
+                                s = if evenI z || rem8is1or7 b then r else negJS r
+                            in jacOL s b o
+  | otherwise = jacOL (if rem4is3 a && rem4is3 b then MinusOne else One) b a
 
 -- numerator positive and smaller than denominator
-{-# SPECIALISE jacPS :: JacobiSymbol -> Integer -> Integer -> JacobiSymbol,
-                        JacobiSymbol -> Int -> Int -> JacobiSymbol,
-                        JacobiSymbol -> Word -> Word -> JacobiSymbol
-  #-}
 jacPS :: (Integral a, Bits a) => JacobiSymbol -> a -> a -> JacobiSymbol
-jacPS j a b
-  | evenI a     = case shiftToOddCount a of
-                    (z,o) | evenI z || unsafeAt jac2 (rem8 b) == 1 ->
-                              jacOL (if rem4 o .&. rem4 b == 3 then (negJS j) else j) b o
-                          | otherwise ->
-                              jacOL (if rem4 o .&. rem4 b == 3 then j else (negJS j)) b o
-  | otherwise   = jacOL (if rem4 a .&. rem4 b == 3 then (negJS j) else j) b a
+jacPS !acc a b
+  | evenI a = case shiftToOddCount a of
+    (z, o)
+      | evenI z || rem8is1or7 b -> jacOL (if rem4is3 o && rem4is3 b then negJS acc else acc) b o
+      | otherwise               -> jacOL (if rem4is3 o && rem4is3 b then acc else negJS acc) b o
+  | otherwise = jacOL (if rem4is3 a && rem4is3 b then negJS acc else acc) b a
 
 -- numerator odd, positive and larger than denominator
-{-# SPECIALISE jacOL :: JacobiSymbol -> Integer -> Integer -> JacobiSymbol,
-                        JacobiSymbol -> Int -> Int -> JacobiSymbol,
-                        JacobiSymbol -> Word -> Word -> JacobiSymbol
-  #-}
 jacOL :: (Integral a, Bits a) => JacobiSymbol -> a -> a -> JacobiSymbol
-jacOL j a b
-  | b == 1    = j
-  | otherwise = case a `rem` b of
-                 0 -> Zero
-                 r -> jacPS j r b
+jacOL !acc _ 1 = acc
+jacOL !acc a b = case a `rem` b of
+  0 -> Zero
+  r -> jacPS acc r b
 
 -- Utilities
 
--- For large Integers, going via Int is much faster than bit-fiddling
--- on the Integer, so we do that.
-{-# SPECIALISE evenI :: Integer -> Bool,
-                        Int -> Bool,
-                        Word -> Bool
-  #-}
-evenI :: Integral a => a -> Bool
-evenI n = fromIntegral n .&. 1 == (0 :: Int)
-
-{-# SPECIALISE rem4 :: Integer -> Int,
-                       Int -> Int,
-                       Word -> Int
-  #-}
-rem4 :: Integral a => a -> Int
-rem4 n = fromIntegral n .&. 3
+-- Sadly, GHC do not optimise `Prelude.even` to a bit test automatically.
+evenI :: Bits a => a -> Bool
+evenI n = not (n `testBit` 0)
 
-{-# SPECIALISE rem8 :: Integer -> Int,
-                       Int -> Int,
-                       Word -> Int
-  #-}
-rem8 :: Integral a => a -> Int
-rem8 n = fromIntegral n .&. 7
+-- For an odd input @n@ test whether n `rem` 4 == 1
+rem4is3 :: Bits a => a -> Bool
+rem4is3 n = n `testBit` 1
 
-jac2 :: UArray Int Int
-jac2 = array (0,7) [(0,0),(1,1),(2,0),(3,-1),(4,0),(5,-1),(6,0),(7,1)]
+-- For an odd input @n@ test whether (n `rem` 8) `elem` [1, 7]
+rem8is1or7 :: Bits a => a -> Bool
+rem8is1or7 n = n `testBit` 1 == n `testBit` 2

Math/NumberTheory/Moduli/Sqrt.hs

@@ -43,7 +43,7 @@ import Math.NumberTheory.Utils (shiftToOddCount, splitOff)
 --   the result is @Nothing@.
 sqrtModP :: Integer -> Integer -> Maybe Integer
 sqrtModP n 2 = Just (n `mod` 2)
-sqrtModP n prime = case jacobi' n prime of
+sqrtModP n prime = case jacobi n prime of
                      MinusOne -> Nothing
                      Zero     -> Just 0
                      One      -> Just (sqrtModP' (n `mod` prime) prime)
@@ -215,7 +215,7 @@ findNonSquare n
       where
         primelist = [3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67]
                         ++ sieveFrom (68 + n `rem` 4) -- prevent sharing
-        search (p:ps) = case jacobi' p n of
+        search (p:ps) = case jacobi p n of
           MinusOne -> p
           _        -> search ps
         search _ = error "Should never have happened, prime list exhausted."

Math/NumberTheory/Primes/Testing/Probabilistic.hs

@@ -156,7 +156,7 @@ lucasTest n
       square = isPossibleSquare2 n && r*r == n
       r = integerSquareRoot n
       d = find True 5
-      find !pos cd = case jacobi' (n `rem` cd) cd of
+      find !pos cd = case jacobi (n `rem` cd) cd of
                        MinusOne -> if pos then cd else (-cd)
                        Zero     -> if cd == n then 1 else 0
                        One      -> find (not pos) (cd+2)

Math/NumberTheory/Utils.hs

@@ -29,6 +29,7 @@ import GHC.Base
 
 import GHC.Integer
 import GHC.Integer.GMP.Internals
+import GHC.Natural
 
 import Data.Bits
 
@@ -42,6 +43,7 @@ uncheckedShiftR (W# w#) (I# i#) = W# (uncheckedShiftRL# w# i#)
 "shiftToOddCount/Int"       shiftToOddCount = shiftOCInt
 "shiftToOddCount/Word"      shiftToOddCount = shiftOCWord
 "shiftToOddCount/Integer"   shiftToOddCount = shiftOCInteger
+"shiftToOddCount/Natural"   shiftToOddCount = shiftOCNatural
   #-}
 {-# INLINE [1] shiftToOddCount #-}
 shiftToOddCount :: Integral a => a -> (Int, a)
@@ -75,6 +77,18 @@ shiftOCInteger n@(Jn# bn#) = case bigNatZeroCount bn# of
                                  0#  -> (0, n)
                                  z#  -> (I# z#, n `shiftRInteger` z#)
 
+-- | Specialised version for @'Natural'@.
+--   Precondition: argument nonzero (not checked).
+shiftOCNatural :: Natural -> (Int, Natural)
+shiftOCNatural n@(NatS# i#) =
+    case shiftToOddCount# i# of
+      (# z#, w# #)
+        | isTrue# (z# ==# 0#) -> (0, n)
+        | otherwise -> (I# z#, NatS# w#)
+shiftOCNatural n@(NatJ# bn#) = case bigNatZeroCount bn# of
+                                 0#  -> (0, n)
+                                 z#  -> (I# z#, NatJ# (bn# `shiftRBigNat` z#))
+
 -- | Count trailing zeros in a @'BigNat'@.
 --   Precondition: argument nonzero (not checked, Integer invariant).
 bigNatZeroCount :: BigNat -> Int#

arithmoi.cabal

@@ -180,6 +180,7 @@ benchmark criterion
   other-modules:
     Math.NumberTheory.ArithmeticFunctionsBench
     Math.NumberTheory.GCDBench
+    Math.NumberTheory.JacobiBench
     Math.NumberTheory.MertensBench
     Math.NumberTheory.PowersBench
     Math.NumberTheory.PrimesBench

benchmark/Bench.hs

@@ -4,6 +4,7 @@ import Gauge.Main
 
 import Math.NumberTheory.ArithmeticFunctionsBench as ArithmeticFunctions
 import Math.NumberTheory.GCDBench as GCD
+import Math.NumberTheory.JacobiBench as Jacobi
 import Math.NumberTheory.MertensBench as Mertens
 import Math.NumberTheory.PowersBench as Powers
 import Math.NumberTheory.PrimesBench as Primes
@@ -13,6 +14,7 @@ import Math.NumberTheory.SieveBlockBench as SieveBlock
 main = defaultMain
   [ ArithmeticFunctions.benchSuite
   , GCD.benchSuite
+  , Jacobi.benchSuite
   , Mertens.benchSuite
   , Powers.benchSuite
   , Primes.benchSuite

benchmark/Math/NumberTheory/JacobiBench.hs

@@ -0,0 +1,22 @@
+{-# OPTIONS_GHC -fno-warn-type-defaults #-}
+
+module Math.NumberTheory.JacobiBench
+  ( benchSuite
+  ) where
+
+import Data.Bits
+import Gauge.Main
+import Numeric.Natural
+
+import Math.NumberTheory.Moduli.Jacobi
+
+doBench :: (Integral a, Bits a) => (a -> a -> JacobiSymbol) -> a -> a
+doBench func lim = sum [ x + y | y <- [3, 5 .. lim], x <- [0..y], func x y == One ]
+
+benchSuite :: Benchmark
+benchSuite = bgroup "Jacobi"
+  [ bench "jacobi/Int"      $ nf (doBench jacobi  :: Int -> Int)         2000
+  , bench "jacobi/Word"     $ nf (doBench jacobi  :: Word -> Word)       2000
+  , bench "jacobi/Integer"  $ nf (doBench jacobi  :: Integer -> Integer) 2000
+  , bench "jacobi/Natural"  $ nf (doBench jacobi  :: Natural -> Natural) 2000
+  ]