b1076eb · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add pinned versions of fromList(N) and write(N)

b1076ebAdithya Kumar3 years ago

core/src/Streamly/Internal/Data/Array/Mut/Type.hs+77−32

compare with master

unifiedsplit

@@ -55,18 +55,22 @@ module Streamly.Internal.Data.Array.Mut.Type

55 , writeNUnsafe

56 , writeNUnsafeAs

57 , writeN

58+ , pinnedWriteN

58 , writeNAs

59 , writeNAs

59 , pinnedWriteNAligned

60 , pinnedWriteNAligned

61 , writeWith

62 , writeWith

62 , write

63 , write

64+ , pinnedWrite

64 , writeRevN

66 , writeRevN

65 -- , writeRev

67 -- , writeRev

67 -- ** From containers

69 -- ** From containers

68 , fromListN

70 , fromListN

71+ , pinnedFromListN

69 , fromList

72 , fromList

73+ , pinnedFromList

70 , fromListRevN

74 , fromListRevN

71 , fromListRev

75 , fromListRev

72 , fromStreamDN

76 , fromStreamDN

@@ -1357,15 +1361,20 @@ chunksOf :: forall m a. (MonadIO m, Unbox a)

1357-- chunksOf n = D.foldMany (writeN n)

1361-- chunksOf n = D.foldMany (writeN n)

1358chunksOf = chunksOfAs Unpinned

1362chunksOf = chunksOfAs Unpinned

1359

1363

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+

1360-- XXX This should take a PinnedState

1371-- XXX This should take a PinnedState

1361-- XXX buffer to a list instead?

1372-- XXX buffer to a list instead?

1362-- | Buffer the stream into arrays in memory.

1373-- | Buffer the stream into arrays in memory.

