@@ -55,18 +55,22 @@ module Streamly.Internal.Data.Array.Mut.Type
5555 , writeNUnsafe
5656 , writeNUnsafeAs
5757 , writeN
58+ , pinnedWriteN
5859 , writeNAs
5960 , pinnedWriteNAligned
6061
6162 , writeWith
6263 , write
64+ , pinnedWrite
6365
6466 , writeRevN
6567 -- , writeRev
6668
6769 -- ** From containers
6870 , fromListN
71+ , pinnedFromListN
6972 , fromList
73+ , pinnedFromList
7074 , fromListRevN
7175 , fromListRev
7276 , fromStreamDN
@@ -1357,15 +1361,20 @@ chunksOf :: forall m a. (MonadIO m, Unbox a)
13571361-- chunksOf n = D.foldMany (writeN n)
13581362chunksOf = chunksOfAs Unpinned
13591363
1364+{-# INLINE arrayStreamKFromStreamDAs #-}
1365+arrayStreamKFromStreamDAs :: forall m a. (MonadIO m, Unbox a) =>
1366+ PinnedState -> D.Stream m a -> m (StreamK m (MutArray a))
1367+arrayStreamKFromStreamDAs ps =
1368+ let n = allocBytesToElemCount (undefined :: a) defaultChunkSize
1369+ in D.foldr K.cons K.nil . chunksOfAs ps n
1370+
13601371-- XXX This should take a PinnedState
13611372-- XXX buffer to a list instead?
13621373-- | Buffer the stream into arrays in memory.
13631374{-# INLINE arrayStreamKFromStreamD #-}
13641375arrayStreamKFromStreamD :: forall m a. (MonadIO m, Unbox a) =>
13651376 D.Stream m a -> m (StreamK m (MutArray a))
1366−arrayStreamKFromStreamD =
1367− let n = allocBytesToElemCount (undefined :: a) defaultChunkSize
1368− in D.foldr K.cons K.nil . chunksOf n
1377+arrayStreamKFromStreamD = arrayStreamKFromStreamDAs Pinned
13691378
13701379-------------------------------------------------------------------------------
13711380-- Streams of arrays - Flattening
@@ -1807,6 +1816,13 @@ writeNAs ps = writeNWith (newAs ps)
18071816writeN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (MutArray a)
18081817writeN = writeNAs Unpinned
18091818
1819+{-# INLINE_NORMAL pinnedWriteN #-}
1820+pinnedWriteN ::
1821+ forall m a. (MonadIO m, Unbox a)
1822+ => Int
1823+ -> Fold m a (MutArray a)
1824+pinnedWriteN = writeNAs Pinned
1825+
18101826-- | Like writeNWithUnsafe but writes the array in reverse order.
18111827--
18121828-- /Internal/
@@ -1874,6 +1890,29 @@ writeChunks :: (MonadIO m, Unbox a) =>
18741890 Int -> Fold m a (StreamK n (MutArray a))
18751891writeChunks n = FL.many (writeN n) FL.toStreamK
18761892
1893+{-# INLINE_NORMAL writeWithAs #-}
1894+writeWithAs :: forall m a. (MonadIO m, Unbox a)
1895+ => PinnedState -> Int -> Fold m a (MutArray a)
1896+-- writeWithAs ps n = FL.rmapM rightSize $ writeAppendWith (* 2) (newAs ps n)
1897+writeWithAs ps elemCount =
1898+ FL.rmapM extract $ FL.foldlM' step initial
1899+
1900+ where
1901+
1902+ initial = do
1903+ when (elemCount < 0) $ error "writeWith: elemCount is negative"
1904+ liftIO $ newAs ps elemCount
1905+
1906+ step arr@(MutArray _ start end bound) x
1907+ | INDEX_NEXT(end,a) > bound = do
1908+ let oldSize = end - start
1909+ newSize = max (oldSize * 2) 1
1910+ arr1 <- liftIO $ reallocExplicit (SIZE_OF(a)) newSize arr
1911+ snocUnsafe arr1 x
1912+ step arr x = snocUnsafe arr x
1913+
1914+ extract = liftIO . rightSize
1915+
18771916-- XXX Compare writeWith with fromStreamD which uses an array of streams
18781917-- implementation. We can write this using writeChunks above if that is faster.
18791918-- If writeWith is faster then we should use that to implement
@@ -1906,24 +1945,7 @@ writeChunks n = FL.many (writeN n) FL.toStreamK
19061945writeWith :: forall m a. (MonadIO m, Unbox a)
19071946 => Int -> Fold m a (MutArray a)
19081947-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (new n)
1909−writeWith elemCount =
1910− FL.rmapM extract $ FL.foldlM' step initial
1911−
1912− where
1913−
1914− initial = do
1915− when (elemCount < 0) $ error "writeWith: elemCount is negative"
1916− liftIO $ new elemCount
1917−
1918− step arr@(MutArray _ start end bound) x
1919− | INDEX_NEXT(end,a) > bound = do
1920− let oldSize = end - start
1921− newSize = max (oldSize * 2) 1
1922− arr1 <- liftIO $ reallocExplicit (SIZE_OF(a)) newSize arr
1923− snocUnsafe arr1 x
1924− step arr x = snocUnsafe arr x
1925−
1926− extract = liftIO . rightSize
1948+writeWith = writeWithAs Unpinned
19271949
19281950-- | Fold the whole input to a single array.
19291951--
@@ -1936,20 +1958,20 @@ writeWith elemCount =
19361958write :: forall m a. (MonadIO m, Unbox a) => Fold m a (MutArray a)
19371959write = writeWith (allocBytesToElemCount (undefined :: a) arrayChunkBytes)
19381960
1961+{-# INLINE pinnedWrite #-}
1962+pinnedWrite :: forall m a. (MonadIO m, Unbox a) => Fold m a (MutArray a)
1963+pinnedWrite =
1964+ writeWithAs Pinned (allocBytesToElemCount (undefined :: a) arrayChunkBytes)
1965+
19391966-------------------------------------------------------------------------------
19401967-- construct from streams, known size
19411968-------------------------------------------------------------------------------
19421969
1943−-- | Use the 'writeN' fold instead.
1944−--
1945−-- >>> fromStreamDN n = Stream.fold (MutArray.writeN n)
1946−--
1947−{-# INLINE_NORMAL fromStreamDN #-}
1948−fromStreamDN :: forall m a. (MonadIO m, Unbox a)
1949− => Int -> D.Stream m a -> m (MutArray a)
1950−-- fromStreamDN n = D.fold (writeN n)
1951−fromStreamDN limit str = do
1952− (arr :: MutArray a) <- liftIO $ new limit
1970+{-# INLINE_NORMAL fromStreamDNAs #-}
1971+fromStreamDNAs :: forall m a. (MonadIO m, Unbox a)
1972+ => PinnedState -> Int -> D.Stream m a -> m (MutArray a)
1973+fromStreamDNAs ps limit str = do
1974+ (arr :: MutArray a) <- liftIO $ newAs ps limit
19531975 end <- D.foldlM' (fwrite (arrContents arr)) (return $ arrEnd arr) $ D.take limit str
19541976 return $ arr {arrEnd = end}
19551977
@@ -1959,6 +1981,16 @@ fromStreamDN limit str = do
19591981 liftIO $ pokeByteIndex ptr arrContents x
19601982 return $ INDEX_NEXT(ptr,a)
19611983
1984+-- | Use the 'writeN' fold instead.
1985+--
1986+-- >>> fromStreamDN n = Stream.fold (MutArray.writeN n)
1987+--
1988+{-# INLINE_NORMAL fromStreamDN #-}
1989+fromStreamDN :: forall m a. (MonadIO m, Unbox a)
1990+ => Int -> D.Stream m a -> m (MutArray a)
1991+-- fromStreamDN n = D.fold (writeN n)
1992+fromStreamDN = fromStreamDNAs Unpinned
1993+
19621994-- | Create a 'MutArray' from the first N elements of a list. The array is
19631995-- allocated to size N, if the list terminates before N elements then the
19641996-- array may hold less than N elements.
@@ -1967,6 +1999,10 @@ fromStreamDN limit str = do
19671999fromListN :: (MonadIO m, Unbox a) => Int -> [a] -> m (MutArray a)
19682000fromListN n xs = fromStreamDN n $ D.fromList xs
19692001
2002+{-# INLINABLE pinnedFromListN #-}
2003+pinnedFromListN :: (MonadIO m, Unbox a) => Int -> [a] -> m (MutArray a)
2004+pinnedFromListN n xs = fromStreamDNAs Pinned n $ D.fromList xs
2005+
19702006-- | Like fromListN but writes the array in reverse order.
19712007--
19722008-- /Pre-release/
@@ -1998,6 +2034,11 @@ fromArrayStreamK as = do
19982034 len <- arrayStreamKLength as
19992035 fromStreamDN len $ D.unfoldMany reader $ D.fromStreamK as
20002036
2037+{-# INLINE fromStreamDAs #-}
2038+fromStreamDAs ::
2039+ (MonadIO m, Unbox a) => PinnedState -> D.Stream m a -> m (MutArray a)
2040+fromStreamDAs ps m = arrayStreamKFromStreamDAs ps m >>= fromArrayStreamK
2041+
20012042-- CAUTION: a very large number (millions) of arrays can degrade performance
20022043-- due to GC overhead because we need to buffer the arrays before we flatten
20032044-- all the arrays.
@@ -2013,7 +2054,7 @@ fromArrayStreamK as = do
20132054--
20142055{-# INLINE fromStreamD #-}
20152056fromStreamD :: (MonadIO m, Unbox a) => D.Stream m a -> m (MutArray a)
2016−fromStreamD m = arrayStreamKFromStreamD m >>= fromArrayStreamK
2057+fromStreamD = fromStreamDAs Unpinned
20172058
20182059-- | Create a 'MutArray' from a list. The list must be of finite size.
20192060--
@@ -2021,6 +2062,10 @@ fromStreamD m = arrayStreamKFromStreamD m >>= fromArrayStreamK
20212062fromList :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)
20222063fromList xs = fromStreamD $ D.fromList xs
20232064
2065+{-# INLINE pinnedFromList #-}
2066+pinnedFromList :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)
2067+pinnedFromList xs = fromStreamDAs Pinned $ D.fromList xs
2068+
20242069-- XXX We are materializing the whole list first for getting the length. Check
20252070-- if the 'fromList' like chunked implementation would fare better.
20262071