Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 67 additions & 29 deletions Data/Primitive/Array.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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 #-}
Expand All @@ -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)
Expand All @@ -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
Expand Down
108 changes: 74 additions & 34 deletions Data/Primitive/SmallArray.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Copy link
Member

@RyanGlScott RyanGlScott Mar 20, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to reuse this datatype across both SmallArray.hs and Array.hs? My (entirely untested) hypothesis is that you could turn this into:

data ArrayStack f a
  = PushArray !(f a) !(ArrayStack a)
  | EmptyStack

and put this somewhere in Data.Primitive.Internal.

= 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

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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) $
Expand Down