diff --git a/Data/Primitive/Array.hs b/Data/Primitive/Array.hs index cfea04d2..5758b7c6 100644 --- a/Data/Primitive/Array.hs +++ b/Data/Primitive/Array.hs @@ -42,7 +42,7 @@ import Data.Primitive.Internal.Compat ( isTrue#, mkNoRepType ) import Control.Monad.ST(ST,runST) import Control.Applicative -import Control.Monad (MonadPlus(..)) +import Control.Monad (MonadPlus(..), when) import Control.Monad.Fix #if MIN_VERSION_base(4,4,0) import Control.Monad.Zip @@ -339,7 +339,9 @@ die fun problem = error $ "Data.Primitive.Array." ++ fun ++ ": " ++ problem instance Eq a => Eq (Array a) where a1 == a2 = sizeofArray a1 == sizeofArray a2 && loop (sizeofArray a1 - 1) where loop i | i < 0 = True - | otherwise = indexArray a1 i == indexArray a2 i && loop (i-1) + | (# x1 #) <- indexArray## a1 i + , (# x2 #) <- indexArray## a2 i + = x1 == x2 && loop (i-1) instance Eq (MutableArray s a) where ma1 == ma2 = isTrue# (sameMutableArray# (marray# ma1) (marray# ma2)) @@ -349,7 +351,10 @@ instance Ord a => Ord (Array a) where where mn = sizeofArray a1 `min` sizeofArray a2 loop i - | i < mn = compare (indexArray a1 i) (indexArray a2 i) `mappend` loop (i+1) + | i < mn + , (# x1 #) <- indexArray## a1 i + , (# x2 #) <- indexArray## a2 i + = compare x1 x2 `mappend` loop (i+1) | otherwise = compare (sizeofArray a1) (sizeofArray a2) instance Foldable Array where @@ -474,9 +479,11 @@ fromList l = fromListN (length l) l instance Functor Array where fmap f a = createArray (sizeofArray a) (die "fmap" "impossible") $ \mb -> - let go i | i < sizeofArray a = return () - | otherwise = writeArray mb i (f $ indexArray a i) - >> go (i+1) + let go i | i == sizeofArray a + = return () + | otherwise + = do x <- indexArrayM a i + writeArray mb i (f x) >> go (i+1) in go 0 #if MIN_VERSION_base(4,8,0) e <$ a = runST $ newArray (sizeofArray a) e >>= unsafeFreezeArray @@ -486,12 +493,15 @@ instance Applicative Array where pure x = runST $ newArray 1 x >>= unsafeFreezeArray ab <*> a = runST $ do mb <- newArray (szab*sza) $ die "<*>" "impossible" - let go1 i - | i < szab = go2 (i*sza) (indexArray ab i) 0 >> go1 (i+1) - | otherwise = return () - go2 off f j - | j < sza = writeArray mb (off + j) (f $ indexArray a j) - | otherwise = return () + let go1 i = when (i < szab) $ + do + f <- indexArrayM ab i + go2 (i*sza) f 0 + go1 (i+1) + go2 off f j = when (j < sza) $ + do + x <- indexArrayM a j + writeArray mb (off + j) (f x) go1 0 unsafeFreezeArray mb where szab = sizeofArray ab ; sza = sizeofArray a @@ -503,7 +513,9 @@ instance Applicative Array where a <* b = createArray (sza*szb) (die "<*" "impossible") $ \ma -> let fill off i e | i < szb = writeArray ma (off+i) e >> fill off (i+1) e | otherwise = return () - go i | i < sza = fill (i*szb) 0 (indexArray a i) >> go (i+1) + go i | i < sza + = do x <- indexArrayM a i + fill (i*szb) 0 x >> go (i+1) | otherwise = return () in go 0 where sza = sizeofArray a ; szb = sizeofArray b @@ -518,20 +530,36 @@ instance Alternative Array where many a | sizeofArray a == 0 = pure [] | otherwise = die "many" "infinite arrays are not well defined" +data ArrayStack a + = PushArray !(Array a) !(ArrayStack a) + | EmptyStack +-- See the note in SmallArray about how we might improve this. + instance Monad Array where return = pure (>>) = (*>) - a >>= f = push 0 [] (sizeofArray a - 1) + + ary >>= f = collect 0 EmptyStack (la-1) where - push !sz bs i - | i < 0 = build sz bs - | otherwise = let b = f $ indexArray a i - in push (sz + sizeofArray b) (b:bs) (i+1) - - build sz stk = createArray sz (die ">>=" "impossible") $ \mb -> - let go off (b:bs) = copyArray mb off b 0 (sizeofArray b) >> go (off + sizeofArray b) bs - go _ [ ] = return () - in go 0 stk + la = sizeofArray ary + collect sz stk i + | i < 0 = createArray sz (die ">>=" "impossible") $ fill 0 stk + | (# x #) <- indexArray## ary i + , let sb = f x + lsb = sizeofArray sb + -- If we don't perform this check, we could end up allocating + -- a stack full of empty arrays if someone is filtering most + -- things out. So we refrain from pushing empty arrays. + = if lsb == 0 + then collect sz stk (i - 1) + else collect (sz + lsb) (PushArray sb stk) (i-1) + + fill _ EmptyStack _ = return () + fill off (PushArray sb sbs) smb + | let lsb = sizeofArray sb + = copyArray smb off sb 0 (lsb) + *> fill (off + lsb) sbs smb + fail _ = empty instance MonadPlus Array where @@ -540,10 +568,13 @@ instance MonadPlus Array where zipW :: String -> (a -> b -> c) -> Array a -> Array b -> Array c zipW s f aa ab = createArray mn (die s "impossible") $ \mc -> - let go i - | i < mn = writeArray mc i (f (indexArray aa i) (indexArray ab i)) - >> go (i+1) - | otherwise = return () + let go i | i < mn + = do + x <- indexArrayM aa i + y <- indexArrayM ab i + writeArray mc i (f x y) + go (i+1) + | otherwise = return () in go 0 where mn = sizeofArray aa `min` sizeofArray ab {-# INLINE zipW #-} @@ -557,7 +588,7 @@ instance MonadZip Array where ma <- newArray sz (die "munzip" "impossible") mb <- newArray sz (die "munzip" "impossible") let go i | i < sz = do - let (a, b) = indexArray aab i + (a, b) <- indexArrayM aab i writeArray ma i a writeArray mb i b go (i+1) @@ -567,7 +598,14 @@ instance MonadZip Array where #endif instance MonadFix Array where - mfix f = let l = mfix (toList . f) in fromListN (length l) l + mfix f = createArray (sizeofArray (f err)) + (die "mfix" "impossible") $ flip fix 0 $ + \r !i !mary -> when (i < sz) $ do + writeArray mary i (fix (\xi -> f xi `indexArray` i)) + r (i + 1) mary + where + sz = sizeofArray (f err) + err = error "mfix for Data.Primitive.Array applied to strict function." #if MIN_VERSION_base(4,9,0) instance Semigroup (Array a) where diff --git a/Data/Primitive/SmallArray.hs b/Data/Primitive/SmallArray.hs index 6a9990e7..b00efb46 100644 --- a/Data/Primitive/SmallArray.hs +++ b/Data/Primitive/SmallArray.hs @@ -75,7 +75,7 @@ import Control.Monad.ST import Control.Monad.Zip #endif import Data.Data -import Data.Foldable +import Data.Foldable as Foldable import Data.Functor.Identity #if !(MIN_VERSION_base(4,11,0)) import Data.Monoid @@ -121,7 +121,7 @@ instance IsList (SmallArray a) where type Item (SmallArray a) = a fromListN n l = SmallArray (fromListN n l) fromList l = SmallArray (fromList l) - toList (SmallArray a) = toList a + toList a = Foldable.toList a #endif #endif @@ -419,19 +419,27 @@ instance Eq a => Eq (SmallArray a) where sa1 == sa2 = length sa1 == length sa2 && loop (length sa1 - 1) where loop i - | i < 0 = True - | otherwise = indexSmallArray sa1 i == indexSmallArray sa2 i && loop (i-1) + | i < 0 + = True + | (# x #) <- indexSmallArray## sa1 i + , (# y #) <- indexSmallArray## sa2 i + = x == y && loop (i-1) instance Eq (SmallMutableArray s a) where SmallMutableArray sma1# == SmallMutableArray sma2# = isTrue# (sameSmallMutableArray# sma1# sma2#) instance Ord a => Ord (SmallArray a) where - compare sl sr = fix ? 0 $ \go i -> - if i < l - then compare (indexSmallArray sl i) (indexSmallArray sr i) <> go (i+1) - else compare (length sl) (length sr) - where l = length sl `min` length sr + compare a1 a2 = loop 0 + where + mn = length a1 `min` length a2 + loop i + | i < mn + , (# x1 #) <- indexSmallArray## a1 i + , (# x2 #) <- indexSmallArray## a2 i + = compare x1 x2 `mappend` loop (i+1) + | otherwise = compare (length a1) (length a2) + instance Foldable SmallArray where -- Note: we perform the array lookups eagerly so we won't @@ -532,8 +540,9 @@ instance Traversable SmallArray where instance Functor SmallArray where fmap f sa = createSmallArray (length sa) (die "fmap" "impossible") $ \smb -> fix ? 0 $ \go i -> - when (i < length sa) $ - writeSmallArray smb i (f $ indexSmallArray sa i) *> go (i+1) + when (i < length sa) $ do + x <- indexSmallArrayM sa i + writeSmallArray smb i (f x) *> go (i+1) {-# INLINE fmap #-} x <$ sa = createSmallArray (length sa) x noOp @@ -548,22 +557,23 @@ instance Applicative SmallArray where where la = length sa ; lb = length sb - sa <* sb = createSmallArray (la*lb) (indexSmallArray sa $ la-1) $ \sma -> - fix ? 0 $ \outer i -> when (i < la-1) $ do - let a = indexSmallArray sa i - fix ? 0 $ \inner j -> - when (j < lb) $ - writeSmallArray sma (la*i + j) a *> inner (j+1) - outer $ i+1 - where - la = length sa ; lb = length sb + a <* b = createSmallArray (sza*szb) (die "<*" "impossible") $ \ma -> + let fill off i e = when (i < szb) $ + writeSmallArray ma (off+i) e >> fill off (i+1) e + go i = when (i < sza) $ do + x <- indexSmallArrayM a i + fill (i*szb) 0 x + go (i+1) + in go 0 + where sza = sizeofSmallArray a ; szb = sizeofSmallArray b sf <*> sx = createSmallArray (lf*lx) (die "<*>" "impossible") $ \smb -> fix ? 0 $ \outer i -> when (i < lf) $ do - let f = indexSmallArray sf i + f <- indexSmallArrayM sf i fix ? 0 $ \inner j -> - when (j < lx) $ - writeSmallArray smb (lf*i + j) (f $ indexSmallArray sx j) + when (j < lx) $ do + x <- indexSmallArrayM sx j + writeSmallArray smb (lf*i + j) (f x) *> inner (j+1) outer $ i+1 where @@ -583,20 +593,41 @@ instance Alternative SmallArray where some sa | null sa = emptySmallArray | otherwise = die "some" "infinite arrays are not well defined" +data ArrayStack a + = PushArray !(SmallArray a) !(ArrayStack a) + | EmptyStack +-- TODO: This isn't terribly efficient. It would be better to wrap +-- ArrayStack with a type like +-- +-- data NES s a = NES !Int !(SmallMutableArray s a) !(ArrayStack a) +-- +-- We'd copy incoming arrays into the mutable array until we would +-- overflow it. Then we'd freeze it, push it on the stack, and continue. +-- Any sufficiently large incoming arrays would go straight on the stack. +-- Such a scheme would make the stack much more compact in the case +-- of many small arrays. + instance Monad SmallArray where return = pure (>>) = (*>) - sa >>= f = collect 0 [] (la-1) + sa >>= f = collect 0 EmptyStack (la-1) where la = length sa collect sz stk i | i < 0 = createSmallArray sz (die ">>=" "impossible") $ fill 0 stk - | otherwise = let sb = f $ indexSmallArray sa i in - collect (sz + length sb) (sb:stk) (i-1) - - fill _ [ ] _ = return () - fill off (sb:sbs) smb = + | (# x #) <- indexSmallArray## sa i + , let sb = f x + lsb = length sb + -- If we don't perform this check, we could end up allocating + -- a stack full of empty arrays if someone is filtering most + -- things out. So we refrain from pushing empty arrays. + = if lsb == 0 + then collect sz stk (i-1) + else collect (sz + lsb) (PushArray sb stk) (i-1) + + fill _ EmptyStack _ = return () + fill off (PushArray sb sbs) smb = copySmallArray smb off sb 0 (length sb) *> fill (off + length sb) sbs smb @@ -609,9 +640,11 @@ instance MonadPlus SmallArray where zipW :: String -> (a -> b -> c) -> SmallArray a -> SmallArray b -> SmallArray c zipW nm = \f sa sb -> let mn = length sa `min` length sb in createSmallArray mn (die nm "impossible") $ \mc -> - fix ? 0 $ \go i -> when (i < mn) $ - writeSmallArray mc i (f (indexSmallArray sa i) (indexSmallArray sb i)) - *> go (i+1) + fix ? 0 $ \go i -> when (i < mn) $ do + x <- indexSmallArrayM sa i + y <- indexSmallArrayM sb i + writeSmallArray mc i (f x y) + go (i+1) {-# INLINE zipW #-} instance MonadZip SmallArray where @@ -631,7 +664,14 @@ instance MonadZip SmallArray where <*> unsafeFreezeSmallArray smb instance MonadFix SmallArray where - mfix f = fromList . mfix $ toList . f + mfix f = createSmallArray (sizeofSmallArray (f err)) + (die "mfix" "impossible") $ flip fix 0 $ + \r !i !mary -> when (i < sz) $ do + writeSmallArray mary i (fix (\xi -> f xi `indexSmallArray` i)) + r (i + 1) mary + where + sz = sizeofSmallArray (f err) + err = error "mfix for Data.Primitive.SmallArray applied to strict function." #if MIN_VERSION_base(4,9,0) instance Sem.Semigroup (SmallArray a) where @@ -658,7 +698,7 @@ instance IsList (SmallArray a) where [] -> pure () x:xs -> writeSmallArray sma i x *> go (i+1) xs fromList l = fromListN (length l) l - toList sa = indexSmallArray sa <$> [0 .. length sa - 1] + toList = Foldable.toList instance Show a => Show (SmallArray a) where showsPrec p sa = showParen (p > 10) $