Upgrade toWheel30, backporting from Chimera

Bodigrim · Dec 7, 2019 · 75a384f · 75a384f
1 parent c260ca6
commit 75a384f
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/Math/NumberTheory/Primes/Sieve/Atkin.hs b/Math/NumberTheory/Primes/Sieve/Atkin.hs
@@ -136,7 +136,7 @@ traverseLatticePoints1 !sp vec (!x0, !y0) =
     -- Step 6
     doActions (!k, !y)
       | k < spLength sp
-      = unsafeFlipBit vec (k `shiftL` 4 + toWheel30 (spDelta sp))
+      = unsafeFlipBit vec (k `shiftL` 4 + toWheel30Int (spDelta sp))
         >> doActions (forwardY (k, y))
       | otherwise
       = pure ()
@@ -182,7 +182,7 @@ traverseLatticePoints2 sp vec (x0, y0) =
     -- Step 6
     doActions (!k, !y)
       | k < spLength sp
-      = unsafeFlipBit vec (k `shiftL` 4 + toWheel30 (spDelta sp))
+      = unsafeFlipBit vec (k `shiftL` 4 + toWheel30Int (spDelta sp))
         >> doActions (forwardY (k, y))
       | otherwise
       = pure ()
@@ -212,7 +212,7 @@ traverseLatticePoints3 sp vec (x0, y0) =
     -- Step 6
     doActions (!k, !x, !y)
       | k >= 0 && y < x
-      = unsafeFlipBit vec (k `shiftL` 4 + toWheel30 (spDelta sp))
+      = unsafeFlipBit vec (k `shiftL` 4 + toWheel30Int (spDelta sp))
         >> (let (k', y') = forwardY (k, y) in doActions (k', x, y'))
       | otherwise
       = pure ()

diff --git a/Math/NumberTheory/Utils.hs b/Math/NumberTheory/Utils.hs
@@ -24,6 +24,7 @@ module Math.NumberTheory.Utils
     , recipMod
 
     , toWheel30
+    , toWheel30Int
     , fromWheel30
     ) where
 
@@ -204,13 +205,42 @@ bigNatToNatural bn
 
 -------------------------------------------------------------------------------
 -- Helpers for mapping to rough numbers and back.
--- Copypasted from Data.BitStream.WheelMapping
+-- Copypasted from Data.Chimera.WheelMapping
 
+bits :: Int
+bits = finiteBitSize (0 :: Word)
+
+-- | Left inverse for 'fromWheel30'. Monotonically non-decreasing function.
+--
+-- prop> toWheel30 . fromWheel30 == id
 toWheel30 :: (Integral a, Bits a) => a -> a
 toWheel30 i = q `shiftL` 3 + (r + r `shiftR` 4) `shiftR` 2
   where
     (q, r) = i `P.quotRem` 30
+{-# INLINE toWheel30 #-}
+
+toWheel30Int :: Int -> Int
+toWheel30Int i@(I# i#) = q `shiftL` 3 + (r + r `shiftR` 4) `shiftR` 2
+  where
+    (q, r) = case bits of
+      64 -> (q64, r64)
+      _  -> i `quotRem` 30
 
+    m# = 9838263505978427529## -- (2^67+7) / 15
+    !(# z1#, _ #) = timesWord2# m# (int2Word# i#)
+    q64 = I# (word2Int# z1#) `shiftR` 4
+    r64 = i - q64 `shiftL` 5 + q64 `shiftL` 1
+
+{-# INLINE toWheel30Int #-}
+
+-- | 'fromWheel30' n is the (n+1)-th positive number, not divisible by 2, 3 or 5.
+-- Sequence <https://oeis.org/A007775 A007775>.
+--
+-- prop> map fromWheel30 [0..] == [ n | n <- [0..], n `gcd` 30 == 1 ]
+--
+-- > > map fromWheel30 [0..9]
+-- > [1,7,11,13,17,19,23,29,31,37]
 fromWheel30 :: (Num a, Bits a) => a -> a
 fromWheel30 i = ((i `shiftL` 2 - i `shiftR` 2) .|. 1)
               + ((i `shiftL` 1 - i `shiftR` 1) .&. 2)
+{-# INLINE fromWheel30 #-}