diff options
Diffstat (limited to 'third_party/bazel/rules_haskell/examples/primitive/Data/Primitive/SmallArray.hs')
| -rw-r--r-- | third_party/bazel/rules_haskell/examples/primitive/Data/Primitive/SmallArray.hs | 967 |
1 files changed, 967 insertions, 0 deletions
diff --git a/third_party/bazel/rules_haskell/examples/primitive/Data/Primitive/SmallArray.hs b/third_party/bazel/rules_haskell/examples/primitive/Data/Primitive/SmallArray.hs new file mode 100644 index 0000000000..3a50cf2183 --- /dev/null +++ b/third_party/bazel/rules_haskell/examples/primitive/Data/Primitive/SmallArray.hs @@ -0,0 +1,967 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module : Data.Primitive.SmallArray +-- Copyright: (c) 2015 Dan Doel +-- License: BSD3 +-- +-- Maintainer: libraries@haskell.org +-- Portability: non-portable +-- +-- Small arrays are boxed (im)mutable arrays. +-- +-- The underlying structure of the 'Array' type contains a card table, allowing +-- segments of the array to be marked as having been mutated. This allows the +-- garbage collector to only re-traverse segments of the array that have been +-- marked during certain phases, rather than having to traverse the entire +-- array. +-- +-- 'SmallArray' lacks this table. This means that it takes up less memory and +-- has slightly faster writes. It is also more efficient during garbage +-- collection so long as the card table would have a single entry covering the +-- entire array. These advantages make them suitable for use as arrays that are +-- known to be small. +-- +-- The card size is 128, so for uses much larger than that, 'Array' would likely +-- be superior. +-- +-- The underlying type, 'SmallArray#', was introduced in GHC 7.10, so prior to +-- that version, this module simply implements small arrays as 'Array'. + +module Data.Primitive.SmallArray + ( SmallArray(..) + , SmallMutableArray(..) + , newSmallArray + , readSmallArray + , writeSmallArray + , copySmallArray + , copySmallMutableArray + , indexSmallArray + , indexSmallArrayM + , indexSmallArray## + , cloneSmallArray + , cloneSmallMutableArray + , freezeSmallArray + , unsafeFreezeSmallArray + , thawSmallArray + , runSmallArray + , unsafeThawSmallArray + , sizeofSmallArray + , sizeofSmallMutableArray + , smallArrayFromList + , smallArrayFromListN + , mapSmallArray' + , traverseSmallArrayP + ) where + + +#if (__GLASGOW_HASKELL__ >= 710) +#define HAVE_SMALL_ARRAY 1 +#endif + +#if MIN_VERSION_base(4,7,0) +import GHC.Exts hiding (toList) +import qualified GHC.Exts +#endif + +import Control.Applicative +import Control.Monad +import Control.Monad.Fix +import Control.Monad.Primitive +import Control.Monad.ST +import Control.Monad.Zip +import Data.Data +import Data.Foldable as Foldable +import Data.Functor.Identity +#if !(MIN_VERSION_base(4,10,0)) +import Data.Monoid +#endif +#if MIN_VERSION_base(4,9,0) +import qualified GHC.ST as GHCST +import qualified Data.Semigroup as Sem +#endif +import Text.ParserCombinators.ReadP +#if MIN_VERSION_base(4,10,0) +import GHC.Exts (runRW#) +#elif MIN_VERSION_base(4,9,0) +import GHC.Base (runRW#) +#endif + +#if !(HAVE_SMALL_ARRAY) +import Data.Primitive.Array +import Data.Traversable +import qualified Data.Primitive.Array as Array +#endif + +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,4,0) +import Data.Functor.Classes (Eq1(..),Ord1(..),Show1(..),Read1(..)) +#endif + +#if HAVE_SMALL_ARRAY +data SmallArray a = SmallArray (SmallArray# a) + deriving Typeable +#else +newtype SmallArray a = SmallArray (Array a) deriving + ( Eq + , Ord + , Show + , Read + , Foldable + , Traversable + , Functor + , Applicative + , Alternative + , Monad + , MonadPlus + , MonadZip + , MonadFix + , Monoid + , Typeable +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,4,0) + , Eq1 + , Ord1 + , Show1 + , Read1 +#endif + ) + +#if MIN_VERSION_base(4,7,0) +instance IsList (SmallArray a) where + type Item (SmallArray a) = a + fromListN n l = SmallArray (fromListN n l) + fromList l = SmallArray (fromList l) + toList a = Foldable.toList a +#endif +#endif + +#if HAVE_SMALL_ARRAY +data SmallMutableArray s a = SmallMutableArray (SmallMutableArray# s a) + deriving Typeable +#else +newtype SmallMutableArray s a = SmallMutableArray (MutableArray s a) + deriving (Eq, Typeable) +#endif + +-- | Create a new small mutable array. +newSmallArray + :: PrimMonad m + => Int -- ^ size + -> a -- ^ initial contents + -> m (SmallMutableArray (PrimState m) a) +#if HAVE_SMALL_ARRAY +newSmallArray (I# i#) x = primitive $ \s -> + case newSmallArray# i# x s of + (# s', sma# #) -> (# s', SmallMutableArray sma# #) +#else +newSmallArray n e = SmallMutableArray `liftM` newArray n e +#endif +{-# INLINE newSmallArray #-} + +-- | Read the element at a given index in a mutable array. +readSmallArray + :: PrimMonad m + => SmallMutableArray (PrimState m) a -- ^ array + -> Int -- ^ index + -> m a +#if HAVE_SMALL_ARRAY +readSmallArray (SmallMutableArray sma#) (I# i#) = + primitive $ readSmallArray# sma# i# +#else +readSmallArray (SmallMutableArray a) = readArray a +#endif +{-# INLINE readSmallArray #-} + +-- | Write an element at the given idex in a mutable array. +writeSmallArray + :: PrimMonad m + => SmallMutableArray (PrimState m) a -- ^ array + -> Int -- ^ index + -> a -- ^ new element + -> m () +#if HAVE_SMALL_ARRAY +writeSmallArray (SmallMutableArray sma#) (I# i#) x = + primitive_ $ writeSmallArray# sma# i# x +#else +writeSmallArray (SmallMutableArray a) = writeArray a +#endif +{-# INLINE writeSmallArray #-} + +-- | Look up an element in an immutable array. +-- +-- The purpose of returning a result using a monad is to allow the caller to +-- avoid retaining references to the array. Evaluating the return value will +-- cause the array lookup to be performed, even though it may not require the +-- element of the array to be evaluated (which could throw an exception). For +-- instance: +-- +-- > data Box a = Box a +-- > ... +-- > +-- > f sa = case indexSmallArrayM sa 0 of +-- > Box x -> ... +-- +-- 'x' is not a closure that references 'sa' as it would be if we instead +-- wrote: +-- +-- > let x = indexSmallArray sa 0 +-- +-- And does not prevent 'sa' from being garbage collected. +-- +-- Note that 'Identity' is not adequate for this use, as it is a newtype, and +-- cannot be evaluated without evaluating the element. +indexSmallArrayM + :: Monad m + => SmallArray a -- ^ array + -> Int -- ^ index + -> m a +#if HAVE_SMALL_ARRAY +indexSmallArrayM (SmallArray sa#) (I# i#) = + case indexSmallArray# sa# i# of + (# x #) -> pure x +#else +indexSmallArrayM (SmallArray a) = indexArrayM a +#endif +{-# INLINE indexSmallArrayM #-} + +-- | Look up an element in an immutable array. +indexSmallArray + :: SmallArray a -- ^ array + -> Int -- ^ index + -> a +#if HAVE_SMALL_ARRAY +indexSmallArray sa i = runIdentity $ indexSmallArrayM sa i +#else +indexSmallArray (SmallArray a) = indexArray a +#endif +{-# INLINE indexSmallArray #-} + +-- | Read a value from the immutable array at the given index, returning +-- the result in an unboxed unary tuple. This is currently used to implement +-- folds. +indexSmallArray## :: SmallArray a -> Int -> (# a #) +#if HAVE_SMALL_ARRAY +indexSmallArray## (SmallArray ary) (I# i) = indexSmallArray# ary i +#else +indexSmallArray## (SmallArray a) = indexArray## a +#endif +{-# INLINE indexSmallArray## #-} + +-- | Create a copy of a slice of an immutable array. +cloneSmallArray + :: SmallArray a -- ^ source + -> Int -- ^ offset + -> Int -- ^ length + -> SmallArray a +#if HAVE_SMALL_ARRAY +cloneSmallArray (SmallArray sa#) (I# i#) (I# j#) = + SmallArray (cloneSmallArray# sa# i# j#) +#else +cloneSmallArray (SmallArray a) i j = SmallArray $ cloneArray a i j +#endif +{-# INLINE cloneSmallArray #-} + +-- | Create a copy of a slice of a mutable array. +cloneSmallMutableArray + :: PrimMonad m + => SmallMutableArray (PrimState m) a -- ^ source + -> Int -- ^ offset + -> Int -- ^ length + -> m (SmallMutableArray (PrimState m) a) +#if HAVE_SMALL_ARRAY +cloneSmallMutableArray (SmallMutableArray sma#) (I# o#) (I# l#) = + primitive $ \s -> case cloneSmallMutableArray# sma# o# l# s of + (# s', smb# #) -> (# s', SmallMutableArray smb# #) +#else +cloneSmallMutableArray (SmallMutableArray ma) i j = + SmallMutableArray `liftM` cloneMutableArray ma i j +#endif +{-# INLINE cloneSmallMutableArray #-} + +-- | Create an immutable array corresponding to a slice of a mutable array. +-- +-- This operation copies the portion of the array to be frozen. +freezeSmallArray + :: PrimMonad m + => SmallMutableArray (PrimState m) a -- ^ source + -> Int -- ^ offset + -> Int -- ^ length + -> m (SmallArray a) +#if HAVE_SMALL_ARRAY +freezeSmallArray (SmallMutableArray sma#) (I# i#) (I# j#) = + primitive $ \s -> case freezeSmallArray# sma# i# j# s of + (# s', sa# #) -> (# s', SmallArray sa# #) +#else +freezeSmallArray (SmallMutableArray ma) i j = + SmallArray `liftM` freezeArray ma i j +#endif +{-# INLINE freezeSmallArray #-} + +-- | Render a mutable array immutable. +-- +-- This operation performs no copying, so care must be taken not to modify the +-- input array after freezing. +unsafeFreezeSmallArray + :: PrimMonad m => SmallMutableArray (PrimState m) a -> m (SmallArray a) +#if HAVE_SMALL_ARRAY +unsafeFreezeSmallArray (SmallMutableArray sma#) = + primitive $ \s -> case unsafeFreezeSmallArray# sma# s of + (# s', sa# #) -> (# s', SmallArray sa# #) +#else +unsafeFreezeSmallArray (SmallMutableArray ma) = + SmallArray `liftM` unsafeFreezeArray ma +#endif +{-# INLINE unsafeFreezeSmallArray #-} + +-- | Create a mutable array corresponding to a slice of an immutable array. +-- +-- This operation copies the portion of the array to be thawed. +thawSmallArray + :: PrimMonad m + => SmallArray a -- ^ source + -> Int -- ^ offset + -> Int -- ^ length + -> m (SmallMutableArray (PrimState m) a) +#if HAVE_SMALL_ARRAY +thawSmallArray (SmallArray sa#) (I# o#) (I# l#) = + primitive $ \s -> case thawSmallArray# sa# o# l# s of + (# s', sma# #) -> (# s', SmallMutableArray sma# #) +#else +thawSmallArray (SmallArray a) off len = + SmallMutableArray `liftM` thawArray a off len +#endif +{-# INLINE thawSmallArray #-} + +-- | Render an immutable array mutable. +-- +-- This operation performs no copying, so care must be taken with its use. +unsafeThawSmallArray + :: PrimMonad m => SmallArray a -> m (SmallMutableArray (PrimState m) a) +#if HAVE_SMALL_ARRAY +unsafeThawSmallArray (SmallArray sa#) = + primitive $ \s -> case unsafeThawSmallArray# sa# s of + (# s', sma# #) -> (# s', SmallMutableArray sma# #) +#else +unsafeThawSmallArray (SmallArray a) = SmallMutableArray `liftM` unsafeThawArray a +#endif +{-# INLINE unsafeThawSmallArray #-} + +-- | Copy a slice of an immutable array into a mutable array. +copySmallArray + :: PrimMonad m + => SmallMutableArray (PrimState m) a -- ^ destination + -> Int -- ^ destination offset + -> SmallArray a -- ^ source + -> Int -- ^ source offset + -> Int -- ^ length + -> m () +#if HAVE_SMALL_ARRAY +copySmallArray + (SmallMutableArray dst#) (I# do#) (SmallArray src#) (I# so#) (I# l#) = + primitive_ $ copySmallArray# src# so# dst# do# l# +#else +copySmallArray (SmallMutableArray dst) i (SmallArray src) = copyArray dst i src +#endif +{-# INLINE copySmallArray #-} + +-- | Copy a slice of one mutable array into another. +copySmallMutableArray + :: PrimMonad m + => SmallMutableArray (PrimState m) a -- ^ destination + -> Int -- ^ destination offset + -> SmallMutableArray (PrimState m) a -- ^ source + -> Int -- ^ source offset + -> Int -- ^ length + -> m () +#if HAVE_SMALL_ARRAY +copySmallMutableArray + (SmallMutableArray dst#) (I# do#) + (SmallMutableArray src#) (I# so#) + (I# l#) = + primitive_ $ copySmallMutableArray# src# so# dst# do# l# +#else +copySmallMutableArray (SmallMutableArray dst) i (SmallMutableArray src) = + copyMutableArray dst i src +#endif +{-# INLINE copySmallMutableArray #-} + +sizeofSmallArray :: SmallArray a -> Int +#if HAVE_SMALL_ARRAY +sizeofSmallArray (SmallArray sa#) = I# (sizeofSmallArray# sa#) +#else +sizeofSmallArray (SmallArray a) = sizeofArray a +#endif +{-# INLINE sizeofSmallArray #-} + +sizeofSmallMutableArray :: SmallMutableArray s a -> Int +#if HAVE_SMALL_ARRAY +sizeofSmallMutableArray (SmallMutableArray sa#) = + I# (sizeofSmallMutableArray# sa#) +#else +sizeofSmallMutableArray (SmallMutableArray ma) = sizeofMutableArray ma +#endif +{-# INLINE sizeofSmallMutableArray #-} + +-- | This is the fastest, most straightforward way to traverse +-- an array, but it only works correctly with a sufficiently +-- "affine" 'PrimMonad' instance. In particular, it must only produce +-- *one* result array. 'Control.Monad.Trans.List.ListT'-transformed +-- monads, for example, will not work right at all. +traverseSmallArrayP + :: PrimMonad m + => (a -> m b) + -> SmallArray a + -> m (SmallArray b) +#if HAVE_SMALL_ARRAY +traverseSmallArrayP f = \ !ary -> + let + !sz = sizeofSmallArray ary + go !i !mary + | i == sz + = unsafeFreezeSmallArray mary + | otherwise + = do + a <- indexSmallArrayM ary i + b <- f a + writeSmallArray mary i b + go (i + 1) mary + in do + mary <- newSmallArray sz badTraverseValue + go 0 mary +#else +traverseSmallArrayP f (SmallArray ar) = SmallArray `liftM` traverseArrayP f ar +#endif +{-# INLINE traverseSmallArrayP #-} + +-- | Strict map over the elements of the array. +mapSmallArray' :: (a -> b) -> SmallArray a -> SmallArray b +#if HAVE_SMALL_ARRAY +mapSmallArray' f sa = createSmallArray (length sa) (die "mapSmallArray'" "impossible") $ \smb -> + fix ? 0 $ \go i -> + when (i < length sa) $ do + x <- indexSmallArrayM sa i + let !y = f x + writeSmallArray smb i y *> go (i+1) +#else +mapSmallArray' f (SmallArray ar) = SmallArray (mapArray' f ar) +#endif +{-# INLINE mapSmallArray' #-} + +#ifndef HAVE_SMALL_ARRAY +runSmallArray + :: (forall s. ST s (SmallMutableArray s a)) + -> SmallArray a +runSmallArray m = SmallArray $ runArray $ + m >>= \(SmallMutableArray mary) -> return mary + +#elif !MIN_VERSION_base(4,9,0) +runSmallArray + :: (forall s. ST s (SmallMutableArray s a)) + -> SmallArray a +runSmallArray m = runST $ m >>= unsafeFreezeSmallArray + +#else +-- This low-level business is designed to work with GHC's worker-wrapper +-- transformation. A lot of the time, we don't actually need an Array +-- constructor. By putting it on the outside, and being careful about +-- how we special-case the empty array, we can make GHC smarter about this. +-- The only downside is that separately created 0-length arrays won't share +-- their Array constructors, although they'll share their underlying +-- Array#s. +runSmallArray + :: (forall s. ST s (SmallMutableArray s a)) + -> SmallArray a +runSmallArray m = SmallArray (runSmallArray# m) + +runSmallArray# + :: (forall s. ST s (SmallMutableArray s a)) + -> SmallArray# a +runSmallArray# m = case runRW# $ \s -> + case unST m s of { (# s', SmallMutableArray mary# #) -> + unsafeFreezeSmallArray# mary# s'} of (# _, ary# #) -> ary# + +unST :: ST s a -> State# s -> (# State# s, a #) +unST (GHCST.ST f) = f + +#endif + +#if HAVE_SMALL_ARRAY +-- See the comment on runSmallArray for why we use emptySmallArray#. +createSmallArray + :: Int + -> a + -> (forall s. SmallMutableArray s a -> ST s ()) + -> SmallArray a +createSmallArray 0 _ _ = SmallArray (emptySmallArray# (# #)) +createSmallArray n x f = runSmallArray $ do + mary <- newSmallArray n x + f mary + pure mary + +emptySmallArray# :: (# #) -> SmallArray# a +emptySmallArray# _ = case emptySmallArray of SmallArray ar -> ar +{-# NOINLINE emptySmallArray# #-} + +die :: String -> String -> a +die fun problem = error $ "Data.Primitive.SmallArray." ++ fun ++ ": " ++ problem + +emptySmallArray :: SmallArray a +emptySmallArray = + runST $ newSmallArray 0 (die "emptySmallArray" "impossible") + >>= unsafeFreezeSmallArray +{-# NOINLINE emptySmallArray #-} + + +infixl 1 ? +(?) :: (a -> b -> c) -> (b -> a -> c) +(?) = flip +{-# INLINE (?) #-} + +noOp :: a -> ST s () +noOp = const $ pure () + +smallArrayLiftEq :: (a -> b -> Bool) -> SmallArray a -> SmallArray b -> Bool +smallArrayLiftEq p sa1 sa2 = length sa1 == length sa2 && loop (length sa1 - 1) + where + loop i + | i < 0 + = True + | (# x #) <- indexSmallArray## sa1 i + , (# y #) <- indexSmallArray## sa2 i + = p x y && loop (i-1) + +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,4,0) +-- | @since 0.6.4.0 +instance Eq1 SmallArray where +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0) + liftEq = smallArrayLiftEq +#else + eq1 = smallArrayLiftEq (==) +#endif +#endif + +instance Eq a => Eq (SmallArray a) where + sa1 == sa2 = smallArrayLiftEq (==) sa1 sa2 + +instance Eq (SmallMutableArray s a) where + SmallMutableArray sma1# == SmallMutableArray sma2# = + isTrue# (sameSmallMutableArray# sma1# sma2#) + +smallArrayLiftCompare :: (a -> b -> Ordering) -> SmallArray a -> SmallArray b -> Ordering +smallArrayLiftCompare elemCompare a1 a2 = loop 0 + where + mn = length a1 `min` length a2 + loop i + | i < mn + , (# x1 #) <- indexSmallArray## a1 i + , (# x2 #) <- indexSmallArray## a2 i + = elemCompare x1 x2 `mappend` loop (i+1) + | otherwise = compare (length a1) (length a2) + +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,4,0) +-- | @since 0.6.4.0 +instance Ord1 SmallArray where +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0) + liftCompare = smallArrayLiftCompare +#else + compare1 = smallArrayLiftCompare compare +#endif +#endif + +-- | Lexicographic ordering. Subject to change between major versions. +instance Ord a => Ord (SmallArray a) where + compare sa1 sa2 = smallArrayLiftCompare compare sa1 sa2 + +instance Foldable SmallArray where + -- Note: we perform the array lookups eagerly so we won't + -- create thunks to perform lookups even if GHC can't see + -- that the folding function is strict. + foldr f = \z !ary -> + let + !sz = sizeofSmallArray ary + go i + | i == sz = z + | (# x #) <- indexSmallArray## ary i + = f x (go (i+1)) + in go 0 + {-# INLINE foldr #-} + foldl f = \z !ary -> + let + go i + | i < 0 = z + | (# x #) <- indexSmallArray## ary i + = f (go (i-1)) x + in go (sizeofSmallArray ary - 1) + {-# INLINE foldl #-} + foldr1 f = \ !ary -> + let + !sz = sizeofSmallArray ary - 1 + go i = + case indexSmallArray## ary i of + (# x #) | i == sz -> x + | otherwise -> f x (go (i+1)) + in if sz < 0 + then die "foldr1" "Empty SmallArray" + else go 0 + {-# INLINE foldr1 #-} + foldl1 f = \ !ary -> + let + !sz = sizeofSmallArray ary - 1 + go i = + case indexSmallArray## ary i of + (# x #) | i == 0 -> x + | otherwise -> f (go (i - 1)) x + in if sz < 0 + then die "foldl1" "Empty SmallArray" + else go sz + {-# INLINE foldl1 #-} + foldr' f = \z !ary -> + let + go i !acc + | i == -1 = acc + | (# x #) <- indexSmallArray## ary i + = go (i-1) (f x acc) + in go (sizeofSmallArray ary - 1) z + {-# INLINE foldr' #-} + foldl' f = \z !ary -> + let + !sz = sizeofSmallArray ary + go i !acc + | i == sz = acc + | (# x #) <- indexSmallArray## ary i + = go (i+1) (f acc x) + in go 0 z + {-# INLINE foldl' #-} + null a = sizeofSmallArray a == 0 + {-# INLINE null #-} + length = sizeofSmallArray + {-# INLINE length #-} + maximum ary | sz == 0 = die "maximum" "Empty SmallArray" + | (# frst #) <- indexSmallArray## ary 0 + = go 1 frst + where + sz = sizeofSmallArray ary + go i !e + | i == sz = e + | (# x #) <- indexSmallArray## ary i + = go (i+1) (max e x) + {-# INLINE maximum #-} + minimum ary | sz == 0 = die "minimum" "Empty SmallArray" + | (# frst #) <- indexSmallArray## ary 0 + = go 1 frst + where sz = sizeofSmallArray ary + go i !e + | i == sz = e + | (# x #) <- indexSmallArray## ary i + = go (i+1) (min e x) + {-# INLINE minimum #-} + sum = foldl' (+) 0 + {-# INLINE sum #-} + product = foldl' (*) 1 + {-# INLINE product #-} + +newtype STA a = STA {_runSTA :: forall s. SmallMutableArray# s a -> ST s (SmallArray a)} + +runSTA :: Int -> STA a -> SmallArray a +runSTA !sz = \ (STA m) -> runST $ newSmallArray_ sz >>= + \ (SmallMutableArray ar#) -> m ar# +{-# INLINE runSTA #-} + +newSmallArray_ :: Int -> ST s (SmallMutableArray s a) +newSmallArray_ !n = newSmallArray n badTraverseValue + +badTraverseValue :: a +badTraverseValue = die "traverse" "bad indexing" +{-# NOINLINE badTraverseValue #-} + +instance Traversable SmallArray where + traverse f = traverseSmallArray f + {-# INLINE traverse #-} + +traverseSmallArray + :: Applicative f + => (a -> f b) -> SmallArray a -> f (SmallArray b) +traverseSmallArray f = \ !ary -> + let + !len = sizeofSmallArray ary + go !i + | i == len + = pure $ STA $ \mary -> unsafeFreezeSmallArray (SmallMutableArray mary) + | (# x #) <- indexSmallArray## ary i + = liftA2 (\b (STA m) -> STA $ \mary -> + writeSmallArray (SmallMutableArray mary) i b >> m mary) + (f x) (go (i + 1)) + in if len == 0 + then pure emptySmallArray + else runSTA len <$> go 0 +{-# INLINE [1] traverseSmallArray #-} + +{-# RULES +"traverse/ST" forall (f :: a -> ST s b). traverseSmallArray f = traverseSmallArrayP f +"traverse/IO" forall (f :: a -> IO b). traverseSmallArray f = traverseSmallArrayP f +"traverse/Id" forall (f :: a -> Identity b). traverseSmallArray f = + (coerce :: (SmallArray a -> SmallArray (Identity b)) + -> SmallArray a -> Identity (SmallArray b)) (fmap f) + #-} + + +instance Functor SmallArray where + fmap f sa = createSmallArray (length sa) (die "fmap" "impossible") $ \smb -> + fix ? 0 $ \go i -> + 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 + +instance Applicative SmallArray where + pure x = createSmallArray 1 x noOp + + sa *> sb = createSmallArray (la*lb) (die "*>" "impossible") $ \smb -> + fix ? 0 $ \go i -> + when (i < la) $ + copySmallArray smb 0 sb 0 lb *> go (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 + + ab <*> a = createSmallArray (szab*sza) (die "<*>" "impossible") $ \mb -> + let go1 i = when (i < szab) $ + do + f <- indexSmallArrayM ab i + go2 (i*sza) f 0 + go1 (i+1) + go2 off f j = when (j < sza) $ + do + x <- indexSmallArrayM a j + writeSmallArray mb (off + j) (f x) + go2 off f (j + 1) + in go1 0 + where szab = sizeofSmallArray ab ; sza = sizeofSmallArray a + +instance Alternative SmallArray where + empty = emptySmallArray + + sl <|> sr = + createSmallArray (length sl + length sr) (die "<|>" "impossible") $ \sma -> + copySmallArray sma 0 sl 0 (length sl) + *> copySmallArray sma (length sl) sr 0 (length sr) + + many sa | null sa = pure [] + | otherwise = die "many" "infinite arrays are not well defined" + + 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 EmptyStack (la-1) + where + la = length sa + collect sz stk i + | i < 0 = createSmallArray sz (die ">>=" "impossible") $ fill 0 stk + | (# 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 + + fail _ = emptySmallArray + +instance MonadPlus SmallArray where + mzero = empty + mplus = (<|>) + +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) $ do + x <- indexSmallArrayM sa i + y <- indexSmallArrayM sb i + writeSmallArray mc i (f x y) + go (i+1) +{-# INLINE zipW #-} + +instance MonadZip SmallArray where + mzip = zipW "mzip" (,) + mzipWith = zipW "mzipWith" + {-# INLINE mzipWith #-} + munzip sab = runST $ do + let sz = length sab + sma <- newSmallArray sz $ die "munzip" "impossible" + smb <- newSmallArray sz $ die "munzip" "impossible" + fix ? 0 $ \go i -> + when (i < sz) $ case indexSmallArray sab i of + (x, y) -> do writeSmallArray sma i x + writeSmallArray smb i y + go $ i+1 + (,) <$> unsafeFreezeSmallArray sma + <*> unsafeFreezeSmallArray smb + +instance MonadFix SmallArray where + 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) +-- | @since 0.6.3.0 +instance Sem.Semigroup (SmallArray a) where + (<>) = (<|>) + sconcat = mconcat . toList +#endif + +instance Monoid (SmallArray a) where + mempty = empty +#if !(MIN_VERSION_base(4,11,0)) + mappend = (<|>) +#endif + mconcat l = createSmallArray n (die "mconcat" "impossible") $ \ma -> + let go !_ [ ] = return () + go off (a:as) = + copySmallArray ma off a 0 (sizeofSmallArray a) >> go (off + sizeofSmallArray a) as + in go 0 l + where n = sum . fmap length $ l + +instance IsList (SmallArray a) where + type Item (SmallArray a) = a + fromListN = smallArrayFromListN + fromList = smallArrayFromList + toList = Foldable.toList + +smallArrayLiftShowsPrec :: (Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> SmallArray a -> ShowS +smallArrayLiftShowsPrec elemShowsPrec elemListShowsPrec p sa = showParen (p > 10) $ + showString "fromListN " . shows (length sa) . showString " " + . listLiftShowsPrec elemShowsPrec elemListShowsPrec 11 (toList sa) + +-- this need to be included for older ghcs +listLiftShowsPrec :: (Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> [a] -> ShowS +listLiftShowsPrec _ sl _ = sl + +instance Show a => Show (SmallArray a) where + showsPrec p sa = smallArrayLiftShowsPrec showsPrec showList p sa + +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,4,0) +-- | @since 0.6.4.0 +instance Show1 SmallArray where +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0) + liftShowsPrec = smallArrayLiftShowsPrec +#else + showsPrec1 = smallArrayLiftShowsPrec showsPrec showList +#endif +#endif + +smallArrayLiftReadsPrec :: (Int -> ReadS a) -> ReadS [a] -> Int -> ReadS (SmallArray a) +smallArrayLiftReadsPrec _ listReadsPrec p = readParen (p > 10) . readP_to_S $ do + () <$ string "fromListN" + skipSpaces + n <- readS_to_P reads + skipSpaces + l <- readS_to_P listReadsPrec + return $ smallArrayFromListN n l + +instance Read a => Read (SmallArray a) where + readsPrec = smallArrayLiftReadsPrec readsPrec readList + +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,4,0) +-- | @since 0.6.4.0 +instance Read1 SmallArray where +#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0) + liftReadsPrec = smallArrayLiftReadsPrec +#else + readsPrec1 = smallArrayLiftReadsPrec readsPrec readList +#endif +#endif + + + +smallArrayDataType :: DataType +smallArrayDataType = + mkDataType "Data.Primitive.SmallArray.SmallArray" [fromListConstr] + +fromListConstr :: Constr +fromListConstr = mkConstr smallArrayDataType "fromList" [] Prefix + +instance Data a => Data (SmallArray a) where + toConstr _ = fromListConstr + dataTypeOf _ = smallArrayDataType + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> die "gunfold" "SmallArray" + gfoldl f z m = z fromList `f` toList m + +instance (Typeable s, Typeable a) => Data (SmallMutableArray s a) where + toConstr _ = die "toConstr" "SmallMutableArray" + gunfold _ _ = die "gunfold" "SmallMutableArray" + dataTypeOf _ = mkNoRepType "Data.Primitive.SmallArray.SmallMutableArray" +#endif + +-- | Create a 'SmallArray' from a list of a known length. If the length +-- of the list does not match the given length, this throws an exception. +smallArrayFromListN :: Int -> [a] -> SmallArray a +#if HAVE_SMALL_ARRAY +smallArrayFromListN n l = + createSmallArray n + (die "smallArrayFromListN" "uninitialized element") $ \sma -> + let go !ix [] = if ix == n + then return () + else die "smallArrayFromListN" "list length less than specified size" + go !ix (x : xs) = if ix < n + then do + writeSmallArray sma ix x + go (ix+1) xs + else die "smallArrayFromListN" "list length greater than specified size" + in go 0 l +#else +smallArrayFromListN n l = SmallArray (Array.fromListN n l) +#endif + +-- | Create a 'SmallArray' from a list. +smallArrayFromList :: [a] -> SmallArray a +smallArrayFromList l = smallArrayFromListN (length l) l |