Small fixes and add sieve bench

Author: Jaro Reinders

bench/Bench.hs

@@ -3,9 +3,11 @@ import qualified Array
 import Test.Tasty.Bench
 import qualified Fleet.Array as Fleet
 import Quicksort (quicksort)
-import qualified QuicksortA
+import qualified QuicksortMA
 import qualified QuicksortIM
 import qualified Data.List as List
+import qualified Eratosthenes as Fleet
+import qualified EratosthenesMA as MA
 
 class Indexable a where
   (!) :: a -> Int -> Int
@@ -49,18 +51,27 @@ main = do
     !arr2 = Fleet.set 0 0 arr3
     !arr1 = Fleet.set 0 0 arr2
     !arr0 = Fleet.set 0 0 arr1
-  !marr <- QuicksortA.fromList list
+  !marr <- QuicksortMA.fromList list
   defaultMain
-    [ bench "array" $ whnf fooA (Array.fromList list)
-    , bench "fleet" $ whnf fooDA arr0
-    , bench "fleet 1" $ whnf fooDA arr1
-    , bench "fleet 2" $ whnf fooDA arr2
-    , bench "fleet 5" $ whnf fooDA arr5
-    , bench "fleet 7" $ whnf fooDA arr7
-    , bench "fleet 10" $ whnf fooDA arr10
-    , bench "quicksort array" $ whnfIO (QuicksortA.clone marr >>= \marr' -> QuicksortA.quicksort marr' 0 9973)
-    , bench "quicksort fleet" $ whnf (Quicksort.quicksort 0 9973) (Fleet.copy arr0)
-    , bench "quicksort fleet copy" $ whnf (Quicksort.quicksort 0 9973 . Fleet.copy) arr0
-    , bench "quicksort intmap" $ whnf (QuicksortIM.quicksort 0 9973) (QuicksortIM.fromList list)
-    , bench "sort" $ nf (\xs -> List.sort xs) list
+    [ bgroup "read"
+      [ bench "array" $ whnf fooA (Array.fromList list)
+      , bench "fleet" $ whnf fooDA arr0
+      -- reading old versions:
+      , bench "fleet 1" $ whnf fooDA arr1
+      , bench "fleet 2" $ whnf fooDA arr2
+      , bench "fleet 5" $ whnf fooDA arr5
+      , bench "fleet 7" $ whnf fooDA arr7
+      , bench "fleet 10" $ whnf fooDA arr10
+      ]
+    , bgroup "quicksort"
+      [ bench "MutArr" $ whnfIO (QuicksortMA.clone marr >>= \marr' -> QuicksortMA.quicksort marr' 0 9973)
+      , bench "Fleet" $ whnf (Quicksort.quicksort 0 9973) (Fleet.copy arr0)
+      , bench "Fleet copy" $ whnf (Quicksort.quicksort 0 9973 . Fleet.copy) arr0
+      , bench "IntMap" $ whnf (QuicksortIM.quicksort 0 9973) (QuicksortIM.fromList list)
+      , bench "List.sort" $ nf (\xs -> List.sort xs) list
+      ]
+    , bgroup "sieve"
+      [ bench "fleet" $ nf Fleet.sieve 100000
+      , bench "MutArr" $ nfIO (MA.sieve 100000)
+      ]
     ]

bench/Eratosthenes.hs

@@ -0,0 +1,14 @@
+{-# OPTIONS_GHC -O2 #-}
+module Eratosthenes where
+
+import Fleet.Array as Fleet
+
+sieve :: Int -> [Int]
+sieve n = go 2 (Fleet.replicate (n + 1) True) where
+  go !p !xs
+    | p > n = []
+    | xs ! p = p : go (p + 1) (go' p (p + p) xs)
+    | otherwise = go (p + 1) xs
+  go' !d !i !xs
+    | i > n = xs
+    | otherwise = go' d (i + d) (set i False xs)

bench/EratosthenesMA.hs

@@ -0,0 +1,27 @@
+{-# OPTIONS_GHC -ddump-simpl -ddump-to-file -dsuppress-all -dno-suppress-type-signatures -dno-typeable-binds -ddump-stg-final -ddump-stg-from-core -ddump-prep #-}
+{-# OPTIONS_GHC -O2 #-}
+module EratosthenesMA where
+
+import MutArr as MA
+
+sieve :: Int -> IO [Int]
+sieve n = do
+  arr <- MA.replicate (n + 1) True
+  go arr 2
+  where
+    go :: MutArr Bool -> Int -> IO [Int]
+    go !ma !p
+      | p > n = pure []
+      | otherwise = do
+        isPrime <- readMA ma p
+        if isPrime then do
+          go' ma p (p + p)
+          xs <- go ma (p + 1)
+          pure (p : xs)
+        else
+          go ma (p + 1)
+    go' ma d i
+      | i > n = pure ()
+      | otherwise = do
+        writeMA ma i False
+        go' ma d (i + d)

bench/QuicksortA.hs renamed to bench/MutArr.hs

@@ -1,9 +1,8 @@
 {-# LANGUAGE MagicHash, UnboxedTuples #-}
-{-# OPTIONS_GHC -Wno-name-shadowing -ddump-simpl -ddump-to-file -dsuppress-all -dno-suppress-type-signatures -dno-typeable-binds #-}
-module QuicksortA (MutArr, fromList, toList, clone, quicksort) where
+module MutArr where
 
-import GHC.Exts hiding (fromList, toList)
-import GHC.IO (IO (IO))
+import GHC.Exts
+import GHC.Base
 
 data MutArr a = MA (MutableArray# RealWorld a)
 
@@ -22,6 +21,11 @@ fromList xs = IO $ \s ->
     go arr _ [] s = (# s, MA arr #)
     go arr i (x:xs) s = go arr (i +# 1#) xs (writeArray# arr i x s)
 
+replicate :: Int -> a -> IO (MutArr a)
+replicate (I# n) x = IO $ \s ->
+  case newArray# n x s of
+    (# s', arr #) -> (# s', MA arr #)
+
 toList :: MutArr a -> IO [a]
 toList (MA arr) = IO $ \s ->
   let
@@ -48,26 +52,3 @@ swap !arr !i !j = do
   y <- readMA arr j
   writeMA arr i y
   writeMA arr j x
-
-{-# INLINEABLE quicksort #-}
-quicksort :: Ord a => MutArr a -> Int -> Int -> IO ()
-quicksort !arr !l !r
-  | r - l <= 1 = pure ()
-  | otherwise = do
-    x <- readMA arr (r - 1)
-    m <- partition arr l (r - 1) x
-    swap arr (r - 1) m
-    quicksort arr (m + 1) r
-    quicksort arr l m
-
-{-# INLINEABLE partition #-}
-partition :: Ord a => MutArr a -> Int -> Int -> a -> IO Int
-partition arr l r x = go arr l l where
-  go !arr !m !i
-    | i == r = pure m
-    | otherwise = do
-      y <- readMA arr i
-      if y <= x then do
-        swap arr i m
-        go arr (m + 1) (i + 1)
-      else go arr m (i + 1)

bench/QuicksortMA.hs

@@ -0,0 +1,28 @@
+{-# LANGUAGE MagicHash, UnboxedTuples #-}
+{-# OPTIONS_GHC -Wno-name-shadowing -ddump-simpl -ddump-to-file -dsuppress-all -dno-suppress-type-signatures -dno-typeable-binds #-}
+module QuicksortMA (MutArr, fromList, toList, clone, quicksort) where
+
+import MutArr
+
+{-# INLINEABLE quicksort #-}
+quicksort :: Ord a => MutArr a -> Int -> Int -> IO ()
+quicksort !arr !l !r
+  | r - l <= 1 = pure ()
+  | otherwise = do
+    x <- readMA arr (r - 1)
+    m <- partition arr l (r - 1) x
+    swap arr (r - 1) m
+    quicksort arr (m + 1) r
+    quicksort arr l m
+
+{-# INLINEABLE partition #-}
+partition :: Ord a => MutArr a -> Int -> Int -> a -> IO Int
+partition arr l r x = go arr l l where
+  go !arr !m !i
+    | i == r = pure m
+    | otherwise = do
+      y <- readMA arr i
+      if y <= x then do
+        swap arr i m
+        go arr (m + 1) (i + 1)
+      else go arr m (i + 1)

fleet-array.cabal

@@ -44,7 +44,13 @@ library
 benchmark fleet-array-bench
     import:           warnings
     main-is:          Bench.hs
-    other-modules:    Array, Quicksort, QuicksortA, QuicksortIM
+    other-modules:    Array
+                    , Quicksort
+                    , QuicksortMA
+                    , QuicksortIM
+                    , Eratosthenes
+                    , EratosthenesMA
+                    , MutArr
     build-depends:    base, fleet-array, tasty-bench ^>= {0.4.1}, containers
     hs-source-dirs:   bench/
     default-language: GHC2021

src/Fleet/Array.hs

@@ -24,6 +24,7 @@ latest version.
 module Fleet.Array
   ( Array
   , fromList
+  , replicate
   , toList
   , (!)
   , index
@@ -33,17 +34,23 @@ module Fleet.Array
   , pseq
   ) where
 
+import Prelude hiding (replicate)
+
 import Data.Tuple (Solo (MkSolo))
 import GHC.Exts hiding (fromList, toList, Lifted)
 
 import Data.Kind (Type)
-import GHC.IO.Unsafe (unsafeDupablePerformIO)
 import GHC.Conc (pseq)
+import GHC.Base (IO (IO))
+-- import GHC.IO.Unsafe (unsafeDupablePerformIO)
 
 import Fleet.Array.MutVar
 import Fleet.Array.Lift
 import Fleet.Array.MutArray
 
+unsafeDupablePerformIO :: IO a -> a
+unsafeDupablePerformIO (IO f) = runRW# (\s -> case f s of (# _ , x #) -> x)
+
 data Op a = Set {-# UNPACK #-} !Int a | Swap {-# UNPACK #-} !Int {-# UNPACK #-} !Int
 
 -- | Fleet arrays.
@@ -75,11 +82,18 @@ fromList xs = unsafeDupablePerformIO $ do
   v <- newMutVar (Current arr0)
   pure (A v)
 
+replicate :: Int -> a -> Array a
+replicate n x = unsafeDupablePerformIO $ do
+  arr <- newMutArray n x
+  v <- newMutVar (Current arr)
+  pure (A v)
+
 copyInternal :: ArrayVar a -> IO (MutArray a)
 copyInternal v = do
   av <- readMutVar v
   case av of
     Current arr -> cloneMutArray arr 0 (sizeofMutArray arr)
+    -- _ -> error "Accessing old version"
     Diff op v' -> do
       clone <- copyInternal v'
       appOp clone op
@@ -134,6 +148,7 @@ A v0 ! i0 = unsafeDupablePerformIO (go v0 i0) where
     dat <- readMutVar v
     case dat of
       Current arr -> readMutArray arr i
+      -- _ -> error "Accessing old version"
       Diff (Set j x) v'
         | i == j -> pure x
         | otherwise -> go v' i
@@ -152,6 +167,7 @@ index i0 (A v0) = unsafeDupablePerformIO (go v0 i0) where
     dat <- readMutVar v
     case dat of
       Current arr -> MkSolo <$> readMutArray arr i
+      -- _ -> error "Accessing old version"
       Diff (Set j x) xs
         | i == j -> pure (MkSolo x)
         | otherwise -> go xs i
@@ -182,12 +198,22 @@ reversePointers v = do
   dat <- readMutVar v
   case dat of
     Current arr -> pure arr
-    Diff op v' -> do
-      arr <- reversePointers v'
-      op' <- invert arr op
-      appOp arr op
-      writeMutVar v' (Diff op' v)
-      pure arr
+    Diff op v' -> reversePointersDiff v op v'
+
+-- this needs to be a separate function, because we want the good weather path
+-- (where dat = Current ...) to inline and optimize. The recursion in this
+-- function, which prevents inlining, thus needs to be extracted from
+-- reversePointers.
+reversePointersDiff :: ArrayVar a -> Op a -> ArrayVar a -> IO (MutArray a)
+reversePointersDiff v op v' = do
+  dat <- readMutVar v'
+  arr <- case dat of
+    Current arr -> pure arr
+    Diff op' v'' -> reversePointersDiff v' op' v''
+  op' <- invert arr op
+  appOp arr op
+  writeMutVar v' (Diff op' v)
+  pure arr
 
 {-# INLINE appDiffOp #-}
 appDiffOp :: Op a -> Array a -> Array a