Benchmarking three versions of inits

inits.hs

{-# LANGUAGE BangPatterns #-}

import Criterion
import Criterion.Main
import Data.List (inits, scanl')
import Data.Primitive.Array
import GHC.Exts (RealWorld)
import System.IO.Unsafe

-- base implementation, using queues
initsQ :: [a] -> [[a]]
initsQ = Data.List.inits

-- using take
initsT :: [a] -> [[a]]
initsT xs = [] : go 1 xs
  where
    go !i [] = []
    go !i (_ : zs) = take i xs : go (i + 1) zs

-- Similar to initsQ (initsQ = map toList . scanl' snoc empty),
-- but the queue is implemented using mutable arrays
-- instead of the banker's queue that uses lists.
-- (unsafePerformIO galore; or linear types FTW (future work))
initsA :: [a] -> [[a]]
initsA xs = unsafePerformIO $ do
  emptyA <- newA
  pure ((map toListA . scanl' snocA emptyA) xs)

-- Queue implemeted using arrays. MUST BE GIVEN TO snoc NO MORE THAN ONCE.
data A a = A
  {-# UNPACK #-} !Int                         -- length n
  {-# UNPACK #-} !(MutableArray RealWorld a)  -- the first n elements are the contents of the queue. Elements can be pushed to the end. YOU MUST NOT PUSH DIFFERENT ELEMENTS TO THE SAME QUEUE (in inits, we only push once).

newA :: IO (A a)
newA = do
  xs <- newArray 1024 undefined
  pure (A 0 xs)

-- Size must be positive
growMA :: MutableArray RealWorld a -> IO (MutableArray RealWorld a)
growMA xs = do
  let n = sizeofMutableArray xs
  ys <- newArray (2 * n) undefined
  copyMutableArray ys 0 xs 0 n
  pure ys

snocA :: A a -> a -> A a
snocA (A n xs) x = unsafePerformIO $ do
  ys <- if n == sizeofMutableArray xs then growMA xs else pure xs
  writeArray ys n x
  pure (A (n + 1) ys)

toListA :: A a -> [a]
toListA (A n xs) = go 0
  where
    go !i | i == n = []
    go i = unsafeDupablePerformIO $ do
      x <- readArray xs i
      pure (x : go (i + 1))

spine :: [a] -> ()
spine = foldr seq ()

spine2 :: [[a]] -> ()
spine2 = foldr (seq . spine) ()

main :: IO ()
main = defaultMain
  [ bench "initsA" $ whnf (spine2 . initsA) [0 .. 10000 :: Int]
  , bench "initsQ" $ whnf (spine2 . inits) [0 .. 10000 :: Int]
  , bench "initsT" $ whnf (spine2 . initsT) [0 .. 10000 :: Int]
  ]