gistfile1.txt

{-# OPTIONS_GHC -ddump-simpl -ddump-asm -ddump-to-file -ddump-hi #-}
{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples, ScopedTypeVariables #-}
module Main (main, firstFuuzetsu, firstTreeOwl) where

import GHC.Exts (SmallArray#, newSmallArray#, readSmallArray#, writeSmallArray#,
                 indexSmallArray#, unsafeFreezeSmallArray#, unsafeThawSmallArray#,
                 SmallMutableArray#, sizeofSmallArray#, copySmallArray#, thawSmallArray#,
                 sizeofSmallMutableArray#, copySmallMutableArray#, cloneSmallMutableArray#)
import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#)
import Gauge.Main
import GHC.Int
import Prelude hiding (filter, foldMap, foldr, foldl, length, map, read, traverse)
import Data.Monoid (Sum(..), First(..))
import System.Random (randomIO)
import Data.Coerce (coerce)
import Control.Monad (replicateM)
import Control.DeepSeq
import Control.Monad.ST
import GHC.ST (ST(..))
import Data.IORef

type Array# a = SmallArray# a
type MutableArray# a = SmallMutableArray# a

sizeofArray# :: SmallArray# a -> Int#
sizeofArray# = sizeofSmallArray#

indexArray# :: SmallArray# a -> Int# -> (# a #)
indexArray# = indexSmallArray#

writeArray# :: SmallMutableArray# d a
            -> Int# -> a -> State# d -> State# d
writeArray# = writeSmallArray#

index# :: Array a -> Int -> (# a #)
index# ary _i@(I# i#) = indexArray# (unArray ary) i#
{-# INLINE index# #-}

newArray# :: Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newArray# = newSmallArray#

unsafeFreezeArray# :: SmallMutableArray# d a
                   -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeArray# = unsafeFreezeSmallArray#

run :: (forall s . ST s (MArray s e)) -> Array e
run act = runST $ act >>= unsafeFreeze
{-# INLINE run #-}

-- | Smart constructor
array :: Array# a -> Int -> Array a
array ary _n = Array ary
{-# INLINE array #-}

unsafeFreeze :: MArray s a -> ST s (Array a)
unsafeFreeze mary
    = ST $ \s -> case unsafeFreezeArray# (unMArray mary) s of
                   (# s', ary #) -> (# s', array ary (lengthM mary) #)
{-# INLINE unsafeFreeze #-}

new_ :: Int -> ST s (MArray s a)
new_ n = new n undefinedElem


-- | Create a new mutable array of specified size, in the specified
-- state thread, with each element containing the specified initial
-- value.
new :: Int -> a -> ST s (MArray s a)
new n@(I# n#) b =
    ST $ \s ->
        case newArray# n# b s of
            (# s', ary #) -> (# s', marray ary n #)
{-# INLINE new #-}

undefinedElem :: a
undefinedElem = error "Data.HashMap.Array: Undefined element"
{-# NOINLINE undefinedElem #-}

data MArray s a = MArray {
      unMArray :: !(MutableArray# s a)
    }


sizeofMutableArray# :: SmallMutableArray# s a -> Int#
sizeofMutableArray# = sizeofSmallMutableArray#


lengthM :: MArray s a -> Int
lengthM mary = I# (sizeofMutableArray# (unMArray mary))
{-# INLINE lengthM #-}

-- | Smart constructor
marray :: MutableArray# s a -> Int -> MArray s a
marray mary _n = MArray mary
{-# INLINE marray #-}

length :: Array a -> Int
length ary = I# (sizeofArray# (unArray ary))
{-# INLINE length #-}

write :: MArray s a -> Int -> a -> ST s ()
write ary _i@(I# i#) b = ST $ \ s ->
        case writeArray# (unMArray ary) i# b s of
            s' -> (# s' , () #)
{-# INLINE write #-}

data Array a = Array {
      unArray :: !(Array# a)
    }

instance NFData (Array a) where
    rnf !_ = ()

{-# INLINE foldMapFuuzetsu #-}
foldMapFuuzetsu :: Monoid m => (a -> m) -> Array a -> m
foldMapFuuzetsu f = \ary0 -> case length ary0 of
  0 -> mempty
  len ->
    let !lst = len - 1
        go i | (# x #) <- index# ary0 i, let fx = f x =
          if i == lst then fx else fx <> go (i + 1)
    in go 0

{-# INLINE foldMapTreeOwl #-}
foldMapTreeOwl :: Monoid m => (a -> m) -> Array a -> m
foldMapTreeOwl f = \ary0 ->
  let len = length ary0
  in if len == 0
     then mempty
     else go ary0 (len - 1) 0
    where
      go ary !lst i
        | (# x #) <- index# ary i
        , let fx = f x
        = if i == lst
          then fx
          else fx `mappend` go ary lst (i + 1)



mkAr :: Int -> IO (Array Int)
mkAr l = fromList l <$> replicateM l randomIO
    where
        fromList :: Int -> [a] -> Array a
        fromList n xs0 = run $ do
            mary <- new_ n
            go xs0 mary 0
          where
            go [] !mary !_   = return mary
            go (x:xs) mary i = do write mary i x
                                  go xs mary (i+1)

firstFuuzetsu :: Array Int -> IO Int
firstFuuzetsu ar = do
    ref <- newIORef 0
    foldMapFuuzetsu (\x -> modifyIORef' ref (+ x)) ar
    readIORef ref

firstTreeOwl :: Array Int -> IO Int
firstTreeOwl ar = do
    ref <- newIORef 0
    foldMapTreeOwl (\x -> modifyIORef' ref (+ x)) ar
    readIORef ref

main :: IO ()
main = defaultMain
    [ env (mkAr len) $ \ar -> bgroup (show len)
        [ bench "treeowl" $ nfIO (firstTreeOwl ar)
        , bench "fuuzetsu" $ nfIO (firstFuuzetsu ar)
        ]
    ]
    where
        len = 1000000