Replace all remaining usages of Data.Array with Data.Vector

Author: Bodigrim

Math/NumberTheory/Primes/Counting/Impl.hs

@@ -28,8 +28,6 @@ import Math.NumberTheory.Roots
 import Math.NumberTheory.Utils.FromIntegral
 
 import Control.Monad.ST
-import Data.Array.Base
-import Data.Array.ST
 import Data.Bit
 import Data.Bits
 import Data.Int
@@ -269,8 +267,8 @@ treat end n old new = do
                 !val <- MU.unsafeRead old (qi+1)
                 let !q0 = key `quot` n
                     !r0 = int64ToInt (q0 `rem` 30030)
-                    !nkey = q0 - int8ToInt64 (cpDfAr `unsafeAt` r0)
-                    nk0 = q0 + int8ToInt64 (cpGpAr `unsafeAt` (r0+1) + 1)
+                    !nkey = q0 - int8ToInt64 (cpDfAr `U.unsafeIndex` r0)
+                    nk0 = q0 + int8ToInt64 (cpGpAr `U.unsafeIndex` (r0+1) + 1)
                     !nlim = n*nk0
                 (wi1,ci1) <- copyTo end nkey old ci new wi
                 ckey <- MU.unsafeRead old ci1
@@ -332,27 +330,27 @@ cp6 :: Int64 -> Integer
 cp6 k =
   case k `quotRem` 30030 of
     (q,r) -> 5760*int64ToInteger q +
-                int16ToInteger (cpCtAr `unsafeAt` int64ToInt r)
+                int16ToInteger (cpCtAr `U.unsafeIndex` int64ToInt r)
 
 cop :: Int64 -> Int64
-cop m = m - int8ToInt64 (cpDfAr `unsafeAt` int64ToInt (m `rem` 30030))
+cop m = m - int8ToInt64 (cpDfAr `U.unsafeIndex` int64ToInt (m `rem` 30030))
 
 
 --------------------------------------------------------------------------------
 --                           Ugly helper arrays                               --
 --------------------------------------------------------------------------------
 
-cpCtAr :: UArray Int Int16
-cpCtAr = runSTUArray $ do
-    ar <- newArray (0,30029) 1
+cpCtAr :: U.Vector Int16
+cpCtAr = runST $ do
+    ar <- MU.replicate 30030 1
     let zilch s i
-            | i < 30030 = unsafeWrite ar i 0 >> zilch s (i+s)
+            | i < 30030 = MU.unsafeWrite ar i 0 >> zilch s (i+s)
             | otherwise = return ()
         accumulate ct i
             | i < 30030 = do
-                v <- unsafeRead ar i
+                v <- MU.unsafeRead ar i
                 let !ct' = ct+v
-                unsafeWrite ar i ct'
+                MU.unsafeWrite ar i ct'
                 accumulate ct' (i+1)
             | otherwise = return ar
     zilch 2 0
@@ -361,20 +359,21 @@ cpCtAr = runSTUArray $ do
     zilch 14 7
     zilch 22 11
     zilch 26 13
-    accumulate 1 2
+    vec <- accumulate 1 2
+    U.unsafeFreeze vec
 
-cpDfAr :: UArray Int Int8
-cpDfAr = runSTUArray $ do
-    ar <- newArray (0,30029) 0
+cpDfAr :: U.Vector Int8
+cpDfAr = runST $ do
+    ar <- MU.replicate 30030 0
     let note s i
-            | i < 30029 = unsafeWrite ar i 1 >> note s (i+s)
+            | i < 30029 = MU.unsafeWrite ar i 1 >> note s (i+s)
             | otherwise = return ()
         accumulate d i
             | i < 30029 = do
-                v <- unsafeRead ar i
+                v <- MU.unsafeRead ar i
                 if v == 0
                     then accumulate 2 (i+2)
-                    else do unsafeWrite ar i d
+                    else do MU.unsafeWrite ar i d
                             accumulate (d+1) (i+1)
             | otherwise = return ar
     note 2 0
@@ -383,30 +382,32 @@ cpDfAr = runSTUArray $ do
     note 14 7
     note 22 11
     note 26 13
-    accumulate 2 3
+    vec <- accumulate 2 3
+    U.unsafeFreeze vec
 
-cpGpAr :: UArray Int Int8
-cpGpAr = runSTUArray $ do
-    ar <- newArray (0,30030) 0
-    unsafeWrite ar 30030 1
+cpGpAr :: U.Vector Int8
+cpGpAr = runST $ do
+    ar <- MU.replicate 30031 0
+    MU.unsafeWrite ar 30030 1
     let note s i
-            | i < 30029 = unsafeWrite ar i 1 >> note s (i+s)
+            | i < 30029 = MU.unsafeWrite ar i 1 >> note s (i+s)
             | otherwise = return ()
         accumulate d i
             | i < 1     = return ar
             | otherwise = do
-                v <- unsafeRead ar i
+                v <- MU.unsafeRead ar i
                 if v == 0
                     then accumulate 2 (i-2)
-                    else do unsafeWrite ar i d
+                    else do MU.unsafeWrite ar i d
                             accumulate (d+1) (i-1)
     note 2 0
     note 6 3
     note 10 5
     note 14 7
     note 22 11
     note 26 13
-    accumulate 2 30027
+    vec <- accumulate 2 30027
+    U.unsafeFreeze vec
 
 -------------------------------------------------------------------------------
 -- Prime counting

Math/NumberTheory/Primes/Sieve/Eratosthenes.hs