1363{-# INLINE arrayStreamKFromStreamD #-}

1374{-# INLINE arrayStreamKFromStreamD #-}

1364arrayStreamKFromStreamD :: forall m a. (MonadIO m, Unbox a) =>

1375arrayStreamKFromStreamD :: forall m a. (MonadIO m, Unbox a) =>

1365 D.Stream m a -> m (StreamK m (MutArray a))

1376 D.Stream m a -> m (StreamK m (MutArray a))

1366−arrayStreamKFromStreamD =

1377+arrayStreamKFromStreamD = arrayStreamKFromStreamDAs Pinned

1367− let n = allocBytesToElemCount (undefined :: a) defaultChunkSize

1368− in D.foldr K.cons K.nil . chunksOf n

1369

1378

1370-------------------------------------------------------------------------------

1379-------------------------------------------------------------------------------

1371-- Streams of arrays - Flattening

1380-- Streams of arrays - Flattening

@@ -1807,6 +1816,13 @@ writeNAs ps = writeNWith (newAs ps)

1807writeN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (MutArray a)

1816writeN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (MutArray a)

1808writeN = writeNAs Unpinned

1817writeN = writeNAs Unpinned

1809

1818

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+

1810-- | Like writeNWithUnsafe but writes the array in reverse order.

1826-- | Like writeNWithUnsafe but writes the array in reverse order.

1811--

1827--

1812-- /Internal/

1828-- /Internal/

@@ -1874,6 +1890,29 @@ writeChunks :: (MonadIO m, Unbox a) =>

1874 Int -> Fold m a (StreamK n (MutArray a))

1890 Int -> Fold m a (StreamK n (MutArray a))

1875writeChunks n = FL.many (writeN n) FL.toStreamK

1891writeChunks n = FL.many (writeN n) FL.toStreamK

1876

1892

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+

1877-- XXX Compare writeWith with fromStreamD which uses an array of streams

1916-- XXX Compare writeWith with fromStreamD which uses an array of streams

1878-- implementation. We can write this using writeChunks above if that is faster.

1917-- implementation. We can write this using writeChunks above if that is faster.

1879-- If writeWith is faster then we should use that to implement

1918-- If writeWith is faster then we should use that to implement

@@ -1906,24 +1945,7 @@ writeChunks n = FL.many (writeN n) FL.toStreamK

1906writeWith :: forall m a. (MonadIO m, Unbox a)

1945writeWith :: forall m a. (MonadIO m, Unbox a)

1907 => Int -> Fold m a (MutArray a)

1946 => Int -> Fold m a (MutArray a)

1908-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (new n)

1947-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (new n)

1909−writeWith elemCount =

1948+writeWith = writeWithAs Unpinned

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

1927

1949

1928-- | Fold the whole input to a single array.

1950-- | Fold the whole input to a single array.

1929--

1951--

@@ -1936,20 +1958,20 @@ writeWith elemCount =

1936write :: forall m a. (MonadIO m, Unbox a) => Fold m a (MutArray a)

1958write :: forall m a. (MonadIO m, Unbox a) => Fold m a (MutArray a)

1937write = writeWith (allocBytesToElemCount (undefined :: a) arrayChunkBytes)

1959write = writeWith (allocBytesToElemCount (undefined :: a) arrayChunkBytes)

1938

1960

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+

1939-------------------------------------------------------------------------------

1966-------------------------------------------------------------------------------

1940-- construct from streams, known size

1967-- construct from streams, known size

1941-------------------------------------------------------------------------------

1968-------------------------------------------------------------------------------

1942

1969

1943−-- | Use the 'writeN' fold instead.

1944−--

1945−-- >>> fromStreamDN n = Stream.fold (MutArray.writeN n)

1946−--

1947−{-# INLINE_NORMAL fromStreamDN #-}

1970+{-# INLINE_NORMAL fromStreamDNAs #-}

1948−fromStreamDN :: forall m a. (MonadIO m, Unbox a)

1971+fromStreamDNAs :: forall m a. (MonadIO m, Unbox a)

1949− => Int -> D.Stream m a -> m (MutArray a)

1972+ => PinnedState -> Int -> D.Stream m a -> m (MutArray a)

1950−-- fromStreamDN n = D.fold (writeN n)

1951−fromStreamDN limit str = do

1973+fromStreamDNAs ps limit str = do

1952− (arr :: MutArray a) <- liftIO $ new limit

1974+ (arr :: MutArray a) <- liftIO $ newAs ps limit

1953 end <- D.foldlM' (fwrite (arrContents arr)) (return $ arrEnd arr) $ D.take limit str

1975 end <- D.foldlM' (fwrite (arrContents arr)) (return $ arrEnd arr) $ D.take limit str

1954 return $ arr {arrEnd = end}

1976 return $ arr {arrEnd = end}

1955

1977

@@ -1959,6 +1981,16 @@ fromStreamDN limit str = do

1959 liftIO $ pokeByteIndex ptr arrContents x

1981 liftIO $ pokeByteIndex ptr arrContents x

1960 return $ INDEX_NEXT(ptr,a)

1982 return $ INDEX_NEXT(ptr,a)

1961

1983

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+

1962-- | Create a 'MutArray' from the first N elements of a list. The array is

1994-- | Create a 'MutArray' from the first N elements of a list. The array is

1963-- allocated to size N, if the list terminates before N elements then the

1995-- allocated to size N, if the list terminates before N elements then the

1964-- array may hold less than N elements.

1996-- array may hold less than N elements.

@@ -1967,6 +1999,10 @@ fromStreamDN limit str = do

1967fromListN :: (MonadIO m, Unbox a) => Int -> [a] -> m (MutArray a)

1999fromListN :: (MonadIO m, Unbox a) => Int -> [a] -> m (MutArray a)

1968fromListN n xs = fromStreamDN n $ D.fromList xs

2000fromListN n xs = fromStreamDN n $ D.fromList xs

1969

2001

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+

1970-- | Like fromListN but writes the array in reverse order.

2006-- | Like fromListN but writes the array in reverse order.

1971--

2007--

1972-- /Pre-release/

2008-- /Pre-release/

@@ -1998,6 +2034,11 @@ fromArrayStreamK as = do

1998 len <- arrayStreamKLength as

2034 len <- arrayStreamKLength as

1999 fromStreamDN len $ D.unfoldMany reader $ D.fromStreamK as

2035 fromStreamDN len $ D.unfoldMany reader $ D.fromStreamK as

2000

2036

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+

2001-- CAUTION: a very large number (millions) of arrays can degrade performance

2042-- CAUTION: a very large number (millions) of arrays can degrade performance

2002-- due to GC overhead because we need to buffer the arrays before we flatten

2043-- due to GC overhead because we need to buffer the arrays before we flatten

2003-- all the arrays.

2044-- all the arrays.

@@ -2013,7 +2054,7 @@ fromArrayStreamK as = do

2013--

2054--

2014{-# INLINE fromStreamD #-}

2055{-# INLINE fromStreamD #-}

2015fromStreamD :: (MonadIO m, Unbox a) => D.Stream m a -> m (MutArray a)

2056fromStreamD :: (MonadIO m, Unbox a) => D.Stream m a -> m (MutArray a)

2016−fromStreamD m = arrayStreamKFromStreamD m >>= fromArrayStreamK

2057+fromStreamD = fromStreamDAs Unpinned

2017

2058

2018-- | Create a 'MutArray' from a list. The list must be of finite size.

2059-- | Create a 'MutArray' from a list. The list must be of finite size.

2019--

2060--

@@ -2021,6 +2062,10 @@ fromStreamD m = arrayStreamKFromStreamD m >>= fromArrayStreamK

2021fromList :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)

2062fromList :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)

2022fromList xs = fromStreamD $ D.fromList xs

2063fromList xs = fromStreamD $ D.fromList xs

2023

2064

2065+{-# INLINE pinnedFromList #-}

2066+pinnedFromList :: (MonadIO m, Unbox a) => [a] -> m (MutArray a)

2067+pinnedFromList xs = fromStreamDAs Pinned $ D.fromList xs

2068+

2024-- XXX We are materializing the whole list first for getting the length. Check

2069-- XXX We are materializing the whole list first for getting the length. Check

2025-- if the 'fromList' like chunked implementation would fare better.

2070-- if the 'fromList' like chunked implementation would fare better.

2026

2071

core/src/Streamly/Internal/Data/Array/Type.hs+22

@@ -30,7 +30,9 @@ module Streamly.Internal.Data.Array.Type

30 , splice

32 , fromList

33+ , pinnedFromList

33 , fromListN

34 , fromListN

35+ , pinnedFromListN

34 , fromListRev

36 , fromListRev

35 , fromListRevN

37 , fromListRevN

36 , fromStreamDN

38 , fromStreamDN

@@ -65,12 +67,14 @@ module Streamly.Internal.Data.Array.Type

65 -- * Folds

67 -- * Folds

66 , writeWith

68 , writeWith

67 , writeN

69 , writeN

70+ , pinnedWriteN

68 , writeNAs

71 , writeNAs

69 , writeNUnsafe

72 , writeNUnsafe

70 , writeNUnsafeAs

73 , writeNUnsafeAs

71 , MA.ArrayUnsafe (..)

74 , MA.ArrayUnsafe (..)

72 , pinnedWriteNAligned

75 , pinnedWriteNAligned

73 , write

76 , write

77+ , pinnedWrite

74 , unsafeMakePure

78 , unsafeMakePure

76 -- * Streams of arrays

80 -- * Streams of arrays

@@ -252,6 +256,12 @@ splice arr1 arr2 =

252fromListN :: Unbox a => Int -> [a] -> Array a

256fromListN :: Unbox a => Int -> [a] -> Array a

253fromListN n xs = unsafePerformIO $ unsafeFreeze <$> MA.fromListN n xs

257fromListN n xs = unsafePerformIO $ unsafeFreeze <$> MA.fromListN n xs

254

258

259+-- | Like 'fromListN' but creates a pinned 'Array'

260+{-# INLINABLE pinnedFromListN #-}

261+pinnedFromListN :: Unbox a => Int -> [a] -> Array a

262+pinnedFromListN n xs =

263+ unsafePerformIO $ unsafeFreeze <$> MA.pinnedFromListN n xs

264+

255-- | Create an 'Array' from the first N elements of a list in reverse order.

265-- | Create an 'Array' from the first N elements of a list in reverse order.

256-- The array is allocated to size N, if the list terminates before N elements

266-- The array is allocated to size N, if the list terminates before N elements

257-- then the array may hold less than N elements.

267-- then the array may hold less than N elements.

@@ -267,6 +277,10 @@ fromListRevN n xs = unsafePerformIO $ unsafeFreeze <$> MA.fromListRevN n xs

267fromList :: Unbox a => [a] -> Array a

277fromList :: Unbox a => [a] -> Array a

268fromList xs = unsafePerformIO $ unsafeFreeze <$> MA.fromList xs

278fromList xs = unsafePerformIO $ unsafeFreeze <$> MA.fromList xs

269

279

280+{-# INLINE pinnedFromList #-}

281+pinnedFromList :: Unbox a => [a] -> Array a

282+pinnedFromList xs = unsafePerformIO $ unsafeFreeze <$> MA.pinnedFromList xs

283+

270-- | Create an 'Array' from a list in reverse order. The list must be of finite

284-- | Create an 'Array' from a list in reverse order. The list must be of finite

271-- size.

285-- size.

272--

286--

@@ -489,6 +503,10 @@ writeNAs ps = fmap unsafeFreeze . MA.writeNAs ps

489writeN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (Array a)

503writeN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (Array a)

490writeN = fmap unsafeFreeze . MA.writeN

504writeN = fmap unsafeFreeze . MA.writeN

491

505

506+{-# INLINE_NORMAL pinnedWriteN #-}

507+pinnedWriteN :: forall m a. (MonadIO m, Unbox a) => Int -> Fold m a (Array a)

508+pinnedWriteN = fmap unsafeFreeze . MA.pinnedWriteN

509+

492-- | @pinnedWriteNAligned alignment n@ folds a maximum of @n@ elements from the input

510-- | @pinnedWriteNAligned alignment n@ folds a maximum of @n@ elements from the input

493-- stream to an 'Array' aligned to the given size.

511-- stream to an 'Array' aligned to the given size.

494--

512--

@@ -529,6 +547,10 @@ writeWith elemCount = unsafeFreeze <$> MA.writeWith elemCount

529write :: forall m a. (MonadIO m, Unbox a) => Fold m a (Array a)

547write :: forall m a. (MonadIO m, Unbox a) => Fold m a (Array a)

530write = fmap unsafeFreeze MA.write

548write = fmap unsafeFreeze MA.write

531

549

550+{-# INLINE pinnedWrite #-}

551+pinnedWrite :: forall m a. (MonadIO m, Unbox a) => Fold m a (Array a)

552+pinnedWrite = fmap unsafeFreeze MA.write

553+

532-- | Fold "step" has a dependency on "initial", and each step is dependent on

554-- | Fold "step" has a dependency on "initial", and each step is dependent on

533-- the previous invocation of step due to state passing, finally extract

555-- the previous invocation of step due to state passing, finally extract

534-- depends on the result of step, therefore, as long as the fold is driven in

556-- depends on the result of step, therefore, as long as the fold is driven in