@@ -26,7 +26,6 @@ module Math.NumberTheory.Primes.Sieve.Eratosthenes
 
 import Control.Monad (when)
 import Control.Monad.ST
-import Data.Array.Base
 import Data.Bit
 import Data.Bits
 import Data.Coerce
@@ -374,17 +373,17 @@ psieveFrom n = makeSieves plim sqlim bitOff valOff cache
 
 {-# INLINE delta #-}
 delta :: Int -> Int
-delta = unsafeAt deltas
+delta = U.unsafeIndex deltas
 
-deltas :: UArray Int Int
-deltas = listArray (0,7) [4,2,4,2,4,6,2,6]
+deltas :: U.Vector Int
+deltas = U.fromList [4,2,4,2,4,6,2,6]
 
 {-# INLINE tau #-}
 tau :: Int -> Int
-tau = unsafeAt taus
+tau = U.unsafeIndex taus
 
-taus :: UArray Int Int
-taus = listArray (0,63)
+taus :: U.Vector Int
+taus = U.fromList
         [  7,  4,  7,  4,  7, 12,  3, 12
         , 12,  6, 11,  6, 12, 18,  5, 18
         , 14,  7, 13,  7, 14, 21,  7, 21
@@ -397,28 +396,28 @@ taus = listArray (0,63)
 
 {-# INLINE byte #-}
 byte :: Int -> Int
-byte = unsafeAt startByte
+byte = U.unsafeIndex startByte
 
-startByte :: UArray Int Int
-startByte = listArray (0,7) [1,3,5,9,11,17,27,31]
+startByte :: U.Vector Int
+startByte = U.fromList [1,3,5,9,11,17,27,31]
 
 {-# INLINE idx #-}
 idx :: Int -> Int
-idx = unsafeAt startIdx
+idx = U.unsafeIndex startIdx
 
-startIdx :: UArray Int Int
-startIdx = listArray (0,7) [4,7,4,4,7,4,7,7]
+startIdx :: U.Vector Int
+startIdx = U.fromList [4,7,4,4,7,4,7,7]
 
 {-# INLINE mu #-}
 mu :: Int -> Int
-mu = unsafeAt mArr
+mu = U.unsafeIndex mArr
 
 {-# INLINE nu #-}
 nu :: Int -> Int
-nu = unsafeAt nArr
+nu = U.unsafeIndex nArr
 
-mArr :: UArray Int Int
-mArr = listArray (0,63)
+mArr :: U.Vector Int
+mArr = U.fromList
         [ 1,  2,  2,  3,  4,  5,  6,  7
         , 2,  3,  4,  6,  6,  8, 10, 11
         , 2,  4,  5,  7,  8,  9, 12, 13
@@ -429,8 +428,8 @@ mArr = listArray (0,63)
         , 7, 11, 13, 17, 19, 23, 29, 31
         ]
 
-nArr :: UArray Int Int
-nArr = listArray (0,63)
+nArr :: U.Vector Int
+nArr = U.fromList
         [ 4, 3, 7, 6, 2, 1, 5, 0
         , 3, 7, 5, 0, 6, 2, 4, 1
         , 7, 5, 4, 1, 0, 6, 3, 2

Math/NumberTheory/Primes/Sieve/Indexing.hs

@@ -12,8 +12,8 @@ module Math.NumberTheory.Primes.Sieve.Indexing
     , rho
     ) where
 
-import Data.Array.Base
 import Data.Bits
+import qualified Data.Vector.Unboxed as U
 
 {-# INLINE idxPr #-}
 idxPr :: Integral a => a -> (Int,Int)
@@ -37,7 +37,7 @@ toPrim ix = 30*fromIntegral k + fromIntegral (rho i)
 
 {-# INLINE rho #-}
 rho :: Int -> Int
-rho = unsafeAt residues
+rho = U.unsafeIndex residues
 
-residues :: UArray Int Int
-residues = listArray (0,7) [7,11,13,17,19,23,29,31]
+residues :: U.Vector Int
+residues = U.fromList [7,11,13,17,19,23,29,31]

arithmoi.cabal

@@ -28,7 +28,6 @@ source-repository head
 library
   build-depends:
     base >=4.15 && <5,
-    array >=0.5 && <0.6,
     bitvec >=1.1 && <1.2,
     containers >=0.5.11 && <0.9,
     chimera >=0.3 && <0.5,
@@ -177,7 +176,6 @@ benchmark arithmoi-bench
   build-depends:
     base,
     arithmoi,
-    array,
     constraints,
     containers,
     deepseq,

benchmark/Math/NumberTheory/SequenceBench.hs

@@ -6,8 +6,8 @@ module Math.NumberTheory.SequenceBench
 
 import Test.Tasty.Bench
 
-import Data.Array.Unboxed
 import Data.Bits
+import qualified Data.Vector.Unboxed as U
 
 import Math.NumberTheory.Primes (Prime(..), nextPrime, precPrime)
 import Math.NumberTheory.Primes.Testing
@@ -43,10 +43,10 @@ benchSuite = bgroup "Sequence"
 -- Utils copypasted from internal modules
 
 rho :: Int -> Int
-rho i = residues ! i
+rho i = residues U.! i
 
-residues :: UArray Int Int
-residues = listArray (0,7) [7,11,13,17,19,23,29,31]
+residues :: U.Vector Int
+residues = U.fromList [7,11,13,17,19,23,29,31]
 
 toIdx :: Integral a => a -> Int
 toIdx n = 8*fromIntegral q+r2