mirror of
https://github.com/composewell/streamly.git
synced 2024-10-26 19:50:19 +03:00
Add Data.Prim.Pinned.Array and extend Data.Prim.Array
This commit is contained in:
adithyaov
parent
81339c8f8b
commit
33476c33ff
@ -16,3 +16,4 @@ src/Streamly/Internal/Data/Prim/Array/Types.hs
|
||||
src/Streamly/Internal/Data/SmallArray/Types.hs
|
||||
src/Streamly/Internal/Data/Unicode/Stream.hs
|
||||
src/Streamly/Internal/Mutable/Prim/Var.hs
|
||||
src/prim-array-types.hs
|
||||
|
||||
11
bench.sh
11
bench.sh
@ -29,7 +29,12 @@ prelude_other_grp="\
|
||||
Prelude.Concurrent \
|
||||
Prelude.Adaptive"
|
||||
|
||||
array_grp="Memory.Array Data.Array Data.Prim.Array Data.SmallArray"
|
||||
array_grp="\
|
||||
Memory.Array \
|
||||
Data.Array \
|
||||
Data.Prim.Array \
|
||||
Data.SmallArray \
|
||||
Data.Prim.Pinned.Array"
|
||||
|
||||
base_parser_grp="Data.Parser.ParserD Data.Parser.ParserK"
|
||||
parser_grp="Data.Fold Data.Parser"
|
||||
@ -54,7 +59,8 @@ base_stream_cmp="Data.Stream.StreamD Data.Stream.StreamK"
|
||||
serial_wserial_cmp="Prelude.Serial Prelude.WSerial"
|
||||
serial_async_cmp="Prelude.Serial Prelude.Async"
|
||||
concurrent_cmp="Prelude.Async Prelude.WAsync Prelude.Ahead Prelude.Parallel"
|
||||
array_cmp="Memory.Array Data.Prim.Array Data.Array"
|
||||
array_cmp="Memory.Array Data.Prim.Array Data.Array Data.Prim.Pinned.Array"
|
||||
pinned_array_cmp="Memory.Array Data.Prim.Pinned.Array"
|
||||
base_parser_cmp=$base_parser_grp
|
||||
COMPARISONS="\
|
||||
base_stream_cmp \
|
||||
@ -62,6 +68,7 @@ COMPARISONS="\
|
||||
serial_async_cmp \
|
||||
concurrent_cmp \
|
||||
array_cmp \
|
||||
pinned_array_cmp \
|
||||
base_parser_cmp"
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
@ -13,6 +13,7 @@ module Main (main) where
|
||||
|
||||
import Control.DeepSeq (NFData(..))
|
||||
import System.Random (randomRIO)
|
||||
import Data.Primitive.Types (Prim)
|
||||
|
||||
import qualified Streamly.Benchmark.Data.Prim.ArrayOps as Ops
|
||||
import qualified Streamly.Internal.Data.Prim.Array as A
|
||||
@ -32,7 +33,7 @@ benchIO name src f = bench name $ nfIO $
|
||||
|
||||
-- Drain a source that generates an array in the IO monad
|
||||
{-# INLINE benchIOSrc #-}
|
||||
benchIOSrc :: A.Prim a => String -> (Int -> IO (Ops.Stream a)) -> Benchmark
|
||||
benchIOSrc :: Prim a => String -> (Int -> IO (Ops.Stream a)) -> Benchmark
|
||||
benchIOSrc name src = benchIO name src id
|
||||
|
||||
{-# INLINE benchPureSink #-}
|
||||
|
||||
@ -15,7 +15,7 @@
|
||||
|
||||
module Streamly.Benchmark.Data.Prim.ArrayOps where
|
||||
|
||||
import Control.Monad.IO.Class (MonadIO)
|
||||
import Control.Monad.Primitive (PrimMonad)
|
||||
import Prelude (Int, Bool, (+), ($), (==), (>), (.), Maybe(..), undefined)
|
||||
import qualified Prelude as P
|
||||
#ifdef DEVBUILD
|
||||
@ -37,10 +37,10 @@ value = 100000
|
||||
-- Stream generation and elimination
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
type Stream = A.PrimArray
|
||||
type Stream = A.Array
|
||||
|
||||
{-# INLINE sourceUnfoldr #-}
|
||||
sourceUnfoldr :: MonadIO m => Int -> m (Stream Int)
|
||||
sourceUnfoldr :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceUnfoldr n = S.fold (A.writeN value) $ S.unfoldr step n
|
||||
where
|
||||
step cnt =
|
||||
@ -49,19 +49,19 @@ sourceUnfoldr n = S.fold (A.writeN value) $ S.unfoldr step n
|
||||
else (Just (cnt, cnt + 1))
|
||||
|
||||
{-# INLINE sourceIntFromTo #-}
|
||||
sourceIntFromTo :: MonadIO m => Int -> m (Stream Int)
|
||||
sourceIntFromTo :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceIntFromTo n = S.fold (A.writeN value) $ S.enumerateFromTo n (n + value)
|
||||
|
||||
{-# INLINE sourceIntFromToFromStream #-}
|
||||
sourceIntFromToFromStream :: MonadIO m => Int -> m (Stream Int)
|
||||
sourceIntFromToFromStream :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceIntFromToFromStream n = S.fold A.write $ S.enumerateFromTo n (n + value)
|
||||
|
||||
{-# INLINE sourceIntFromToFromList #-}
|
||||
sourceIntFromToFromList :: MonadIO m => Int -> m (Stream Int)
|
||||
sourceIntFromToFromList :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceIntFromToFromList n = P.return $ A.fromList $ [n..n + value]
|
||||
|
||||
{-# INLINE sourceFromList #-}
|
||||
sourceFromList :: MonadIO m => Int -> m (Stream Int)
|
||||
sourceFromList :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceFromList n = S.fold (A.writeN value) $ S.fromList [n..n+value]
|
||||
{-
|
||||
{-# INLINE sourceIsList #-}
|
||||
@ -92,11 +92,11 @@ composeN n f x =
|
||||
{-# INLINE map #-}
|
||||
|
||||
scanl', scanl1', map
|
||||
:: MonadIO m => Int -> Stream Int -> m (Stream Int)
|
||||
:: PrimMonad m => Int -> Stream Int -> m (Stream Int)
|
||||
|
||||
{-# INLINE onArray #-}
|
||||
onArray
|
||||
:: MonadIO m => (S.SerialT m Int -> S.SerialT m Int)
|
||||
:: PrimMonad m => (S.SerialT m Int -> S.SerialT m Int)
|
||||
-> Stream Int
|
||||
-> m (Stream Int)
|
||||
onArray f arr = S.fold (A.writeN value) $ f $ (S.unfold A.read arr)
|
||||
@ -136,7 +136,7 @@ readInstance str =
|
||||
-}
|
||||
|
||||
{-# INLINE pureFoldl' #-}
|
||||
pureFoldl' :: MonadIO m => Stream Int -> m Int
|
||||
pureFoldl' :: PrimMonad m => Stream Int -> m Int
|
||||
pureFoldl' = S.foldl' (+) 0 . S.unfold A.read
|
||||
|
||||
#if 0
|
||||
|
||||
251
benchmark/Streamly/Benchmark/Data/Prim/Pinned/Array.hs
Normal file
251
benchmark/Streamly/Benchmark/Data/Prim/Pinned/Array.hs
Normal file
@ -0,0 +1,251 @@
|
||||
{-# LANGUAGE CPP #-}
|
||||
-- |
|
||||
-- Module : Main
|
||||
-- Copyright : (c) 2018 Harendra Kumar
|
||||
--
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
|
||||
import Control.DeepSeq (NFData(..), deepseq)
|
||||
import System.Random (randomRIO)
|
||||
import Data.Primitive.Types (Prim)
|
||||
|
||||
import qualified GHC.Exts as GHC
|
||||
|
||||
import qualified Streamly.Benchmark.Data.Prim.Pinned.ArrayOps as Ops
|
||||
import qualified Streamly.Internal.Data.Prim.Pinned.Array as A
|
||||
import qualified Streamly.Prelude as S
|
||||
|
||||
import Streamly.Benchmark.Common hiding (benchPureSrc)
|
||||
import Gauge
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
--
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- Drain a source that generates a pure array
|
||||
{-# INLINE benchPureSrc #-}
|
||||
benchPureSrc :: (Prim a)
|
||||
=> String -> (Int -> Ops.Stream a) -> Benchmark
|
||||
benchPureSrc name src = benchPure name src id
|
||||
|
||||
{-# INLINE benchIO #-}
|
||||
benchIO :: NFData b => String -> (Int -> IO a) -> (a -> b) -> Benchmark
|
||||
benchIO name src f = bench name $ nfIO $
|
||||
randomRIO (1,1) >>= src >>= return . f
|
||||
|
||||
-- Drain a source that generates an array in the IO monad
|
||||
{-# INLINE benchIOSrc #-}
|
||||
benchIOSrc :: (Prim a)
|
||||
=> String -> (Int -> IO (Ops.Stream a)) -> Benchmark
|
||||
benchIOSrc name src = benchIO name src id
|
||||
|
||||
{-# INLINE benchPureSink #-}
|
||||
benchPureSink :: NFData b => String -> (Ops.Stream Int -> b) -> Benchmark
|
||||
benchPureSink name f = benchIO name Ops.sourceIntFromTo f
|
||||
|
||||
{-# INLINE benchIO' #-}
|
||||
benchIO' :: NFData b => String -> (Int -> IO a) -> (a -> IO b) -> Benchmark
|
||||
benchIO' name src f = bench name $ nfIO $
|
||||
randomRIO (1,1) >>= src >>= f
|
||||
|
||||
{-# INLINE benchIOSink #-}
|
||||
benchIOSink :: NFData b => String -> (Ops.Stream Int -> IO b) -> Benchmark
|
||||
benchIOSink name f = benchIO' name Ops.sourceIntFromTo f
|
||||
|
||||
testStr :: String
|
||||
testStr = mkListString Ops.value
|
||||
|
||||
moduleName :: String
|
||||
moduleName = "Data.Prim.Pinned.Array"
|
||||
|
||||
o_1_space :: [Benchmark]
|
||||
o_1_space =
|
||||
[ bgroup (o_1_space_prefix moduleName)
|
||||
[ bgroup "generation"
|
||||
[ benchIOSrc "writeN . intFromTo" Ops.sourceIntFromTo
|
||||
, benchIOSrc "write . intFromTo" Ops.sourceIntFromToFromStream
|
||||
, benchIOSrc "fromList . intFromTo" Ops.sourceIntFromToFromList
|
||||
, benchIOSrc "writeN . unfoldr" Ops.sourceUnfoldr
|
||||
, benchIOSrc "writeN . fromList" Ops.sourceFromList
|
||||
, benchPureSrc "writeN . IsList.fromList" Ops.sourceIsList
|
||||
, benchPureSrc "writeN . IsString.fromString" Ops.sourceIsString
|
||||
, testStr `deepseq` (bench "read" $ nf Ops.readInstance testStr)
|
||||
, benchPureSink "show" Ops.showInstance
|
||||
]
|
||||
, bgroup "elimination"
|
||||
[ benchPureSink "id" id
|
||||
-- , benchPureSink "eqBy" Ops.eqBy
|
||||
, benchPureSink "==" Ops.eqInstance
|
||||
, benchPureSink "/=" Ops.eqInstanceNotEq
|
||||
{-
|
||||
, benchPureSink "cmpBy" Ops.cmpBy
|
||||
-}
|
||||
, benchPureSink "<" Ops.ordInstance
|
||||
, benchPureSink "min" Ops.ordInstanceMin
|
||||
-- length is used to check for foldr/build fusion
|
||||
, benchPureSink "length . IsList.toList" (length . GHC.toList)
|
||||
, benchIOSink "foldl'" Ops.pureFoldl'
|
||||
, benchIOSink "read" (S.drain . S.unfold A.read)
|
||||
, benchIOSink "toStreamRev" (S.drain . A.toStreamRev)
|
||||
#ifdef DEVBUILD
|
||||
, benchPureSink "foldable/foldl'" Ops.foldableFoldl'
|
||||
, benchPureSink "foldable/sum" Ops.foldableSum
|
||||
-- , benchPureSinkIO "traversable/mapM" Ops.traversableMapM
|
||||
#endif
|
||||
]
|
||||
|
||||
{-
|
||||
[ benchPureSink "uncons" Ops.uncons
|
||||
, benchPureSink "toNull" $ Ops.toNull serially
|
||||
, benchPureSink "mapM_" Ops.mapM_
|
||||
|
||||
, benchPureSink "init" Ops.init
|
||||
, benchPureSink "tail" Ops.tail
|
||||
, benchPureSink "nullHeadTail" Ops.nullHeadTail
|
||||
|
||||
-- this is too low and causes all benchmarks reported in ns
|
||||
-- , benchPureSink "head" Ops.head
|
||||
, benchPureSink "last" Ops.last
|
||||
-- , benchPureSink "lookup" Ops.lookup
|
||||
, benchPureSink "find" Ops.find
|
||||
, benchPureSink "findIndex" Ops.findIndex
|
||||
, benchPureSink "elemIndex" Ops.elemIndex
|
||||
|
||||
-- this is too low and causes all benchmarks reported in ns
|
||||
-- , benchPureSink "null" Ops.null
|
||||
, benchPureSink "elem" Ops.elem
|
||||
, benchPureSink "notElem" Ops.notElem
|
||||
, benchPureSink "all" Ops.all
|
||||
, benchPureSink "any" Ops.any
|
||||
, benchPureSink "and" Ops.and
|
||||
, benchPureSink "or" Ops.or
|
||||
|
||||
, benchPureSink "length" Ops.length
|
||||
, benchPureSink "sum" Ops.sum
|
||||
, benchPureSink "product" Ops.product
|
||||
|
||||
, benchPureSink "maximumBy" Ops.maximumBy
|
||||
, benchPureSink "maximum" Ops.maximum
|
||||
, benchPureSink "minimumBy" Ops.minimumBy
|
||||
, benchPureSink "minimum" Ops.minimum
|
||||
|
||||
, benchPureSink "toList" Ops.toList
|
||||
, benchPureSink "toRevList" Ops.toRevList
|
||||
]
|
||||
-}
|
||||
, bgroup "transformation"
|
||||
[ benchIOSink "scanl'" (Ops.scanl' 1)
|
||||
, benchIOSink "scanl1'" (Ops.scanl1' 1)
|
||||
, benchIOSink "map" (Ops.map 1)
|
||||
{-
|
||||
, benchPureSink "fmap" (Ops.fmap 1)
|
||||
, benchPureSink "mapM" (Ops.mapM serially 1)
|
||||
, benchPureSink "mapMaybe" (Ops.mapMaybe 1)
|
||||
, benchPureSink "mapMaybeM" (Ops.mapMaybeM 1)
|
||||
, bench "sequence" $ nfIO $ randomRIO (1,1000) >>= \n ->
|
||||
Ops.sequence serially (Ops.sourceUnfoldrMAction n)
|
||||
, benchPureSink "findIndices" (Ops.findIndices 1)
|
||||
, benchPureSink "elemIndices" (Ops.elemIndices 1)
|
||||
, benchPureSink "reverse" (Ops.reverse 1)
|
||||
, benchPureSink "foldrS" (Ops.foldrS 1)
|
||||
, benchPureSink "foldrSMap" (Ops.foldrSMap 1)
|
||||
, benchPureSink "foldrT" (Ops.foldrT 1)
|
||||
, benchPureSink "foldrTMap" (Ops.foldrTMap 1)
|
||||
-}
|
||||
]
|
||||
, bgroup "transformationX4"
|
||||
[ benchIOSink "scanl'" (Ops.scanl' 4)
|
||||
, benchIOSink "scanl1'" (Ops.scanl1' 4)
|
||||
, benchIOSink "map" (Ops.map 4)
|
||||
{-
|
||||
, benchPureSink "fmap" (Ops.fmap 4)
|
||||
, benchPureSink "mapM" (Ops.mapM serially 4)
|
||||
, benchPureSink "mapMaybe" (Ops.mapMaybe 4)
|
||||
, benchPureSink "mapMaybeM" (Ops.mapMaybeM 4)
|
||||
-- , bench "sequence" $ nfIO $ randomRIO (1,1000) >>= \n ->
|
||||
-- Ops.sequence serially (Ops.sourceUnfoldrMAction n)
|
||||
, benchPureSink "findIndices" (Ops.findIndices 4)
|
||||
, benchPureSink "elemIndices" (Ops.elemIndices 4)
|
||||
-}
|
||||
]
|
||||
{-
|
||||
, bgroup "filtering"
|
||||
[ benchPureSink "filter-even" (Ops.filterEven 1)
|
||||
, benchPureSink "filter-all-out" (Ops.filterAllOut 1)
|
||||
, benchPureSink "filter-all-in" (Ops.filterAllIn 1)
|
||||
, benchPureSink "take-all" (Ops.takeAll 1)
|
||||
, benchPureSink "takeWhile-true" (Ops.takeWhileTrue 1)
|
||||
--, benchPureSink "takeWhileM-true" (Ops.takeWhileMTrue 1)
|
||||
, benchPureSink "drop-one" (Ops.dropOne 1)
|
||||
, benchPureSink "drop-all" (Ops.dropAll 1)
|
||||
, benchPureSink "dropWhile-true" (Ops.dropWhileTrue 1)
|
||||
--, benchPureSink "dropWhileM-true" (Ops.dropWhileMTrue 1)
|
||||
, benchPureSink "dropWhile-false" (Ops.dropWhileFalse 1)
|
||||
, benchPureSink "deleteBy" (Ops.deleteBy 1)
|
||||
, benchPureSink "insertBy" (Ops.insertBy 1)
|
||||
]
|
||||
, bgroup "filteringX4"
|
||||
[ benchPureSink "filter-even" (Ops.filterEven 4)
|
||||
, benchPureSink "filter-all-out" (Ops.filterAllOut 4)
|
||||
, benchPureSink "filter-all-in" (Ops.filterAllIn 4)
|
||||
, benchPureSink "take-all" (Ops.takeAll 4)
|
||||
, benchPureSink "takeWhile-true" (Ops.takeWhileTrue 4)
|
||||
--, benchPureSink "takeWhileM-true" (Ops.takeWhileMTrue 4)
|
||||
, benchPureSink "drop-one" (Ops.dropOne 4)
|
||||
, benchPureSink "drop-all" (Ops.dropAll 4)
|
||||
, benchPureSink "dropWhile-true" (Ops.dropWhileTrue 4)
|
||||
--, benchPureSink "dropWhileM-true" (Ops.dropWhileMTrue 4)
|
||||
, benchPureSink "dropWhile-false" (Ops.dropWhileFalse 4)
|
||||
, benchPureSink "deleteBy" (Ops.deleteBy 4)
|
||||
, benchPureSink "insertBy" (Ops.insertBy 4)
|
||||
]
|
||||
, bgroup "multi-stream"
|
||||
[ benchPureSink "eqBy" Ops.eqBy
|
||||
, benchPureSink "cmpBy" Ops.cmpBy
|
||||
, benchPureSink "zip" Ops.zip
|
||||
, benchPureSink "zipM" Ops.zipM
|
||||
, benchPureSink "mergeBy" Ops.mergeBy
|
||||
, benchPureSink "isPrefixOf" Ops.isPrefixOf
|
||||
, benchPureSink "isSubsequenceOf" Ops.isSubsequenceOf
|
||||
, benchPureSink "stripPrefix" Ops.stripPrefix
|
||||
, benchPureSrc serially "concatMap" Ops.concatMap
|
||||
]
|
||||
-- scanl-map and foldl-map are equivalent to the scan and fold in the foldl
|
||||
-- library. If scan/fold followed by a map is efficient enough we may not
|
||||
-- need monolithic implementations of these.
|
||||
, bgroup "mixed"
|
||||
[ benchPureSink "scanl-map" (Ops.scanMap 1)
|
||||
, benchPureSink "foldl-map" Ops.foldl'ReduceMap
|
||||
, benchPureSink "sum-product-fold" Ops.sumProductFold
|
||||
, benchPureSink "sum-product-scan" Ops.sumProductScan
|
||||
]
|
||||
, bgroup "mixedX4"
|
||||
[ benchPureSink "scan-map" (Ops.scanMap 4)
|
||||
, benchPureSink "drop-map" (Ops.dropMap 4)
|
||||
, benchPureSink "drop-scan" (Ops.dropScan 4)
|
||||
, benchPureSink "take-drop" (Ops.takeDrop 4)
|
||||
, benchPureSink "take-scan" (Ops.takeScan 4)
|
||||
, benchPureSink "take-map" (Ops.takeMap 4)
|
||||
, benchPureSink "filter-drop" (Ops.filterDrop 4)
|
||||
, benchPureSink "filter-take" (Ops.filterTake 4)
|
||||
, benchPureSink "filter-scan" (Ops.filterScan 4)
|
||||
, benchPureSink "filter-scanl1" (Ops.filterScanl1 4)
|
||||
, benchPureSink "filter-map" (Ops.filterMap 4)
|
||||
]
|
||||
, bgroup "iterated"
|
||||
[ benchPureSrc serially "mapM" Ops.iterateMapM
|
||||
, benchPureSrc serially "scan(1/100)" Ops.iterateScan
|
||||
, benchPureSrc serially "scanl1(1/100)" Ops.iterateScanl1
|
||||
, benchPureSrc serially "filterEven" Ops.iterateFilterEven
|
||||
, benchPureSrc serially "takeAll" Ops.iterateTakeAll
|
||||
, benchPureSrc serially "dropOne" Ops.iterateDropOne
|
||||
, benchPureSrc serially "dropWhileFalse" Ops.iterateDropWhileFalse
|
||||
, benchPureSrc serially "dropWhileTrue" Ops.iterateDropWhileTrue
|
||||
]
|
||||
-}
|
||||
]
|
||||
]
|
||||
|
||||
main :: IO ()
|
||||
main = defaultMain $ concat [o_1_space]
|
||||
531
benchmark/Streamly/Benchmark/Data/Prim/Pinned/ArrayOps.hs
Normal file
531
benchmark/Streamly/Benchmark/Data/Prim/Pinned/ArrayOps.hs
Normal file
@ -0,0 +1,531 @@
|
||||
-- |
|
||||
-- Module : ArrayOps
|
||||
-- Copyright : (c) 2018 Harendra Kumar
|
||||
--
|
||||
-- License : MIT
|
||||
-- Maintainer : streamly@composewell.com
|
||||
|
||||
{-# LANGUAGE CPP #-}
|
||||
{-# LANGUAGE FlexibleContexts #-}
|
||||
{-# LANGUAGE ScopedTypeVariables #-}
|
||||
{-# LANGUAGE DeriveAnyClass #-}
|
||||
{-# LANGUAGE DeriveGeneric #-}
|
||||
|
||||
module Streamly.Benchmark.Data.Prim.Pinned.ArrayOps where
|
||||
|
||||
-- import Control.Monad (when)
|
||||
import Control.Monad.Primitive (PrimMonad)
|
||||
-- import Data.Maybe (fromJust)
|
||||
import Prelude (Int, Bool, (+), ($), (==), (>), (.), Maybe(..), undefined)
|
||||
import qualified Prelude as P
|
||||
#ifdef DEVBUILD
|
||||
import qualified Data.Foldable as F
|
||||
#endif
|
||||
import qualified GHC.Exts as GHC
|
||||
-- import Control.DeepSeq (NFData)
|
||||
-- import GHC.Generics (Generic)
|
||||
|
||||
import qualified Streamly as S hiding (foldMapWith, runStream)
|
||||
import qualified Streamly.Internal.Data.Prim.Pinned.Array as A
|
||||
import qualified Streamly.Prelude as S
|
||||
|
||||
value, maxValue :: Int
|
||||
#ifdef LINEAR_ASYNC
|
||||
value = 10000
|
||||
#else
|
||||
value = 100000
|
||||
#endif
|
||||
maxValue = value + 1
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Benchmark ops
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Stream generation and elimination
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
type Stream = A.Array
|
||||
|
||||
{-# INLINE sourceUnfoldr #-}
|
||||
sourceUnfoldr :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceUnfoldr n = S.fold (A.writeN value) $ S.unfoldr step n
|
||||
where
|
||||
step cnt =
|
||||
if cnt > n + value
|
||||
then Nothing
|
||||
else (Just (cnt, cnt + 1))
|
||||
|
||||
{-# INLINE sourceIntFromTo #-}
|
||||
sourceIntFromTo :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceIntFromTo n = S.fold (A.writeN value) $ S.enumerateFromTo n (n + value)
|
||||
|
||||
{-# INLINE sourceIntFromToFromStream #-}
|
||||
sourceIntFromToFromStream :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceIntFromToFromStream n = S.fold A.write $ S.enumerateFromTo n (n + value)
|
||||
|
||||
sourceIntFromToFromList :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceIntFromToFromList n = P.return $ A.fromList $ [n..n + value]
|
||||
|
||||
{-# INLINE sourceFromList #-}
|
||||
sourceFromList :: PrimMonad m => Int -> m (Stream Int)
|
||||
sourceFromList n = S.fold (A.writeN value) $ S.fromList [n..n+value]
|
||||
|
||||
{-# INLINE sourceIsList #-}
|
||||
sourceIsList :: Int -> Stream Int
|
||||
sourceIsList n = GHC.fromList [n..n+value]
|
||||
|
||||
{-# INLINE sourceIsString #-}
|
||||
sourceIsString :: Int -> Stream P.Char
|
||||
sourceIsString n = GHC.fromString (P.replicate (n + value) 'a')
|
||||
|
||||
{-
|
||||
-------------------------------------------------------------------------------
|
||||
-- Elimination
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE runStream #-}
|
||||
runStream :: Monad m => Stream m a -> m ()
|
||||
runStream = S.runStream
|
||||
|
||||
{-# INLINE toList #-}
|
||||
toList :: Monad m => Stream m Int -> m [Int]
|
||||
|
||||
{-# INLINE head #-}
|
||||
{-# INLINE last #-}
|
||||
{-# INLINE maximum #-}
|
||||
{-# INLINE minimum #-}
|
||||
{-# INLINE find #-}
|
||||
{-# INLINE findIndex #-}
|
||||
{-# INLINE elemIndex #-}
|
||||
{-# INLINE foldl1'Reduce #-}
|
||||
head, last, minimum, maximum, find, findIndex, elemIndex, foldl1'Reduce
|
||||
:: Monad m => Stream m Int -> m (Maybe Int)
|
||||
|
||||
{-# INLINE minimumBy #-}
|
||||
{-# INLINE maximumBy #-}
|
||||
minimumBy, maximumBy :: Monad m => Stream m Int -> m (Maybe Int)
|
||||
|
||||
{-# INLINE foldl'Reduce #-}
|
||||
{-# INLINE foldl'ReduceMap #-}
|
||||
{-# INLINE foldlM'Reduce #-}
|
||||
{-# INLINE foldrMReduce #-}
|
||||
{-# INLINE length #-}
|
||||
{-# INLINE sum #-}
|
||||
{-# INLINE product #-}
|
||||
foldl'Reduce, foldl'ReduceMap, foldlM'Reduce, foldrMReduce, length, sum, product
|
||||
:: Monad m
|
||||
=> Stream m Int -> m Int
|
||||
|
||||
{-# INLINE foldl'Build #-}
|
||||
{-# INLINE foldlM'Build #-}
|
||||
{-# INLINE foldrMBuild #-}
|
||||
foldrMBuild, foldl'Build, foldlM'Build
|
||||
:: Monad m
|
||||
=> Stream m Int -> m [Int]
|
||||
|
||||
{-# INLINE all #-}
|
||||
{-# INLINE any #-}
|
||||
{-# INLINE and #-}
|
||||
{-# INLINE or #-}
|
||||
{-# INLINE null #-}
|
||||
{-# INLINE elem #-}
|
||||
{-# INLINE notElem #-}
|
||||
null, elem, notElem, all, any, and, or :: Monad m => Stream m Int -> m Bool
|
||||
|
||||
{-# INLINE toNull #-}
|
||||
toNull :: Monad m => (t m a -> S.SerialT m a) -> t m a -> m ()
|
||||
toNull t = runStream . t
|
||||
|
||||
{-# INLINE uncons #-}
|
||||
uncons :: Monad m => Stream m Int -> m ()
|
||||
uncons s = do
|
||||
r <- S.uncons s
|
||||
case r of
|
||||
Nothing -> return ()
|
||||
Just (_, t) -> uncons t
|
||||
|
||||
{-# INLINE init #-}
|
||||
init :: Monad m => Stream m a -> m ()
|
||||
init s = S.init s >>= Prelude.mapM_ S.runStream
|
||||
|
||||
{-# INLINE tail #-}
|
||||
tail :: Monad m => Stream m a -> m ()
|
||||
tail s = S.tail s >>= Prelude.mapM_ tail
|
||||
|
||||
{-# INLINE nullHeadTail #-}
|
||||
nullHeadTail :: Monad m => Stream m Int -> m ()
|
||||
nullHeadTail s = do
|
||||
r <- S.null s
|
||||
when (not r) $ do
|
||||
_ <- S.head s
|
||||
S.tail s >>= Prelude.mapM_ nullHeadTail
|
||||
|
||||
{-# INLINE mapM_ #-}
|
||||
mapM_ :: Monad m => Stream m Int -> m ()
|
||||
mapM_ = S.mapM_ (\_ -> return ())
|
||||
|
||||
toList = S.toList
|
||||
|
||||
{-# INLINE toRevList #-}
|
||||
toRevList :: Monad m => Stream m Int -> m [Int]
|
||||
toRevList = S.toRevList
|
||||
|
||||
foldrMBuild = S.foldrM (\x xs -> xs >>= return . (x :)) (return [])
|
||||
foldl'Build = S.foldl' (flip (:)) []
|
||||
foldlM'Build = S.foldlM' (\xs x -> return $ x : xs) []
|
||||
|
||||
foldrMReduce = S.foldrM (\x xs -> xs >>= return . (x +)) (return 0)
|
||||
foldl'Reduce = S.foldl' (+) 0
|
||||
foldl'ReduceMap = P.fmap (+1) . S.foldl' (+) 0
|
||||
foldl1'Reduce = S.foldl1' (+)
|
||||
foldlM'Reduce = S.foldlM' (\xs a -> return $ a + xs) 0
|
||||
|
||||
last = S.last
|
||||
null = S.null
|
||||
head = S.head
|
||||
elem = S.elem maxValue
|
||||
notElem = S.notElem maxValue
|
||||
length = S.length
|
||||
all = S.all (<= maxValue)
|
||||
any = S.any (> maxValue)
|
||||
and = S.and . S.map (<= maxValue)
|
||||
or = S.or . S.map (> maxValue)
|
||||
find = S.find (== maxValue)
|
||||
findIndex = S.findIndex (== maxValue)
|
||||
elemIndex = S.elemIndex maxValue
|
||||
maximum = S.maximum
|
||||
minimum = S.minimum
|
||||
sum = S.sum
|
||||
product = S.product
|
||||
minimumBy = S.minimumBy compare
|
||||
maximumBy = S.maximumBy compare
|
||||
-}
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Transformation
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-
|
||||
{-# INLINE transform #-}
|
||||
transform :: Stream a -> Stream a
|
||||
transform = P.id
|
||||
-}
|
||||
|
||||
{-# INLINE composeN #-}
|
||||
composeN :: P.Monad m
|
||||
=> Int -> (Stream Int -> m (Stream Int)) -> Stream Int -> m (Stream Int)
|
||||
composeN n f x =
|
||||
case n of
|
||||
1 -> f x
|
||||
2 -> f x P.>>= f
|
||||
3 -> f x P.>>= f P.>>= f
|
||||
4 -> f x P.>>= f P.>>= f P.>>= f
|
||||
_ -> undefined
|
||||
|
||||
{-# INLINE scanl' #-}
|
||||
{-# INLINE scanl1' #-}
|
||||
{-# INLINE map #-}
|
||||
{-
|
||||
{-# INLINE fmap #-}
|
||||
{-# INLINE mapMaybe #-}
|
||||
{-# INLINE filterEven #-}
|
||||
{-# INLINE filterAllOut #-}
|
||||
{-# INLINE filterAllIn #-}
|
||||
{-# INLINE takeOne #-}
|
||||
{-# INLINE takeAll #-}
|
||||
{-# INLINE takeWhileTrue #-}
|
||||
{-# INLINE takeWhileMTrue #-}
|
||||
{-# INLINE dropOne #-}
|
||||
{-# INLINE dropAll #-}
|
||||
{-# INLINE dropWhileTrue #-}
|
||||
{-# INLINE dropWhileMTrue #-}
|
||||
{-# INLINE dropWhileFalse #-}
|
||||
{-# INLINE findIndices #-}
|
||||
{-# INLINE elemIndices #-}
|
||||
{-# INLINE insertBy #-}
|
||||
{-# INLINE deleteBy #-}
|
||||
{-# INLINE reverse #-}
|
||||
{-# INLINE foldrS #-}
|
||||
{-# INLINE foldrSMap #-}
|
||||
{-# INLINE foldrT #-}
|
||||
{-# INLINE foldrTMap #-}
|
||||
-}
|
||||
scanl' , scanl1', map{-, fmap, mapMaybe, filterEven, filterAllOut,
|
||||
filterAllIn, takeOne, takeAll, takeWhileTrue, takeWhileMTrue, dropOne,
|
||||
dropAll, dropWhileTrue, dropWhileMTrue, dropWhileFalse,
|
||||
findIndices, elemIndices, insertBy, deleteBy, reverse,
|
||||
foldrS, foldrSMap, foldrT, foldrTMap -}
|
||||
:: PrimMonad m => Int -> Stream Int -> m (Stream Int)
|
||||
|
||||
{-
|
||||
{-# INLINE mapMaybeM #-}
|
||||
mapMaybeM :: S.MonadAsync m => Int -> Stream m Int -> m ()
|
||||
|
||||
{-# INLINE mapM #-}
|
||||
{-# INLINE map' #-}
|
||||
{-# INLINE fmap' #-}
|
||||
mapM, map' :: (S.IsStream t, S.MonadAsync m)
|
||||
=> (t m Int -> S.SerialT m Int) -> Int -> t m Int -> m ()
|
||||
|
||||
fmap' :: (S.IsStream t, S.MonadAsync m, P.Functor (t m))
|
||||
=> (t m Int -> S.SerialT m Int) -> Int -> t m Int -> m ()
|
||||
|
||||
{-# INLINE sequence #-}
|
||||
sequence :: (S.IsStream t, S.MonadAsync m)
|
||||
=> (t m Int -> S.SerialT m Int) -> t m (m Int) -> m ()
|
||||
-}
|
||||
|
||||
{-# INLINE onArray #-}
|
||||
onArray
|
||||
:: PrimMonad m => (S.SerialT m Int -> S.SerialT m Int)
|
||||
-> Stream Int
|
||||
-> m (Stream Int)
|
||||
onArray f arr = S.fold (A.writeN value) $ f $ (S.unfold A.read arr)
|
||||
|
||||
scanl' n = composeN n $ onArray $ S.scanl' (+) 0
|
||||
scanl1' n = composeN n $ onArray $ S.scanl1' (+)
|
||||
map n = composeN n $ onArray $ S.map (+1)
|
||||
-- map n = composeN n $ A.map (+1)
|
||||
{-
|
||||
fmap n = composeN n $ Prelude.fmap (+1)
|
||||
fmap' t n = composeN' n $ t . Prelude.fmap (+1)
|
||||
map' t n = composeN' n $ t . S.map (+1)
|
||||
mapM t n = composeN' n $ t . S.mapM return
|
||||
mapMaybe n = composeN n $ S.mapMaybe
|
||||
(\x -> if Prelude.odd x then Nothing else Just x)
|
||||
mapMaybeM n = composeN n $ S.mapMaybeM
|
||||
(\x -> if Prelude.odd x then return Nothing else return $ Just x)
|
||||
sequence t = transform . t . S.sequence
|
||||
filterEven n = composeN n $ S.filter even
|
||||
filterAllOut n = composeN n $ S.filter (> maxValue)
|
||||
filterAllIn n = composeN n $ S.filter (<= maxValue)
|
||||
takeOne n = composeN n $ S.take 1
|
||||
takeAll n = composeN n $ S.take maxValue
|
||||
takeWhileTrue n = composeN n $ S.takeWhile (<= maxValue)
|
||||
takeWhileMTrue n = composeN n $ S.takeWhileM (return . (<= maxValue))
|
||||
dropOne n = composeN n $ S.drop 1
|
||||
dropAll n = composeN n $ S.drop maxValue
|
||||
dropWhileTrue n = composeN n $ S.dropWhile (<= maxValue)
|
||||
dropWhileMTrue n = composeN n $ S.dropWhileM (return . (<= maxValue))
|
||||
dropWhileFalse n = composeN n $ S.dropWhile (> maxValue)
|
||||
findIndices n = composeN n $ S.findIndices (== maxValue)
|
||||
elemIndices n = composeN n $ S.elemIndices maxValue
|
||||
insertBy n = composeN n $ S.insertBy compare maxValue
|
||||
deleteBy n = composeN n $ S.deleteBy (>=) maxValue
|
||||
reverse n = composeN n $ S.reverse
|
||||
foldrS n = composeN n $ S.foldrS S.cons S.nil
|
||||
foldrSMap n = composeN n $ S.foldrS (\x xs -> x + 1 `S.cons` xs) S.nil
|
||||
foldrT n = composeN n $ S.foldrT S.cons S.nil
|
||||
foldrTMap n = composeN n $ S.foldrT (\x xs -> x + 1 `S.cons` xs) S.nil
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Iteration
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
iterStreamLen, maxIters :: Int
|
||||
iterStreamLen = 10
|
||||
maxIters = 10000
|
||||
|
||||
{-# INLINE iterateSource #-}
|
||||
iterateSource
|
||||
:: S.MonadAsync m
|
||||
=> (Stream m Int -> Stream m Int) -> Int -> Int -> Stream m Int
|
||||
iterateSource g i n = f i (sourceUnfoldrMN iterStreamLen n)
|
||||
where
|
||||
f (0 :: Int) m = g m
|
||||
f x m = g (f (x P.- 1) m)
|
||||
|
||||
{-# INLINE iterateMapM #-}
|
||||
{-# INLINE iterateScan #-}
|
||||
{-# INLINE iterateScanl1 #-}
|
||||
{-# INLINE iterateFilterEven #-}
|
||||
{-# INLINE iterateTakeAll #-}
|
||||
{-# INLINE iterateDropOne #-}
|
||||
{-# INLINE iterateDropWhileFalse #-}
|
||||
{-# INLINE iterateDropWhileTrue #-}
|
||||
iterateMapM, iterateScan, iterateScanl1, iterateFilterEven, iterateTakeAll,
|
||||
iterateDropOne, iterateDropWhileFalse, iterateDropWhileTrue
|
||||
:: S.MonadAsync m
|
||||
=> Int -> Stream m Int
|
||||
|
||||
-- this is quadratic
|
||||
iterateScan = iterateSource (S.scanl' (+) 0) (maxIters `div` 10)
|
||||
-- so is this
|
||||
iterateScanl1 = iterateSource (S.scanl1' (+)) (maxIters `div` 10)
|
||||
|
||||
iterateMapM = iterateSource (S.mapM return) maxIters
|
||||
iterateFilterEven = iterateSource (S.filter even) maxIters
|
||||
iterateTakeAll = iterateSource (S.take maxValue) maxIters
|
||||
iterateDropOne = iterateSource (S.drop 1) maxIters
|
||||
iterateDropWhileFalse = iterateSource (S.dropWhile (> maxValue)) maxIters
|
||||
iterateDropWhileTrue = iterateSource (S.dropWhile (<= maxValue)) maxIters
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Zipping and concat
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE zip #-}
|
||||
{-# INLINE zipM #-}
|
||||
{-# INLINE mergeBy #-}
|
||||
zip, zipM, mergeBy :: Monad m => Stream m Int -> m ()
|
||||
|
||||
zip src = do
|
||||
r <- S.tail src
|
||||
let src1 = fromJust r
|
||||
transform (S.zipWith (,) src src1)
|
||||
zipM src = do
|
||||
r <- S.tail src
|
||||
let src1 = fromJust r
|
||||
transform (S.zipWithM (curry return) src src1)
|
||||
|
||||
mergeBy src = do
|
||||
r <- S.tail src
|
||||
let src1 = fromJust r
|
||||
transform (S.mergeBy P.compare src src1)
|
||||
|
||||
{-# INLINE isPrefixOf #-}
|
||||
{-# INLINE isSubsequenceOf #-}
|
||||
isPrefixOf, isSubsequenceOf :: Monad m => Stream m Int -> m Bool
|
||||
|
||||
isPrefixOf src = S.isPrefixOf src src
|
||||
isSubsequenceOf src = S.isSubsequenceOf src src
|
||||
|
||||
{-# INLINE stripPrefix #-}
|
||||
stripPrefix :: Monad m => Stream m Int -> m ()
|
||||
stripPrefix src = do
|
||||
_ <- S.stripPrefix src src
|
||||
return ()
|
||||
|
||||
{-# INLINE zipAsync #-}
|
||||
{-# INLINE zipAsyncM #-}
|
||||
{-# INLINE zipAsyncAp #-}
|
||||
zipAsync, zipAsyncAp, zipAsyncM :: S.MonadAsync m => Stream m Int -> m ()
|
||||
|
||||
zipAsync src = do
|
||||
r <- S.tail src
|
||||
let src1 = fromJust r
|
||||
transform (S.zipAsyncWith (,) src src1)
|
||||
|
||||
zipAsyncM src = do
|
||||
r <- S.tail src
|
||||
let src1 = fromJust r
|
||||
transform (S.zipAsyncWithM (curry return) src src1)
|
||||
|
||||
zipAsyncAp src = do
|
||||
r <- S.tail src
|
||||
let src1 = fromJust r
|
||||
transform (S.zipAsyncly $ (,) <$> S.serially src
|
||||
<*> S.serially src1)
|
||||
|
||||
{-# INLINE eqBy #-}
|
||||
eqBy :: (Monad m, P.Eq a) => Stream m a -> m P.Bool
|
||||
eqBy src = S.eqBy (==) src src
|
||||
|
||||
{-# INLINE cmpBy #-}
|
||||
cmpBy :: (Monad m, P.Ord a) => Stream m a -> m P.Ordering
|
||||
cmpBy src = S.cmpBy P.compare src src
|
||||
|
||||
concatStreamLen, maxNested :: Int
|
||||
concatStreamLen = 1
|
||||
maxNested = 100000
|
||||
|
||||
{-# INLINE concatMap #-}
|
||||
concatMap :: S.MonadAsync m => Int -> Stream m Int
|
||||
concatMap n = S.concatMap (\_ -> sourceUnfoldrMN maxNested n)
|
||||
(sourceUnfoldrMN concatStreamLen n)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Mixed Composition
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE scanMap #-}
|
||||
{-# INLINE dropMap #-}
|
||||
{-# INLINE dropScan #-}
|
||||
{-# INLINE takeDrop #-}
|
||||
{-# INLINE takeScan #-}
|
||||
{-# INLINE takeMap #-}
|
||||
{-# INLINE filterDrop #-}
|
||||
{-# INLINE filterTake #-}
|
||||
{-# INLINE filterScan #-}
|
||||
{-# INLINE filterScanl1 #-}
|
||||
{-# INLINE filterMap #-}
|
||||
scanMap, dropMap, dropScan, takeDrop, takeScan, takeMap, filterDrop,
|
||||
filterTake, filterScan, filterScanl1, filterMap
|
||||
:: Monad m => Int -> Stream m Int -> m ()
|
||||
|
||||
scanMap n = composeN n $ S.map (subtract 1) . S.scanl' (+) 0
|
||||
dropMap n = composeN n $ S.map (subtract 1) . S.drop 1
|
||||
dropScan n = composeN n $ S.scanl' (+) 0 . S.drop 1
|
||||
takeDrop n = composeN n $ S.drop 1 . S.take maxValue
|
||||
takeScan n = composeN n $ S.scanl' (+) 0 . S.take maxValue
|
||||
takeMap n = composeN n $ S.map (subtract 1) . S.take maxValue
|
||||
filterDrop n = composeN n $ S.drop 1 . S.filter (<= maxValue)
|
||||
filterTake n = composeN n $ S.take maxValue . S.filter (<= maxValue)
|
||||
filterScan n = composeN n $ S.scanl' (+) 0 . S.filter (<= maxBound)
|
||||
filterScanl1 n = composeN n $ S.scanl1' (+) . S.filter (<= maxBound)
|
||||
filterMap n = composeN n $ S.map (subtract 1) . S.filter (<= maxValue)
|
||||
|
||||
data Pair a b = Pair !a !b deriving (Generic, NFData)
|
||||
|
||||
{-# INLINE sumProductFold #-}
|
||||
sumProductFold :: Monad m => Stream m Int -> m (Int, Int)
|
||||
sumProductFold = S.foldl' (\(s,p) x -> (s + x, p P.* x)) (0,1)
|
||||
|
||||
{-# INLINE sumProductScan #-}
|
||||
sumProductScan :: Monad m => Stream m Int -> m (Pair Int Int)
|
||||
sumProductScan = S.foldl' (\(Pair _ p) (s0,x) -> Pair s0 (p P.* x)) (Pair 0 1)
|
||||
. S.scanl' (\(s,_) x -> (s + x,x)) (0,0)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Pure stream operations
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-}
|
||||
{-# INLINE eqInstance #-}
|
||||
eqInstance :: Stream Int -> Bool
|
||||
eqInstance src = src == src
|
||||
|
||||
{-# INLINE eqInstanceNotEq #-}
|
||||
eqInstanceNotEq :: Stream Int -> Bool
|
||||
eqInstanceNotEq src = src P./= src
|
||||
|
||||
{-# INLINE ordInstance #-}
|
||||
ordInstance :: Stream Int -> Bool
|
||||
ordInstance src = src P.< src
|
||||
|
||||
{-# INLINE ordInstanceMin #-}
|
||||
ordInstanceMin :: Stream Int -> Stream Int
|
||||
ordInstanceMin src = P.min src src
|
||||
|
||||
{-# INLINE showInstance #-}
|
||||
showInstance :: Stream Int -> P.String
|
||||
showInstance src = P.show src
|
||||
|
||||
{-# INLINE readInstance #-}
|
||||
readInstance :: P.String -> Stream Int
|
||||
readInstance str =
|
||||
let r = P.reads str
|
||||
in case r of
|
||||
[(x,"")] -> x
|
||||
_ -> P.error "readInstance: no parse"
|
||||
|
||||
{-# INLINE pureFoldl' #-}
|
||||
pureFoldl' :: PrimMonad m => Stream Int -> m Int
|
||||
pureFoldl' = S.foldl' (+) 0 . S.unfold A.read
|
||||
|
||||
#ifdef DEVBUILD
|
||||
{-# INLINE foldableFoldl' #-}
|
||||
foldableFoldl' :: Stream Int -> Int
|
||||
foldableFoldl' = F.foldl' (+) 0
|
||||
|
||||
{-# INLINE foldableSum #-}
|
||||
foldableSum :: Stream Int -> Int
|
||||
foldableSum = P.sum
|
||||
#endif
|
||||
|
||||
{-
|
||||
{-# INLINE traversableMapM #-}
|
||||
traversableMapM :: Stream Identity Int -> IO (Stream Identity Int)
|
||||
traversableMapM = P.mapM return
|
||||
-}
|
||||
@ -376,6 +376,7 @@ benchmark Data.Prim.Array
|
||||
hs-source-dirs: .
|
||||
main-is: Streamly/Benchmark/Data/Prim/Array.hs
|
||||
other-modules: Streamly.Benchmark.Data.Prim.ArrayOps
|
||||
build-depends: primitive
|
||||
|
||||
benchmark Data.SmallArray
|
||||
import: bench-options
|
||||
@ -384,6 +385,14 @@ benchmark Data.SmallArray
|
||||
main-is: Streamly/Benchmark/Data/SmallArray.hs
|
||||
other-modules: Streamly.Benchmark.Data.SmallArrayOps
|
||||
|
||||
benchmark Data.Prim.Pinned.Array
|
||||
import: bench-options
|
||||
type: exitcode-stdio-1.0
|
||||
hs-source-dirs: .
|
||||
main-is: Streamly/Benchmark/Data/Prim/Pinned/Array.hs
|
||||
other-modules: Streamly.Benchmark.Data.Prim.Pinned.ArrayOps
|
||||
build-depends: primitive
|
||||
|
||||
benchmark Memory.Array
|
||||
import: bench-options
|
||||
type: exitcode-stdio-1.0
|
||||
|
||||
@ -8,8 +8,7 @@
|
||||
-- Portability : GHC
|
||||
--
|
||||
module Streamly.Data.Prim.Array
|
||||
( PrimArray
|
||||
, Prim
|
||||
( Array
|
||||
|
||||
-- * Construction
|
||||
, A.fromListN
|
||||
@ -39,6 +38,6 @@ module Streamly.Data.Prim.Array
|
||||
)
|
||||
where
|
||||
|
||||
import Streamly.Internal.Data.Prim.Array (PrimArray, Prim)
|
||||
import Streamly.Internal.Data.Prim.Array (Array)
|
||||
|
||||
import qualified Streamly.Internal.Data.Prim.Array as A
|
||||
|
||||
@ -6,196 +6,105 @@
|
||||
-- Module : Streamly.Internal.Data.Prim.Array
|
||||
-- Copyright : (c) 2019 Composewell Technologies
|
||||
--
|
||||
-- License : BSD-3-Clause
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
-- Stability : experimental
|
||||
-- Portability : GHC
|
||||
--
|
||||
module Streamly.Internal.Data.Prim.Array
|
||||
(
|
||||
Array
|
||||
|
||||
-- XXX should it be just Array instead? We should be able to replace one
|
||||
-- array type with another easily.
|
||||
PrimArray(..)
|
||||
-- * Construction
|
||||
|
||||
-- XXX Prim should be exported from Data.Prim module?
|
||||
, Prim(..)
|
||||
|
||||
, foldl'
|
||||
, foldr
|
||||
|
||||
, length
|
||||
|
||||
, writeN
|
||||
, write
|
||||
|
||||
, toStreamD
|
||||
, toStreamDRev
|
||||
|
||||
, toStream
|
||||
, toStreamRev
|
||||
, read
|
||||
, readSlice
|
||||
|
||||
, fromListN
|
||||
, fromList
|
||||
, fromStreamDN
|
||||
, fromStreamD
|
||||
-- Pure List APIs
|
||||
, A.fromListN
|
||||
, A.fromList
|
||||
|
||||
-- Stream Folds
|
||||
, fromStreamN
|
||||
, fromStream
|
||||
|
||||
-- Monadic APIs
|
||||
-- , newArray
|
||||
, A.writeN -- drop new
|
||||
, A.write -- full buffer
|
||||
|
||||
-- * Elimination
|
||||
|
||||
, A.toList
|
||||
, toStream
|
||||
, toStreamRev
|
||||
, read
|
||||
, unsafeRead
|
||||
|
||||
-- * Random Access
|
||||
, length
|
||||
, null
|
||||
, last
|
||||
-- , (!!)
|
||||
, readIndex
|
||||
, A.unsafeIndex
|
||||
-- , readIndices
|
||||
-- , readRanges
|
||||
|
||||
-- , readFrom -- read from a given position to the end of file
|
||||
-- , readFromRev -- read from a given position to the beginning of file
|
||||
-- , readTo -- read from beginning up to the given position
|
||||
-- , readToRev -- read from end to the given position in file
|
||||
-- , readFromTo
|
||||
-- , readFromThenTo
|
||||
|
||||
-- , readChunksOfFrom
|
||||
-- , ...
|
||||
|
||||
-- , writeIndex
|
||||
-- , writeFrom -- start writing at the given position
|
||||
-- , writeFromRev
|
||||
-- , writeTo -- write from beginning up to the given position
|
||||
-- , writeToRev
|
||||
-- , writeFromTo
|
||||
-- , writeFromThenTo
|
||||
--
|
||||
-- , writeChunksOfFrom
|
||||
-- , ...
|
||||
|
||||
-- , writeIndex
|
||||
-- , writeIndices
|
||||
-- , writeRanges
|
||||
|
||||
-- -- * Search
|
||||
-- , bsearch
|
||||
-- , bsearchIndex
|
||||
-- , find
|
||||
-- , findIndex
|
||||
-- , findIndices
|
||||
|
||||
-- -- * In-pace mutation (for Mutable Array type)
|
||||
-- , partitionBy
|
||||
-- , shuffleBy
|
||||
-- , foldtWith
|
||||
-- , foldbWith
|
||||
|
||||
-- * Immutable Transformations
|
||||
-- , streamTransform
|
||||
|
||||
-- * Folding Arrays
|
||||
, streamFold
|
||||
, fold
|
||||
|
||||
-- * Folds with Array as the container
|
||||
-- , D.lastN
|
||||
|
||||
-- * Streaming array operations
|
||||
|
||||
, concat
|
||||
, compact
|
||||
|
||||
)
|
||||
where
|
||||
|
||||
import Prelude hiding (foldr, length, read)
|
||||
import Control.DeepSeq (NFData(..))
|
||||
import Control.Monad (when)
|
||||
import Control.Monad.IO.Class (liftIO, MonadIO)
|
||||
import GHC.IO (unsafePerformIO)
|
||||
import Data.Primitive.Types (Prim(..))
|
||||
import Streamly.Internal.Data.Prim.Array.Types (Array(..), length)
|
||||
import qualified Streamly.Internal.Data.Prim.Array.Types as A
|
||||
|
||||
import Streamly.Internal.Data.Prim.Array.Types
|
||||
import Streamly.Internal.Data.Unfold.Types (Unfold(..))
|
||||
import Streamly.Internal.Data.Fold.Types (Fold(..))
|
||||
import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
|
||||
import Streamly.Internal.Data.Stream.Serial (SerialT)
|
||||
|
||||
import qualified Streamly.Internal.Data.Stream.StreamD as D
|
||||
|
||||
{-# INLINE_NORMAL toStreamD #-}
|
||||
toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
|
||||
toStreamD arr = D.Stream step 0
|
||||
where
|
||||
{-# INLINE_LATE step #-}
|
||||
step _ i
|
||||
| i == sizeofPrimArray arr = return D.Stop
|
||||
step _ i = return $ D.Yield (indexPrimArray arr i) (i + 1)
|
||||
|
||||
{-# INLINE length #-}
|
||||
length :: Prim a => PrimArray a -> Int
|
||||
length = sizeofPrimArray
|
||||
|
||||
{-# INLINE_NORMAL toStreamDRev #-}
|
||||
toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
|
||||
toStreamDRev arr = D.Stream step (sizeofPrimArray arr - 1)
|
||||
where
|
||||
{-# INLINE_LATE step #-}
|
||||
step _ i
|
||||
| i < 0 = return D.Stop
|
||||
step _ i = return $ D.Yield (indexPrimArray arr i) (i - 1)
|
||||
|
||||
{-# INLINE_NORMAL foldl' #-}
|
||||
foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
|
||||
foldl' = foldlPrimArray'
|
||||
|
||||
{-# INLINE_NORMAL foldr #-}
|
||||
foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
|
||||
foldr = foldrPrimArray
|
||||
|
||||
-- writeN n = S.evertM (fromStreamDN n)
|
||||
{-# INLINE_NORMAL writeN #-}
|
||||
writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
|
||||
writeN limit = Fold step initial extract
|
||||
where
|
||||
initial = do
|
||||
marr <- liftIO $ newPrimArray limit
|
||||
return (marr, 0)
|
||||
step (marr, i) x
|
||||
| i == limit = return (marr, i)
|
||||
| otherwise = do
|
||||
liftIO $ writePrimArray marr i x
|
||||
return (marr, i + 1)
|
||||
extract (marr, _) = liftIO $ unsafeFreezePrimArray marr
|
||||
|
||||
{-# INLINE_NORMAL write #-}
|
||||
write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
|
||||
write = Fold step initial extract
|
||||
where
|
||||
initial = do
|
||||
marr <- liftIO $ newPrimArray 0
|
||||
return (marr, 0, 0)
|
||||
step (marr, i, capacity) x
|
||||
| i == capacity =
|
||||
let newCapacity = max (capacity * 2) 1
|
||||
in do newMarr <- liftIO $ resizeMutablePrimArray marr newCapacity
|
||||
liftIO $ writePrimArray newMarr i x
|
||||
return (newMarr, i + 1, newCapacity)
|
||||
| otherwise = do
|
||||
liftIO $ writePrimArray marr i x
|
||||
return (marr, i + 1, capacity)
|
||||
extract (marr, len, _) = do liftIO $ shrinkMutablePrimArray marr len
|
||||
liftIO $ unsafeFreezePrimArray marr
|
||||
|
||||
{-# INLINE_NORMAL fromStreamDN #-}
|
||||
fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
|
||||
fromStreamDN limit str = do
|
||||
marr <- liftIO $ newPrimArray (max limit 0)
|
||||
_ <-
|
||||
D.foldlM'
|
||||
(\i x -> i `seq` liftIO (writePrimArray marr i x) >> return (i + 1))
|
||||
(return 0) $
|
||||
D.take limit str
|
||||
liftIO $ unsafeFreezePrimArray marr
|
||||
|
||||
{-# INLINE fromStreamD #-}
|
||||
fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
|
||||
fromStreamD = D.runFold write
|
||||
|
||||
{-# INLINABLE fromListN #-}
|
||||
fromListN :: Prim a => Int -> [a] -> PrimArray a
|
||||
fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
|
||||
|
||||
{-# INLINABLE fromList #-}
|
||||
fromList :: Prim a => [a] -> PrimArray a
|
||||
fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
|
||||
|
||||
instance Prim a => NFData (PrimArray a) where
|
||||
{-# INLINE rnf #-}
|
||||
rnf = foldl' (\_ _ -> ()) ()
|
||||
|
||||
{-# INLINE fromStreamN #-}
|
||||
fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
|
||||
fromStreamN n m = do
|
||||
when (n < 0) $ error "fromStreamN: negative write count specified"
|
||||
fromStreamDN n $ D.toStreamD m
|
||||
|
||||
{-# INLINE fromStream #-}
|
||||
fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
|
||||
fromStream m = fromStreamD $ D.toStreamD m
|
||||
|
||||
{-# INLINE_EARLY toStream #-}
|
||||
toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
|
||||
toStream = D.fromStreamD . toStreamD
|
||||
|
||||
{-# INLINE_EARLY toStreamRev #-}
|
||||
toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
|
||||
toStreamRev = D.fromStreamD . toStreamDRev
|
||||
|
||||
{-# INLINE fold #-}
|
||||
fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
|
||||
fold f arr = D.runFold f (toStreamD arr)
|
||||
|
||||
{-# INLINE streamFold #-}
|
||||
streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
|
||||
streamFold f arr = f (toStream arr)
|
||||
|
||||
{-# INLINE_NORMAL read #-}
|
||||
read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
|
||||
read = Unfold step inject
|
||||
where
|
||||
inject arr = return (arr, 0)
|
||||
step (arr, i)
|
||||
| i == length arr = return D.Stop
|
||||
step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)
|
||||
|
||||
{-# INLINE_NORMAL readSlice #-}
|
||||
readSlice :: (Prim a, Monad m) => Int -> Int -> Unfold m (PrimArray a) a
|
||||
readSlice off len = Unfold step inject
|
||||
where
|
||||
inject arr = return (arr, off)
|
||||
step (arr, i)
|
||||
| i == min (off + len) (length arr) = return D.Stop
|
||||
step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)
|
||||
#include "prim-array.hs"
|
||||
|
||||
@ -1,207 +1,98 @@
|
||||
{-# LANGUAGE UnboxedTuples #-}
|
||||
|
||||
#include "inline.hs"
|
||||
|
||||
-- |
|
||||
-- Module : Streamly.Internal.Data.Prim.Array.Types
|
||||
-- Copyright : (c) Roman Leshchinskiy 2009-2012
|
||||
-- License : BSD-style
|
||||
-- Copyright : (c) 2019 Composewell Technologies
|
||||
--
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
-- Portability : non-portable
|
||||
-- Stability : experimental
|
||||
-- Portability : GHC
|
||||
--
|
||||
-- Arrays of unboxed primitive types. The function provided by this module
|
||||
-- match the behavior of those provided by @Data.Primitive.ByteArray@, and
|
||||
-- the underlying types and primops that back them are the same.
|
||||
-- However, the type constructors 'PrimArray' and 'MutablePrimArray' take one additional
|
||||
-- argument than their respective counterparts 'ByteArray' and 'MutableByteArray'.
|
||||
-- This argument is used to designate the type of element in the array.
|
||||
-- Consequently, all function this modules accepts length and incides in
|
||||
-- terms of elements, not bytes.
|
||||
--
|
||||
-- @since 0.6.4.0
|
||||
module Streamly.Internal.Data.Prim.Array.Types
|
||||
( -- * Types
|
||||
PrimArray(..)
|
||||
, MutablePrimArray(..)
|
||||
-- * Allocation
|
||||
, newPrimArray
|
||||
, resizeMutablePrimArray
|
||||
, shrinkMutablePrimArray
|
||||
-- * Element Access
|
||||
, writePrimArray
|
||||
, indexPrimArray
|
||||
-- * Freezing and Thawing
|
||||
, unsafeFreezePrimArray
|
||||
-- * Information
|
||||
, sizeofPrimArray
|
||||
-- * Folding
|
||||
, foldrPrimArray
|
||||
, foldlPrimArray'
|
||||
) where
|
||||
(
|
||||
Array (..)
|
||||
, unsafeFreeze
|
||||
, unsafeFreezeWithShrink
|
||||
-- , unsafeThaw
|
||||
, defaultChunkSize
|
||||
, nil
|
||||
|
||||
import GHC.Exts
|
||||
-- * Construction
|
||||
, spliceTwo
|
||||
|
||||
import Data.Primitive.Types
|
||||
import Data.Primitive.ByteArray (ByteArray(..))
|
||||
import Control.Monad.Primitive
|
||||
import qualified Data.Primitive.ByteArray as PB
|
||||
, fromList
|
||||
, fromListN
|
||||
, fromStreamDN
|
||||
, fromStreamD
|
||||
|
||||
-- | Arrays of unboxed elements. This accepts types like 'Double', 'Char',
|
||||
-- 'Int', and 'Word', as well as their fixed-length variants ('Word8',
|
||||
-- 'Word16', etc.). Since the elements are unboxed, a 'PrimArray' is strict
|
||||
-- in its elements. This differs from the behavior of 'Array', which is lazy
|
||||
-- in its elements.
|
||||
data PrimArray a = PrimArray ByteArray#
|
||||
-- * Streams of arrays
|
||||
, fromStreamDArraysOf
|
||||
, FlattenState (..) -- for inspection testing
|
||||
, flattenArrays
|
||||
, flattenArraysRev
|
||||
, SpliceState (..) -- for inspection testing
|
||||
, packArraysChunksOf
|
||||
, lpackArraysChunksOf
|
||||
#if !defined(mingw32_HOST_OS)
|
||||
-- , groupIOVecsOf
|
||||
#endif
|
||||
, splitOn
|
||||
, breakOn
|
||||
|
||||
-- | Mutable primitive arrays associated with a primitive state token.
|
||||
-- These can be written to and read from in a monadic context that supports
|
||||
-- sequencing such as 'IO' or 'ST'. Typically, a mutable primitive array will
|
||||
-- be built and then convert to an immutable primitive array using
|
||||
-- 'unsafeFreezePrimArray'. However, it is also acceptable to simply discard
|
||||
-- a mutable primitive array since it lives in managed memory and will be
|
||||
-- garbage collected when no longer referenced.
|
||||
data MutablePrimArray s a = MutablePrimArray (MutableByteArray# s)
|
||||
-- * Elimination
|
||||
, unsafeIndex
|
||||
, byteLength
|
||||
, length
|
||||
|
||||
sameByteArray :: ByteArray# -> ByteArray# -> Bool
|
||||
sameByteArray ba1 ba2 =
|
||||
case reallyUnsafePtrEquality# (unsafeCoerce# ba1 :: ()) (unsafeCoerce# ba2 :: ()) of
|
||||
r -> isTrue# r
|
||||
, foldl'
|
||||
, foldr
|
||||
, foldr'
|
||||
, foldlM'
|
||||
, splitAt
|
||||
|
||||
-- | @since 0.6.4.0
|
||||
instance (Eq a, Prim a) => Eq (PrimArray a) where
|
||||
a1@(PrimArray ba1#) == a2@(PrimArray ba2#)
|
||||
| sameByteArray ba1# ba2# = True
|
||||
| sz1 /= sz2 = False
|
||||
| otherwise = loop (quot sz1 (sizeOf (undefined :: a)) - 1)
|
||||
where
|
||||
-- Here, we take the size in bytes, not in elements. We do this
|
||||
-- since it allows us to defer performing the division to
|
||||
-- calculate the size in elements.
|
||||
sz1 = PB.sizeofByteArray (ByteArray ba1#)
|
||||
sz2 = PB.sizeofByteArray (ByteArray ba2#)
|
||||
loop !i
|
||||
| i < 0 = True
|
||||
| otherwise = indexPrimArray a1 i == indexPrimArray a2 i && loop (i-1)
|
||||
{-# INLINE (==) #-}
|
||||
, toStreamD
|
||||
, toStreamDRev
|
||||
, toStreamK
|
||||
, toStreamKRev
|
||||
, toList
|
||||
-- , toArrayMinChunk
|
||||
, writeN
|
||||
, MA.ArrayUnsafe(..)
|
||||
, writeNUnsafe
|
||||
, write
|
||||
|
||||
-- | Lexicographic ordering. Subject to change between major versions.
|
||||
--
|
||||
-- @since 0.6.4.0
|
||||
instance (Ord a, Prim a) => Ord (PrimArray a) where
|
||||
compare a1@(PrimArray ba1#) a2@(PrimArray ba2#)
|
||||
| sameByteArray ba1# ba2# = EQ
|
||||
| otherwise = loop 0
|
||||
where
|
||||
cmp LT _ = LT
|
||||
cmp EQ y = y
|
||||
cmp GT _ = GT
|
||||
sz1 = PB.sizeofByteArray (ByteArray ba1#)
|
||||
sz2 = PB.sizeofByteArray (ByteArray ba2#)
|
||||
sz = quot (min sz1 sz2) (sizeOf (undefined :: a))
|
||||
loop !i
|
||||
| i < sz = compare (indexPrimArray a1 i) (indexPrimArray a2 i) `cmp` loop (i+1)
|
||||
| otherwise = compare sz1 sz2
|
||||
{-# INLINE compare #-}
|
||||
|
||||
-- | @since 0.6.4.0
|
||||
instance (Show a, Prim a) => Show (PrimArray a) where
|
||||
showsPrec p a = showParen (p > 10) $
|
||||
showString "fromListN " . shows (sizeofPrimArray a) . showString " "
|
||||
. shows (primArrayToList a)
|
||||
|
||||
-- | Convert the primitive array to a list.
|
||||
{-# INLINE primArrayToList #-}
|
||||
primArrayToList :: forall a. Prim a => PrimArray a -> [a]
|
||||
primArrayToList xs = build (\c n -> foldrPrimArray c n xs)
|
||||
|
||||
-- | Create a new mutable primitive array of the given length. The
|
||||
-- underlying memory is left uninitialized.
|
||||
newPrimArray :: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
|
||||
{-# INLINE newPrimArray #-}
|
||||
newPrimArray (I# n#)
|
||||
= primitive (\s# ->
|
||||
case newByteArray# (n# *# sizeOf# (undefined :: a)) s# of
|
||||
(# s'#, arr# #) -> (# s'#, MutablePrimArray arr# #)
|
||||
, unlines
|
||||
)
|
||||
where
|
||||
|
||||
-- | Resize a mutable primitive array. The new size is given in elements.
|
||||
--
|
||||
-- This will either resize the array in-place or, if not possible, allocate the
|
||||
-- contents into a new, unpinned array and copy the original array\'s contents.
|
||||
--
|
||||
-- To avoid undefined behaviour, the original 'MutablePrimArray' shall not be
|
||||
-- accessed anymore after a 'resizeMutablePrimArray' has been performed.
|
||||
-- Moreover, no reference to the old one should be kept in order to allow
|
||||
-- garbage collection of the original 'MutablePrimArray' in case a new
|
||||
-- 'MutablePrimArray' had to be allocated.
|
||||
resizeMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
|
||||
=> MutablePrimArray (PrimState m) a
|
||||
-> Int -- ^ new size
|
||||
-> m (MutablePrimArray (PrimState m) a)
|
||||
{-# INLINE resizeMutablePrimArray #-}
|
||||
resizeMutablePrimArray (MutablePrimArray arr#) (I# n#)
|
||||
= primitive (\s# -> case resizeMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)) s# of
|
||||
(# s'#, arr'# #) -> (# s'#, MutablePrimArray arr'# #))
|
||||
import qualified Streamly.Internal.Data.Prim.Mutable.Array.Types as MA
|
||||
|
||||
-- Although it is possible to shim resizeMutableByteArray for old GHCs, this
|
||||
-- is not the case with shrinkMutablePrimArray.
|
||||
#include "prim-array-types.hs"
|
||||
|
||||
-- | Shrink a mutable primitive array. The new size is given in elements.
|
||||
-- It must be smaller than the old size. The array will be resized in place.
|
||||
-- This function is only available when compiling with GHC 7.10 or newer.
|
||||
shrinkMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
|
||||
=> MutablePrimArray (PrimState m) a
|
||||
-> Int -- ^ new size
|
||||
-> m ()
|
||||
{-# INLINE shrinkMutablePrimArray #-}
|
||||
shrinkMutablePrimArray (MutablePrimArray arr#) (I# n#)
|
||||
= primitive_ (shrinkMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)))
|
||||
-- Drops the separator byte
|
||||
-- Inefficient compared to Memory Array
|
||||
{-# INLINE breakOn #-}
|
||||
breakOn ::
|
||||
PrimMonad m
|
||||
=> Word8
|
||||
-> Array Word8
|
||||
-> m (Array Word8, Maybe (Array Word8))
|
||||
breakOn sep arr@(Array arr# off len) =
|
||||
case loc of
|
||||
Left _ -> return (arr, Nothing)
|
||||
Right len1 -> do
|
||||
let len2 = len - len1 - 1
|
||||
return (Array arr# off len1, Just $ Array arr# (off + len1 + 1) len2)
|
||||
|
||||
-- | Write an element to the given index.
|
||||
writePrimArray ::
|
||||
(Prim a, PrimMonad m)
|
||||
=> MutablePrimArray (PrimState m) a -- ^ array
|
||||
-> Int -- ^ index
|
||||
-> a -- ^ element
|
||||
-> m ()
|
||||
{-# INLINE writePrimArray #-}
|
||||
writePrimArray (MutablePrimArray arr#) (I# i#) x
|
||||
= primitive_ (writeByteArray# arr# i# x)
|
||||
where
|
||||
|
||||
-- | Convert a mutable byte array to an immutable one without copying. The
|
||||
-- array should not be modified after the conversion.
|
||||
unsafeFreezePrimArray
|
||||
:: PrimMonad m => MutablePrimArray (PrimState m) a -> m (PrimArray a)
|
||||
{-# INLINE unsafeFreezePrimArray #-}
|
||||
unsafeFreezePrimArray (MutablePrimArray arr#)
|
||||
= primitive (\s# -> case unsafeFreezeByteArray# arr# s# of
|
||||
(# s'#, arr'# #) -> (# s'#, PrimArray arr'# #))
|
||||
loc = foldl' chk (Left 0) arr
|
||||
|
||||
-- | Read a primitive value from the primitive array.
|
||||
indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
|
||||
{-# INLINE indexPrimArray #-}
|
||||
indexPrimArray (PrimArray arr#) (I# i#) = indexByteArray# arr# i#
|
||||
|
||||
-- | Get the size, in elements, of the primitive array.
|
||||
sizeofPrimArray :: forall a. Prim a => PrimArray a -> Int
|
||||
{-# INLINE sizeofPrimArray #-}
|
||||
sizeofPrimArray (PrimArray arr#) = I# (quotInt# (sizeofByteArray# arr#) (sizeOf# (undefined :: a)))
|
||||
|
||||
-- | Lazy right-associated fold over the elements of a 'PrimArray'.
|
||||
{-# INLINE foldrPrimArray #-}
|
||||
foldrPrimArray :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
|
||||
foldrPrimArray f z arr = go 0
|
||||
where
|
||||
!sz = sizeofPrimArray arr
|
||||
go !i
|
||||
| sz > i = f (indexPrimArray arr i) (go (i+1))
|
||||
| otherwise = z
|
||||
|
||||
-- | Strict left-associated fold over the elements of a 'PrimArray'.
|
||||
{-# INLINE foldlPrimArray' #-}
|
||||
foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
|
||||
foldlPrimArray' f z0 arr = go 0 z0
|
||||
where
|
||||
!sz = sizeofPrimArray arr
|
||||
go !i !acc
|
||||
| i < sz = go (i + 1) (f acc (indexPrimArray arr i))
|
||||
| otherwise = acc
|
||||
chk (Left i) a =
|
||||
if a == sep
|
||||
then Right i
|
||||
else Left (i + 1)
|
||||
chk r _ = r
|
||||
|
||||
92
src/Streamly/Internal/Data/Prim/Mutable/Array/Types.hs
Normal file
92
src/Streamly/Internal/Data/Prim/Mutable/Array/Types.hs
Normal file
@ -0,0 +1,92 @@
|
||||
{-# LANGUAGE UnboxedTuples #-}
|
||||
|
||||
#include "inline.hs"
|
||||
|
||||
-- |
|
||||
-- Module : Streamly.Internal.Data.Prim.Mutable.Array.Types
|
||||
-- Copyright : (c) 2019 Composewell Technologies
|
||||
--
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
-- Stability : experimental
|
||||
-- Portability : GHC
|
||||
--
|
||||
module Streamly.Internal.Data.Prim.Mutable.Array.Types
|
||||
(
|
||||
Array (..)
|
||||
|
||||
-- * Construction
|
||||
, newArray
|
||||
, unsafeWriteIndex
|
||||
|
||||
, spliceTwo
|
||||
, unsafeCopy
|
||||
|
||||
, fromListM
|
||||
, fromListNM
|
||||
, fromStreamDN
|
||||
, fromStreamD
|
||||
|
||||
-- * Streams of arrays
|
||||
, fromStreamDArraysOf
|
||||
|
||||
, packArraysChunksOf
|
||||
, lpackArraysChunksOf
|
||||
|
||||
#if !defined(mingw32_HOST_OS)
|
||||
-- , groupIOVecsOf
|
||||
#endif
|
||||
|
||||
-- * Elimination
|
||||
, unsafeReadIndex
|
||||
, length
|
||||
, byteLength
|
||||
|
||||
, writeN
|
||||
, ArrayUnsafe(..)
|
||||
, writeNUnsafe
|
||||
, write
|
||||
|
||||
-- * Utilities
|
||||
, resizeArray
|
||||
, shrinkArray
|
||||
)
|
||||
where
|
||||
|
||||
#include "mutable-prim-array-types.hs"
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Allocation (Unpinned)
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Allocate an array that is unpinned and can hold 'count' items. The memory
|
||||
-- of the array is uninitialized.
|
||||
--
|
||||
-- Note that this is internal routine, the reference to this array cannot be
|
||||
-- given out until the array has been written to and frozen.
|
||||
{-# INLINE newArray #-}
|
||||
newArray ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> m (Array (PrimState m) a)
|
||||
newArray (I# n#) =
|
||||
primitive $ \s# ->
|
||||
let bytes = n# *# sizeOf# (undefined :: a)
|
||||
in case newByteArray# bytes s# of
|
||||
(# s1#, arr# #) -> (# s1#, Array arr# #)
|
||||
|
||||
-- | Resize (unpinned) mutable byte array to new specified size (in elem
|
||||
-- count). The returned array is either the original array resized in-place or,
|
||||
-- if not possible, a newly allocated (unpinned) array (with the original
|
||||
-- content copied over).
|
||||
{-# INLINE resizeArray #-}
|
||||
resizeArray ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Array (PrimState m) a
|
||||
-> Int -- ^ new size in elem count
|
||||
-> m (Array (PrimState m) a)
|
||||
resizeArray (Array arr#) (I# n#) =
|
||||
primitive $ \s# ->
|
||||
let bytes = n# *# sizeOf# (undefined :: a)
|
||||
in case resizeMutableByteArray# arr# bytes s# of
|
||||
(# s1#, arr1# #) -> (# s1#, Array arr1# #)
|
||||
110
src/Streamly/Internal/Data/Prim/Pinned/Array.hs
Normal file
110
src/Streamly/Internal/Data/Prim/Pinned/Array.hs
Normal file
@ -0,0 +1,110 @@
|
||||
{-# LANGUAGE UnboxedTuples #-}
|
||||
|
||||
#include "inline.hs"
|
||||
|
||||
-- |
|
||||
-- Module : Streamly.Internal.Data.Prim.Pinned.Array
|
||||
-- Copyright : (c) 2019 Composewell Technologies
|
||||
--
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
-- Stability : experimental
|
||||
-- Portability : GHC
|
||||
--
|
||||
module Streamly.Internal.Data.Prim.Pinned.Array
|
||||
(
|
||||
Array
|
||||
|
||||
-- * Construction
|
||||
|
||||
-- Pure List APIs
|
||||
, A.fromListN
|
||||
, A.fromList
|
||||
|
||||
-- Stream Folds
|
||||
, fromStreamN
|
||||
, fromStream
|
||||
|
||||
-- Monadic APIs
|
||||
-- , newArray
|
||||
, A.writeN -- drop new
|
||||
, A.write -- full buffer
|
||||
|
||||
-- * Elimination
|
||||
|
||||
, A.toList
|
||||
, toStream
|
||||
, toStreamRev
|
||||
, read
|
||||
, unsafeRead
|
||||
|
||||
-- * Random Access
|
||||
, length
|
||||
, null
|
||||
, last
|
||||
-- , (!!)
|
||||
, readIndex
|
||||
, A.unsafeIndex
|
||||
-- , readIndices
|
||||
-- , readRanges
|
||||
|
||||
-- , readFrom -- read from a given position to the end of file
|
||||
-- , readFromRev -- read from a given position to the beginning of file
|
||||
-- , readTo -- read from beginning up to the given position
|
||||
-- , readToRev -- read from end to the given position in file
|
||||
-- , readFromTo
|
||||
-- , readFromThenTo
|
||||
|
||||
-- , readChunksOfFrom
|
||||
-- , ...
|
||||
|
||||
-- , writeIndex
|
||||
-- , writeFrom -- start writing at the given position
|
||||
-- , writeFromRev
|
||||
-- , writeTo -- write from beginning up to the given position
|
||||
-- , writeToRev
|
||||
-- , writeFromTo
|
||||
-- , writeFromThenTo
|
||||
--
|
||||
-- , writeChunksOfFrom
|
||||
-- , ...
|
||||
|
||||
-- , writeIndex
|
||||
-- , writeIndices
|
||||
-- , writeRanges
|
||||
|
||||
-- -- * Search
|
||||
-- , bsearch
|
||||
-- , bsearchIndex
|
||||
-- , find
|
||||
-- , findIndex
|
||||
-- , findIndices
|
||||
|
||||
-- -- * In-pace mutation (for Mutable Array type)
|
||||
-- , partitionBy
|
||||
-- , shuffleBy
|
||||
-- , foldtWith
|
||||
-- , foldbWith
|
||||
|
||||
-- * Immutable Transformations
|
||||
-- , streamTransform
|
||||
|
||||
-- * Folding Arrays
|
||||
, streamFold
|
||||
, fold
|
||||
|
||||
-- * Folds with Array as the container
|
||||
-- , D.lastN
|
||||
|
||||
-- * Streaming array operations
|
||||
|
||||
, concat
|
||||
, compact
|
||||
|
||||
)
|
||||
where
|
||||
|
||||
import Streamly.Internal.Data.Prim.Pinned.Array.Types (Array(..), length)
|
||||
import qualified Streamly.Internal.Data.Prim.Pinned.Array.Types as A
|
||||
|
||||
#include "prim-array.hs"
|
||||
127
src/Streamly/Internal/Data/Prim/Pinned/Array/Types.hs
Normal file
127
src/Streamly/Internal/Data/Prim/Pinned/Array/Types.hs
Normal file
@ -0,0 +1,127 @@
|
||||
{-# LANGUAGE UnboxedTuples #-}
|
||||
|
||||
#include "inline.hs"
|
||||
|
||||
-- |
|
||||
-- Module : Streamly.Internal.Data.Prim.Pinned.Array.Types
|
||||
-- Copyright : (c) 2019 Composewell Technologies
|
||||
--
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
-- Stability : experimental
|
||||
-- Portability : GHC
|
||||
--
|
||||
module Streamly.Internal.Data.Prim.Pinned.Array.Types
|
||||
(
|
||||
Array (..)
|
||||
, unsafeFreeze
|
||||
, unsafeFreezeWithShrink
|
||||
-- , unsafeThaw
|
||||
, defaultChunkSize
|
||||
, nil
|
||||
|
||||
-- * Construction
|
||||
, spliceTwo
|
||||
|
||||
, fromList
|
||||
, fromListN
|
||||
, fromStreamDN
|
||||
, fromStreamD
|
||||
|
||||
-- * Streams of arrays
|
||||
, fromStreamDArraysOf
|
||||
, FlattenState (..) -- for inspection testing
|
||||
, flattenArrays
|
||||
, flattenArraysRev
|
||||
, SpliceState (..) -- for inspection testing
|
||||
, packArraysChunksOf
|
||||
, lpackArraysChunksOf
|
||||
#if !defined(mingw32_HOST_OS)
|
||||
-- , groupIOVecsOf
|
||||
#endif
|
||||
, splitOn
|
||||
, breakOn
|
||||
|
||||
-- * Elimination
|
||||
, unsafeIndex
|
||||
, byteLength
|
||||
, length
|
||||
|
||||
, foldl'
|
||||
, foldr
|
||||
, foldr'
|
||||
, foldlM'
|
||||
, splitAt
|
||||
|
||||
, toStreamD
|
||||
, toStreamDRev
|
||||
, toStreamK
|
||||
, toStreamKRev
|
||||
, toList
|
||||
-- , toArrayMinChunk
|
||||
, writeN
|
||||
, MA.ArrayUnsafe(..)
|
||||
, writeNUnsafe
|
||||
, write
|
||||
|
||||
, unlines
|
||||
|
||||
, toPtr
|
||||
|
||||
, touchArray
|
||||
, withArrayAsPtr
|
||||
)
|
||||
where
|
||||
|
||||
import Foreign.C.Types (CSize(..))
|
||||
import GHC.IO (IO(..))
|
||||
import Foreign.Ptr (minusPtr, nullPtr, plusPtr)
|
||||
|
||||
import qualified Streamly.Internal.Data.Prim.Pinned.Mutable.Array.Types as MA
|
||||
|
||||
#include "prim-array-types.hs"
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Utility functions
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
foreign import ccall unsafe "string.h memchr" c_memchr
|
||||
:: Ptr Word8 -> Word8 -> CSize -> IO (Ptr Word8)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Using as a Pointer
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- Change name later.
|
||||
{-# INLINE toPtr #-}
|
||||
toPtr :: Array a -> Ptr a
|
||||
toPtr (Array arr# off _) = Ptr (byteArrayContents# arr#) `plusPtr` off
|
||||
|
||||
{-# INLINE touchArray #-}
|
||||
touchArray :: Array a -> IO ()
|
||||
touchArray (Array arr# _ _) = IO $ \s -> case touch# arr# s of s1 -> (# s1, () #)
|
||||
|
||||
{-# INLINE withArrayAsPtr #-}
|
||||
withArrayAsPtr :: Array a -> (Ptr a -> IO b) -> IO b
|
||||
withArrayAsPtr arr f = do
|
||||
r <- f (toPtr arr)
|
||||
touchArray arr
|
||||
return r
|
||||
|
||||
-- Drops the separator byte
|
||||
{-# INLINE breakOn #-}
|
||||
breakOn ::
|
||||
PrimMonad m
|
||||
=> Word8
|
||||
-> Array Word8
|
||||
-> m (Array Word8, Maybe (Array Word8))
|
||||
breakOn sep arr@(Array arr# off len) = do
|
||||
let p = toPtr arr
|
||||
loc = unsafePerformIO $ c_memchr p sep (fromIntegral (byteLength arr))
|
||||
len1 = loc `minusPtr` p
|
||||
len2 = len - len1 - 1
|
||||
return $
|
||||
if loc == nullPtr
|
||||
then (arr, Nothing)
|
||||
else ( Array arr# off len1
|
||||
, Just $ Array arr# (off + len1 + 1) len2)
|
||||
177
src/Streamly/Internal/Data/Prim/Pinned/Mutable/Array/Types.hs
Normal file
177
src/Streamly/Internal/Data/Prim/Pinned/Mutable/Array/Types.hs
Normal file
@ -0,0 +1,177 @@
|
||||
{-# LANGUAGE UnboxedTuples #-}
|
||||
|
||||
#include "inline.hs"
|
||||
|
||||
-- |
|
||||
-- Module : Streamly.Internal.Data.Prim.Pinned.Mutable.Array.Types
|
||||
-- Copyright : (c) 2019 Composewell Technologies
|
||||
--
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
-- Stability : experimental
|
||||
-- Portability : GHC
|
||||
--
|
||||
module Streamly.Internal.Data.Prim.Pinned.Mutable.Array.Types
|
||||
(
|
||||
Array (..)
|
||||
|
||||
-- * Construction
|
||||
, newArray
|
||||
, newAlignedArray
|
||||
, unsafeWriteIndex
|
||||
|
||||
, spliceTwo
|
||||
, unsafeCopy
|
||||
|
||||
, fromListM
|
||||
, fromListNM
|
||||
, fromStreamDN
|
||||
, fromStreamD
|
||||
|
||||
-- * Streams of arrays
|
||||
, fromStreamDArraysOf
|
||||
|
||||
, packArraysChunksOf
|
||||
, lpackArraysChunksOf
|
||||
|
||||
#if !defined(mingw32_HOST_OS)
|
||||
-- , groupIOVecsOf
|
||||
#endif
|
||||
|
||||
-- * Elimination
|
||||
, unsafeReadIndex
|
||||
, length
|
||||
, byteLength
|
||||
|
||||
, writeN
|
||||
, ArrayUnsafe(..)
|
||||
, writeNUnsafe
|
||||
, writeNAligned
|
||||
, write
|
||||
|
||||
-- * Utilities
|
||||
, resizeArray
|
||||
, shrinkArray
|
||||
|
||||
, touchArray
|
||||
, withArrayAsPtr
|
||||
)
|
||||
where
|
||||
|
||||
import GHC.IO (IO(..))
|
||||
|
||||
#include "mutable-prim-array-types.hs"
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Allocation (Pinned)
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- XXX we can use a single newArray routine which accepts an allocation
|
||||
-- function which could be newByteArray#, newPinnedByteArray# or
|
||||
-- newAlignedPinnedByteArray#. That function can go in the common include file.
|
||||
--
|
||||
-- | Allocate an array that is pinned and can hold 'count' items. The memory of
|
||||
-- the array is uninitialized.
|
||||
--
|
||||
-- Note that this is internal routine, the reference to this array cannot be
|
||||
-- given out until the array has been written to and frozen.
|
||||
{-# INLINE newArray #-}
|
||||
newArray ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> m (Array (PrimState m) a)
|
||||
newArray (I# n#) =
|
||||
primitive $ \s# ->
|
||||
let bytes = n# *# sizeOf# (undefined :: a)
|
||||
in case newPinnedByteArray# bytes s# of
|
||||
(# s1#, arr# #) -> (# s1#, Array arr# #)
|
||||
|
||||
-- Change order of args?
|
||||
-- | Allocate a new array aligned to the specified alignment and using pinned
|
||||
-- memory.
|
||||
{-# INLINE newAlignedArray #-}
|
||||
newAlignedArray ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int -- size
|
||||
-> Int -- Alignment
|
||||
-> m (Array (PrimState m) a)
|
||||
newAlignedArray (I# n#) (I# a#) =
|
||||
primitive $ \s# ->
|
||||
let bytes = n# *# sizeOf# (undefined :: a)
|
||||
in case newAlignedPinnedByteArray# bytes a# s# of
|
||||
(# s1#, arr# #) -> (# s1#, Array arr# #)
|
||||
|
||||
-- | Resize (pinned) mutable byte array to new specified size (in elem
|
||||
-- count). The returned array is either the original array resized in-place or,
|
||||
-- if not possible, a newly allocated (pinned) array (with the original content
|
||||
-- copied over).
|
||||
{-# INLINE resizeArray #-}
|
||||
resizeArray ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Array (PrimState m) a
|
||||
-> Int -- ^ new size
|
||||
-> m (Array (PrimState m) a)
|
||||
resizeArray arr i =
|
||||
if len == i
|
||||
then return arr
|
||||
else if i < len
|
||||
then shrinkArray arr i >> return arr
|
||||
else do
|
||||
nArr <- newArray i
|
||||
unsafeCopy nArr 0 arr 0 len
|
||||
return nArr
|
||||
|
||||
where
|
||||
|
||||
len = length arr
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Aligned Construction
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- XXX we can also factor out common code in writeN and writeNAligned in the
|
||||
-- same way as suggested above.
|
||||
--
|
||||
{-# INLINE_NORMAL writeNAligned #-}
|
||||
writeNAligned ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> Int
|
||||
-> Fold m a (Array (PrimState m) a)
|
||||
writeNAligned align limit = Fold step initial extract
|
||||
|
||||
where
|
||||
|
||||
initial = do
|
||||
marr <- newAlignedArray limit align
|
||||
return (marr, 0)
|
||||
|
||||
step (marr, i) x
|
||||
| i == limit = return (marr, i)
|
||||
| otherwise = do
|
||||
unsafeWriteIndex marr i x
|
||||
return (marr, i + 1)
|
||||
|
||||
extract (marr, len) = shrinkArray marr len >> return marr
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Mutation with pointers
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- XXX This section can probably go in a common include file for pinned arrays.
|
||||
|
||||
-- Change name later.
|
||||
{-# INLINE toPtr #-}
|
||||
toPtr :: Array s a -> Ptr a
|
||||
toPtr (Array arr#) = Ptr (byteArrayContents# (unsafeCoerce# arr#))
|
||||
|
||||
{-# INLINE touchArray #-}
|
||||
touchArray :: Array s a -> IO ()
|
||||
touchArray arr = IO $ \s -> case touch# arr s of s' -> (# s', () #)
|
||||
|
||||
{-# INLINE withArrayAsPtr #-}
|
||||
withArrayAsPtr :: Array s a -> (Ptr a -> IO b) -> IO b
|
||||
withArrayAsPtr arr f = do
|
||||
r <- f (toPtr arr)
|
||||
touchArray arr
|
||||
return r
|
||||
84
src/Streamly/Internal/Data/Prim/Pinned/Unicode/Array.hs
Normal file
84
src/Streamly/Internal/Data/Prim/Pinned/Unicode/Array.hs
Normal file
@ -0,0 +1,84 @@
|
||||
-- |
|
||||
-- Module : Streamly.Internal.Data.Prim.Pinned.Unicode.Array
|
||||
-- Copyright : (c) 2018 Composewell Technologies
|
||||
--
|
||||
-- License : BSD3
|
||||
-- Maintainer : streamly@composewell.com
|
||||
-- Stability : experimental
|
||||
-- Portability : GHC
|
||||
--
|
||||
module Streamly.Internal.Data.Prim.Pinned.Unicode.Array
|
||||
(
|
||||
-- * Streams of Strings
|
||||
lines
|
||||
, words
|
||||
, unlines
|
||||
, unwords
|
||||
)
|
||||
where
|
||||
|
||||
import Control.Monad.Primitive (PrimMonad)
|
||||
import Streamly (IsStream, MonadAsync)
|
||||
import Prelude hiding (String, lines, words, unlines, unwords)
|
||||
import Streamly.Internal.Data.Prim.Pinned.Array (Array)
|
||||
|
||||
import qualified Streamly.Internal.Data.Unicode.Stream as S
|
||||
import qualified Streamly.Internal.Data.Prim.Pinned.Array as A
|
||||
|
||||
-- | Break a string up into a stream of strings at newline characters.
|
||||
-- The resulting strings do not contain newlines.
|
||||
--
|
||||
-- > lines = S.lines A.write
|
||||
--
|
||||
-- >>> S.toList $ lines $ S.fromList "lines\nthis\nstring\n\n\n"
|
||||
-- ["lines","this","string","",""]
|
||||
--
|
||||
{-# INLINE lines #-}
|
||||
lines :: (PrimMonad m, IsStream t) => t m Char -> t m (Array Char)
|
||||
lines = S.lines A.write
|
||||
|
||||
-- | Break a string up into a stream of strings, which were delimited
|
||||
-- by characters representing white space.
|
||||
--
|
||||
-- > words = S.words A.write
|
||||
--
|
||||
-- >>> S.toList $ words $ S.fromList "A newline\nis considered white space?"
|
||||
-- ["A", "newline", "is", "considered", "white", "space?"]
|
||||
--
|
||||
{-# INLINE words #-}
|
||||
words :: (PrimMonad m, IsStream t) => t m Char -> t m (Array Char)
|
||||
words = S.words A.write
|
||||
|
||||
-- | Flattens the stream of @Array Char@, after appending a terminating
|
||||
-- newline to each string.
|
||||
--
|
||||
-- 'unlines' is an inverse operation to 'lines'.
|
||||
--
|
||||
-- >>> S.toList $ unlines $ S.fromList ["lines", "this", "string"]
|
||||
-- "lines\nthis\nstring\n"
|
||||
--
|
||||
-- > unlines = S.unlines A.read
|
||||
--
|
||||
-- Note that, in general
|
||||
--
|
||||
-- > unlines . lines /= id
|
||||
{-# INLINE unlines #-}
|
||||
unlines :: (MonadAsync m, PrimMonad m, IsStream t) => t m (Array Char) -> t m Char
|
||||
unlines = S.unlines A.read
|
||||
|
||||
-- | Flattens the stream of @Array Char@, after appending a separating
|
||||
-- space to each string.
|
||||
--
|
||||
-- 'unwords' is an inverse operation to 'words'.
|
||||
--
|
||||
-- >>> S.toList $ unwords $ S.fromList ["unwords", "this", "string"]
|
||||
-- "unwords this string"
|
||||
--
|
||||
-- > unwords = S.unwords A.read
|
||||
--
|
||||
-- Note that, in general
|
||||
--
|
||||
-- > unwords . words /= id
|
||||
{-# INLINE unwords #-}
|
||||
unwords :: (MonadAsync m, PrimMonad m, IsStream t) => t m (Array Char) -> t m Char
|
||||
unwords = S.unwords A.read
|
||||
409
src/mutable-prim-array-types.hs
Normal file
409
src/mutable-prim-array-types.hs
Normal file
@ -0,0 +1,409 @@
|
||||
-- MOVE THIS TO A DIFFERENT LOCATION
|
||||
|
||||
import Control.Monad (when)
|
||||
import Control.Monad.Primitive
|
||||
import Data.Primitive.Types
|
||||
import Streamly.Internal.Data.Fold.Types (Fold(..))
|
||||
import Streamly.Internal.Data.SVar (adaptState)
|
||||
import Streamly.Internal.Data.Strict (Tuple'(..), Tuple3'(..))
|
||||
|
||||
import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
|
||||
|
||||
import GHC.Exts
|
||||
import Prelude hiding (length, unlines)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Array Data Type
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
data Array s a = Array (MutableByteArray# s)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Utilities
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Copy a range of the first array to the specified region in the second
|
||||
-- array. Both arrays must fully contain the specified ranges, but this is not
|
||||
-- checked. The regions are allowed to overlap, although this is only possible
|
||||
-- when the same array is provided as both the source and the destination.
|
||||
{-# INLINE unsafeCopy #-}
|
||||
unsafeCopy ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Array (PrimState m) a -- ^ destination array
|
||||
-> Int -- ^ offset into destination array
|
||||
-> Array (PrimState m) a -- ^ source array
|
||||
-> Int -- ^ offset into source array
|
||||
-> Int -- ^ number of elements to copy
|
||||
-> m ()
|
||||
unsafeCopy (Array dst#) (I# doff#) (Array src#) (I# soff#) (I# n#) =
|
||||
let toBytes cnt# = cnt# *# (sizeOf# (undefined :: a))
|
||||
in primitive_ $
|
||||
copyMutableByteArray#
|
||||
src# (toBytes soff#) dst# (toBytes doff#) (toBytes n#)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Length
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- XXX rename to byteCount?
|
||||
{-# INLINE byteLength #-}
|
||||
byteLength :: Array s a -> Int
|
||||
byteLength (Array arr#) = I# (sizeofMutableByteArray# arr#)
|
||||
|
||||
-- XXX Rename length to elemCount so that there is no confusion bout what it
|
||||
-- means.
|
||||
--
|
||||
-- XXX Since size of 'a' is statically known, we can replace `quot` with shift
|
||||
-- when it is power of 2. Though it may not matter unless length is used too
|
||||
-- often.
|
||||
--
|
||||
{-# INLINE length #-}
|
||||
length :: forall s a. Prim a => Array s a -> Int
|
||||
length arr = byteLength arr `quot` (sizeOf (undefined :: a))
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Random Access
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE unsafeReadIndex #-}
|
||||
unsafeReadIndex ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Array (PrimState m) a
|
||||
-> Int
|
||||
-> m a
|
||||
unsafeReadIndex (Array arr#) (I# i#) = primitive (readByteArray# arr# i#)
|
||||
|
||||
{-# INLINE unsafeWriteIndex #-}
|
||||
unsafeWriteIndex ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Array (PrimState m) a -- ^ array
|
||||
-> Int -- ^ index
|
||||
-> a -- ^ element
|
||||
-> m ()
|
||||
unsafeWriteIndex (Array arr#) (I# i#) x = primitive_ (writeByteArray# arr# i# x)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Construction
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- Note: We do not store the actual length of the array in the Array
|
||||
-- constructor. Therefore, for "length" API to work correctly we need to match
|
||||
-- the ByteArray length with the used length by shrinking it.
|
||||
--
|
||||
-- However, it may be expensive to always shrink the array. We may want to
|
||||
-- shrink only if significant space is being wasted. If we want to do that then
|
||||
-- we will have to store the used length separately. Or does GHC take care of
|
||||
-- that?
|
||||
-- Although the docs are not explicit about it, given how the signature is,
|
||||
-- the shrinking must me inplace. "resizeMutableByteArray#" shrinks the
|
||||
-- array inplace.
|
||||
{-# INLINE shrinkArray #-}
|
||||
shrinkArray ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Array (PrimState m) a
|
||||
-> Int -- ^ new size
|
||||
-> m ()
|
||||
shrinkArray (Array arr#) (I# n#) =
|
||||
let bytes = n# *# (sizeOf# (undefined :: a))
|
||||
in primitive_ (shrinkMutableByteArray# arr# bytes)
|
||||
|
||||
-- | Fold the whole input to a single array.
|
||||
--
|
||||
-- /Caution! Do not use this on infinite streams./
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE_NORMAL write #-}
|
||||
write ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Fold m a (Array (PrimState m) a)
|
||||
write = Fold step initial extract
|
||||
|
||||
where
|
||||
|
||||
initial = do
|
||||
marr <- newArray 0
|
||||
return $ Tuple3' marr 0 0
|
||||
|
||||
step (Tuple3' marr i capacity) x
|
||||
| i == capacity = do
|
||||
let newCapacity = max (capacity * 2) 1
|
||||
newMarr <- resizeArray marr newCapacity
|
||||
unsafeWriteIndex newMarr i x
|
||||
return $ Tuple3' newMarr (i + 1) newCapacity
|
||||
| otherwise = do
|
||||
unsafeWriteIndex marr i x
|
||||
return $ Tuple3' marr (i + 1) capacity
|
||||
|
||||
extract (Tuple3' marr len _) = shrinkArray marr len >> return marr
|
||||
|
||||
-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
|
||||
-- 'Array'.
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE_NORMAL writeN #-}
|
||||
writeN ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> Fold m a (Array (PrimState m) a)
|
||||
writeN limit = Fold step initial extract
|
||||
|
||||
where
|
||||
|
||||
initial = do
|
||||
marr <- newArray limit
|
||||
return $ Tuple' marr 0
|
||||
|
||||
step (Tuple' marr i) x
|
||||
| i == limit = return $ Tuple' marr i
|
||||
| otherwise = do
|
||||
unsafeWriteIndex marr i x
|
||||
return $ Tuple' marr (i + 1)
|
||||
|
||||
extract (Tuple' marr len) = shrinkArray marr len >> return marr
|
||||
|
||||
-- Use Tuple' instead?
|
||||
data ArrayUnsafe s a = ArrayUnsafe
|
||||
{-# UNPACK #-} !(Array s a)
|
||||
{-# UNPACK #-} !Int
|
||||
|
||||
-- | Like 'writeN' but does not check the array bounds when writing. The fold
|
||||
-- driver must not call the step function more than 'n' times otherwise it will
|
||||
-- corrupt the memory and crash. This function exists mainly because any
|
||||
-- conditional in the step function blocks fusion causing 10x performance
|
||||
-- slowdown.
|
||||
--
|
||||
-- @since 0.7.0
|
||||
{-# INLINE_NORMAL writeNUnsafe #-}
|
||||
writeNUnsafe ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> Fold m a (Array (PrimState m) a)
|
||||
writeNUnsafe n = Fold step initial extract
|
||||
|
||||
where
|
||||
|
||||
initial = do
|
||||
arr <- newArray (max n 0)
|
||||
return $ ArrayUnsafe arr 0
|
||||
step (ArrayUnsafe marr i) x = do
|
||||
unsafeWriteIndex marr i x
|
||||
return $ ArrayUnsafe marr (i + 1)
|
||||
extract (ArrayUnsafe marr i) = shrinkArray marr i >> return marr
|
||||
|
||||
{-# INLINE_NORMAL fromStreamDN #-}
|
||||
fromStreamDN ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> D.Stream m a
|
||||
-> m (Array (PrimState m) a)
|
||||
fromStreamDN limit str = do
|
||||
marr <- newArray (max limit 0)
|
||||
let step i x = i `seq` (unsafeWriteIndex marr i x) >> return (i + 1)
|
||||
n <- D.foldlM' step (return 0) $ D.take limit str
|
||||
shrinkArray marr n
|
||||
return marr
|
||||
|
||||
{-# INLINE runFold #-}
|
||||
runFold :: (Monad m) => Fold m a b -> D.Stream m a -> m b
|
||||
runFold (Fold step begin done) = D.foldlMx' step begin done
|
||||
|
||||
{-# INLINE fromStreamD #-}
|
||||
fromStreamD ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> D.Stream m a
|
||||
-> m (Array (PrimState m) a)
|
||||
fromStreamD str = runFold write str
|
||||
|
||||
{-# INLINABLE fromListNM #-}
|
||||
fromListNM ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> [a]
|
||||
-> m (Array (PrimState m) a)
|
||||
fromListNM n xs = fromStreamDN n $ D.fromList xs
|
||||
|
||||
{-# INLINABLE fromListM #-}
|
||||
fromListM ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> [a]
|
||||
-> m (Array (PrimState m) a)
|
||||
fromListM xs = fromStreamD $ D.fromList xs
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Combining
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- Splice two mutable arrays creating a new array.
|
||||
{-# INLINE spliceTwo #-}
|
||||
spliceTwo ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Array (PrimState m) a
|
||||
-> Array (PrimState m) a
|
||||
-> m (Array (PrimState m) a)
|
||||
spliceTwo a1 a2 = do
|
||||
a3 <- resizeArray a1 (l1 + l2)
|
||||
unsafeCopy a2 0 a3 l1 l2
|
||||
return a3
|
||||
|
||||
where
|
||||
|
||||
l1 = length a1
|
||||
l2 = length a2
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Stream of Arrays
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
data GroupState s t a
|
||||
= GroupStart s
|
||||
| GroupBuffer s (Array t a) Int
|
||||
| GroupYield (Array t a) s
|
||||
| GroupLastYield (Array t a) Int
|
||||
| GroupFinish
|
||||
|
||||
-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
|
||||
-- arrays of size n.
|
||||
{-# INLINE_NORMAL fromStreamDArraysOf #-}
|
||||
fromStreamDArraysOf ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> D.Stream m a
|
||||
-> D.Stream m (Array (PrimState m) a)
|
||||
-- fromStreamDArraysOf n str = D.groupsOf n (writeN n) str
|
||||
fromStreamDArraysOf n (D.Stream step state) = D.Stream step' (GroupStart state)
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step' #-}
|
||||
step' _ (GroupStart st) = do
|
||||
when (n <= 0) $
|
||||
-- XXX we can pass the module string from the higher level API
|
||||
error $
|
||||
"Streamly.Internal.Memory.Mutable.Array.Types.fromStreamDArraysOf: the size of " ++
|
||||
"arrays [" ++ show n ++ "] must be a natural number"
|
||||
arr <- newArray n
|
||||
return $ D.Skip (GroupBuffer st arr 0)
|
||||
step' gst (GroupBuffer st arr i)
|
||||
| i < n = do
|
||||
r <- step (adaptState gst) st
|
||||
case r of
|
||||
D.Yield x s -> do
|
||||
unsafeWriteIndex arr i x
|
||||
return $ D.Skip (GroupBuffer s arr (i + 1))
|
||||
D.Skip s -> return $ D.Skip (GroupBuffer s arr i)
|
||||
D.Stop -> return $ D.Skip (GroupLastYield arr i)
|
||||
| otherwise = return $ D.Skip (GroupYield arr st)
|
||||
step' _ (GroupYield arr st) = do
|
||||
nArr <- newArray n
|
||||
return $ D.Yield arr (GroupBuffer st nArr 0)
|
||||
step' _ (GroupLastYield arr i)
|
||||
| i == 0 = return D.Stop
|
||||
| otherwise = do
|
||||
shrinkArray arr i
|
||||
return $ D.Yield arr GroupFinish
|
||||
step' _ GroupFinish = return D.Stop
|
||||
|
||||
data SpliceState s arr
|
||||
= SpliceInitial s
|
||||
| SpliceBuffering s arr
|
||||
| SpliceYielding arr (SpliceState s arr)
|
||||
| SpliceFinish
|
||||
|
||||
-- XXX can use general grouping combinators to achieve this?
|
||||
-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
|
||||
-- maximum specified size in bytes. Note that if a single array is bigger than
|
||||
-- the specified size we do not split it to fit. When we coalesce multiple
|
||||
-- arrays if the size would exceed the specified size we do not coalesce
|
||||
-- therefore the actual array size may be less than the specified chunk size.
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE_NORMAL packArraysChunksOf #-}
|
||||
packArraysChunksOf ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> D.Stream m (Array (PrimState m) a)
|
||||
-> D.Stream m (Array (PrimState m) a)
|
||||
packArraysChunksOf n (D.Stream step state) =
|
||||
D.Stream step' (SpliceInitial state)
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step' #-}
|
||||
step' gst (SpliceInitial st) = do
|
||||
when (n <= 0) $
|
||||
-- XXX we can pass the module string from the higher level API
|
||||
error $ "Streamly.Internal.Memory.Mutable.Array.Types.packArraysChunksOf: the size of "
|
||||
++ "arrays [" ++ show n ++ "] must be a natural number"
|
||||
r <- step gst st
|
||||
case r of
|
||||
D.Yield arr s -> return $
|
||||
let len = byteLength arr
|
||||
in if len >= n
|
||||
then D.Skip (SpliceYielding arr (SpliceInitial s))
|
||||
else D.Skip (SpliceBuffering s arr)
|
||||
D.Skip s -> return $ D.Skip (SpliceInitial s)
|
||||
D.Stop -> return $ D.Stop
|
||||
|
||||
step' gst (SpliceBuffering st buf) = do
|
||||
r <- step gst st
|
||||
case r of
|
||||
D.Yield arr s -> do
|
||||
let len = byteLength buf + byteLength arr
|
||||
if len > n
|
||||
then return $
|
||||
D.Skip (SpliceYielding buf (SpliceBuffering s arr))
|
||||
else do
|
||||
buf' <- spliceTwo buf arr
|
||||
return $ D.Skip (SpliceBuffering s buf')
|
||||
D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
|
||||
D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
|
||||
|
||||
step' _ SpliceFinish = return D.Stop
|
||||
|
||||
step' _ (SpliceYielding arr next) = return $ D.Yield arr next
|
||||
|
||||
{-# INLINE_NORMAL lpackArraysChunksOf #-}
|
||||
lpackArraysChunksOf ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> Fold m (Array (PrimState m) a) ()
|
||||
-> Fold m (Array (PrimState m) a) ()
|
||||
lpackArraysChunksOf n (Fold step1 initial1 extract1) =
|
||||
Fold step initial extract
|
||||
|
||||
where
|
||||
|
||||
initial = do
|
||||
when (n <= 0) $
|
||||
-- XXX we can pass the module string from the higher level API
|
||||
error $ "Streamly.Internal.Memory.Mutable.Array.Types.packArraysChunksOf: the size of "
|
||||
++ "arrays [" ++ show n ++ "] must be a natural number"
|
||||
r1 <- initial1
|
||||
return (Tuple' Nothing r1)
|
||||
|
||||
extract (Tuple' Nothing r1) = extract1 r1
|
||||
extract (Tuple' (Just buf) r1) = do
|
||||
r <- step1 r1 buf
|
||||
extract1 r
|
||||
|
||||
step (Tuple' Nothing r1) arr = do
|
||||
let len = byteLength arr
|
||||
in if len >= n
|
||||
then do
|
||||
r <- step1 r1 arr
|
||||
extract1 r
|
||||
r1' <- initial1
|
||||
return (Tuple' Nothing r1')
|
||||
else return (Tuple' (Just arr) r1)
|
||||
|
||||
step (Tuple' (Just buf) r1) arr = do
|
||||
let len = byteLength buf + byteLength arr
|
||||
buf' <- spliceTwo buf arr
|
||||
|
||||
if len >= n
|
||||
then do
|
||||
r <- step1 r1 buf'
|
||||
extract1 r
|
||||
r1' <- initial1
|
||||
return (Tuple' Nothing r1')
|
||||
else return (Tuple' (Just buf') r1)
|
||||
745
src/prim-array-types.hs
Normal file
745
src/prim-array-types.hs
Normal file
@ -0,0 +1,745 @@
|
||||
import Control.DeepSeq (NFData(..))
|
||||
import Control.Monad (when)
|
||||
import Control.Monad.Primitive
|
||||
(PrimMonad(primitive), PrimState, primitive_)
|
||||
import Control.Monad.ST (ST, runST)
|
||||
#if __GLASGOW_HASKELL__ < 808
|
||||
import Data.Semigroup (Semigroup(..))
|
||||
#endif
|
||||
import Data.Word (Word8)
|
||||
import Streamly.Internal.Data.Strict (Tuple3'(..), Maybe'(..))
|
||||
import Streamly.Internal.Data.Fold.Types (Fold(..))
|
||||
import Streamly.Internal.Data.SVar (adaptState)
|
||||
import Text.Read (readPrec, readListPrec, readListPrecDefault)
|
||||
|
||||
import System.IO.Unsafe (unsafePerformIO)
|
||||
|
||||
import qualified GHC.Exts as Exts
|
||||
import qualified Prelude as P
|
||||
import qualified Streamly.Internal.Data.Fold as FL
|
||||
import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
|
||||
import qualified Streamly.Internal.Data.Stream.StreamK as K
|
||||
|
||||
import Data.Primitive.Types
|
||||
import GHC.Exts hiding (fromListN, fromList, toList)
|
||||
import Prelude hiding (length, unlines, foldr)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Array Data Type
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
data Array a = Array ByteArray# Int Int
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Utilities
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Both arrays must fully contain the specified ranges, but this is not
|
||||
-- checked. The two arrays must not be the same array in different states, but
|
||||
-- this is not checked either.
|
||||
{-# INLINE unsafeCopy #-}
|
||||
unsafeCopy ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> MA.Array (PrimState m) a -- ^ destination array
|
||||
-> Int -- ^ offset into destination array
|
||||
-> Array a -- ^ source array
|
||||
-> Int -- ^ offset into source array
|
||||
-> Int -- ^ number of elements to copy
|
||||
-> m ()
|
||||
unsafeCopy (MA.Array dst#) (I# doff#) (Array src# (I# off#) _) (I# soff#) (I# n#) =
|
||||
let toBytes cnt# = cnt# *# (sizeOf# (undefined :: a))
|
||||
in primitive_ $
|
||||
copyByteArray# src# (toBytes (off# +# soff#)) dst# (toBytes doff#) (toBytes n#)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Basic Byte Array Operations
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE unsafeFreeze #-}
|
||||
unsafeFreeze ::
|
||||
forall a m. (Prim a, PrimMonad m)
|
||||
=> MA.Array (PrimState m) a
|
||||
-> m (Array a)
|
||||
unsafeFreeze (MA.Array arr#) =
|
||||
primitive $ \s# ->
|
||||
case unsafeFreezeByteArray# arr# s# of
|
||||
(# s1#, arr1# #) ->
|
||||
(# s1#
|
||||
, Array
|
||||
arr1#
|
||||
0
|
||||
(I# (quotInt#
|
||||
(sizeofByteArray# arr1#)
|
||||
(sizeOf# (undefined :: a))))#)
|
||||
|
||||
{-# INLINE unsafeFreezeWithShrink #-}
|
||||
unsafeFreezeWithShrink ::
|
||||
forall a m. (Prim a, PrimMonad m)
|
||||
=> MA.Array (PrimState m) a
|
||||
-> Int
|
||||
-> m (Array a)
|
||||
unsafeFreezeWithShrink arr@(MA.Array arr#) n = do
|
||||
MA.shrinkArray arr n
|
||||
primitive $ \s# ->
|
||||
case unsafeFreezeByteArray# arr# s# of
|
||||
(# s1#, arr1# #) ->
|
||||
(# s1#
|
||||
, Array
|
||||
arr1#
|
||||
0
|
||||
(I# (quotInt#
|
||||
(sizeofByteArray# arr1#)
|
||||
(sizeOf# (undefined :: a))))#)
|
||||
|
||||
{-
|
||||
-- Should never be used in general
|
||||
{-# INLINE unsafeThaw #-}
|
||||
unsafeThaw :: PrimMonad m => Array a -> m (MA.Array (PrimState m) a)
|
||||
unsafeThaw (Array arr#) =
|
||||
primitive $ \s# -> (# s#, MA.Array (unsafeCoerce# arr#) #)
|
||||
-}
|
||||
|
||||
-- Unsafe because the index bounds are not checked
|
||||
{-# INLINE unsafeIndex #-}
|
||||
unsafeIndex :: Prim a => Array a -> Int -> a
|
||||
unsafeIndex (Array arr# (I# off#) _) (I# i#) = indexByteArray# arr# (off# +# i#)
|
||||
|
||||
-- unsafe
|
||||
sameByteArray :: ByteArray# -> ByteArray# -> Bool
|
||||
sameByteArray ba1 ba2 =
|
||||
case reallyUnsafePtrEquality#
|
||||
(unsafeCoerce# ba1 :: ()) (unsafeCoerce# ba2 :: ()) of
|
||||
r -> isTrue# r
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Chunk Size
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- XXX move this section to mutable array module?
|
||||
|
||||
mkChunkSizeKB :: Int -> Int
|
||||
mkChunkSizeKB n = n * k
|
||||
where k = 1024
|
||||
|
||||
-- | Default maximum buffer size in bytes, for reading from and writing to IO
|
||||
-- devices, the value is 32KB minus GHC allocation overhead, which is a few
|
||||
-- bytes, so that the actual allocation is 32KB.
|
||||
defaultChunkSize :: Int
|
||||
defaultChunkSize = mkChunkSizeKB 32
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Length
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- XXX rename to byteCount?
|
||||
{-# INLINE byteLength #-}
|
||||
byteLength :: forall a. Prim a => Array a -> Int
|
||||
byteLength (Array _ _ len) = len * sizeOf (undefined :: a)
|
||||
|
||||
-- XXX Also, rename to elemCount
|
||||
-- XXX I would prefer length to keep the API consistent
|
||||
-- XXX Also, re-export sizeOf from Primitive
|
||||
{-# INLINE length #-}
|
||||
length :: Array a -> Int
|
||||
length (Array _ _ len) = len
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Construction
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Use a slice of an array as another array. Note that this is unsafe and does
|
||||
-- not check the bounds
|
||||
slice :: Array a -> Int -> Int -> Array a
|
||||
slice (Array arr# off _) off1 len1 = Array arr# (off + off1) len1
|
||||
|
||||
nil :: forall a. Prim a => Array a
|
||||
nil = runST run
|
||||
where
|
||||
run :: forall s. ST s (Array a)
|
||||
run = do
|
||||
arr <- MA.newArray 0
|
||||
unsafeFreeze arr
|
||||
|
||||
-- | Fold the whole input to a single array.
|
||||
--
|
||||
-- /Caution! Do not use this on infinite streams./
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE_NORMAL write #-}
|
||||
write :: (PrimMonad m, Prim a) => Fold m a (Array a)
|
||||
write = FL.mapM unsafeFreeze MA.write
|
||||
|
||||
-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
|
||||
-- 'Array'.
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE_NORMAL writeN #-}
|
||||
writeN :: (PrimMonad m, Prim a) => Int -> Fold m a (Array a)
|
||||
writeN limit = FL.mapM unsafeFreeze (MA.writeN limit)
|
||||
|
||||
-- | Like 'writeN' but does not check the array bounds when writing. The fold
|
||||
-- driver must not call the step function more than 'n' times otherwise it will
|
||||
-- corrupt the memory and crash. This function exists mainly because any
|
||||
-- conditional in the step function blocks fusion causing 10x performance
|
||||
-- slowdown.
|
||||
--
|
||||
-- @since 0.7.0
|
||||
{-# INLINE_NORMAL writeNUnsafe #-}
|
||||
writeNUnsafe ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> Fold m a (Array a)
|
||||
writeNUnsafe limit = FL.mapM unsafeFreeze (MA.writeNUnsafe limit)
|
||||
|
||||
{-# INLINE_NORMAL fromStreamDN #-}
|
||||
fromStreamDN ::
|
||||
(PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> D.Stream m a
|
||||
-> m (Array a)
|
||||
fromStreamDN limit str = MA.fromStreamDN limit str >>= unsafeFreeze
|
||||
|
||||
{-# INLINE fromStreamD #-}
|
||||
fromStreamD ::
|
||||
(PrimMonad m, Prim a) => D.Stream m a -> m (Array a)
|
||||
fromStreamD str = MA.fromStreamD str >>= unsafeFreeze
|
||||
|
||||
-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
|
||||
-- arrays of size n.
|
||||
{-# INLINE_NORMAL fromStreamDArraysOf #-}
|
||||
fromStreamDArraysOf ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> D.Stream m a
|
||||
-> D.Stream m (Array a)
|
||||
fromStreamDArraysOf n str = D.mapM unsafeFreeze (MA.fromStreamDArraysOf n str)
|
||||
|
||||
-- XXX derive from MA.fromListN?
|
||||
{-# INLINE fromListN #-}
|
||||
fromListN :: forall a. Prim a => Int -> [a] -> Array a
|
||||
fromListN len xs = unsafePerformIO $ MA.fromListNM len xs >>= unsafeFreeze
|
||||
|
||||
-- XXX derive from MA.fromList?
|
||||
{-# INLINE fromList #-}
|
||||
fromList :: Prim a => [a] -> Array a
|
||||
fromList xs = fromListN (P.length xs) xs
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Combining
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Splice two immutable arrays creating a new immutable array.
|
||||
{-# INLINE spliceTwo #-}
|
||||
spliceTwo :: (PrimMonad m, Prim a) => Array a -> Array a -> m (Array a)
|
||||
spliceTwo a1 a2 = do
|
||||
let l1 = length a1
|
||||
l2 = length a2
|
||||
a3 <- MA.newArray (l1 + l2)
|
||||
unsafeCopy a3 0 a1 0 l1
|
||||
unsafeCopy a3 l1 a2 0 l2
|
||||
unsafeFreeze a3 -- Use `unsafeFreezeWith off len`?
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Elimination
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE_LATE toListFB #-}
|
||||
toListFB :: forall a b. Prim a => (a -> b -> b) -> b -> Array a -> b
|
||||
toListFB c n arr = go 0
|
||||
where
|
||||
len = length arr
|
||||
go p | p == len = n
|
||||
go p =
|
||||
let !x = unsafeIndex arr p
|
||||
in c x (go (p + 1))
|
||||
|
||||
-- | Convert an 'Array' into a list.
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE toList #-}
|
||||
toList :: Prim a => Array a -> [a]
|
||||
toList s = build (\c n -> toListFB c n s)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Instances
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
instance (Eq a, Prim a) => Eq (Array a) where
|
||||
{-# INLINE (==) #-}
|
||||
a1@(Array ba1# _ len1) == a2@(Array ba2# _ len2)
|
||||
| sameByteArray ba1# ba2# = True
|
||||
| len1 /= len2 = False
|
||||
| otherwise = loop (len1 - 1)
|
||||
|
||||
where
|
||||
|
||||
loop !i
|
||||
| i < 0 = True
|
||||
| otherwise = unsafeIndex a1 i == unsafeIndex a2 i && loop (i - 1)
|
||||
|
||||
-- | Lexicographic ordering. Subject to change between major versions.
|
||||
instance (Ord a, Prim a) => Ord (Array a) where
|
||||
{-# INLINE compare #-}
|
||||
compare a1@(Array ba1# _ len1) a2@(Array ba2# _ len2)
|
||||
| sameByteArray ba1# ba2# = EQ
|
||||
| otherwise = loop 0
|
||||
|
||||
where
|
||||
|
||||
sz = min len1 len2
|
||||
|
||||
loop !i
|
||||
| i < sz =
|
||||
compare (unsafeIndex a1 i) (unsafeIndex a2 i) <> loop (i + 1)
|
||||
| otherwise = compare len1 len2
|
||||
|
||||
instance Prim a => Semigroup (Array a) where
|
||||
-- XXX can't we use runST instead of inlineIO?
|
||||
-- XXX I plan to remove PrimMonad and replace it with IO
|
||||
a <> b = unsafePerformIO (spliceTwo a b :: IO (Array a))
|
||||
|
||||
instance Prim a => Monoid (Array a) where
|
||||
mempty = nil
|
||||
mappend = (<>)
|
||||
|
||||
instance NFData (Array a) where
|
||||
{-# INLINE rnf #-}
|
||||
rnf _ = ()
|
||||
|
||||
|
||||
-- XXX check if this is compatible with Memory.Array?
|
||||
-- XXX It isn't. I might prefer this Show instance though
|
||||
-- XXX Memory.Array: showsPrec _ = shows . toList
|
||||
instance (Show a, Prim a) => Show (Array a) where
|
||||
showsPrec p a =
|
||||
showParen (p > 10) $
|
||||
showString "fromListN "
|
||||
. shows (length a)
|
||||
. showString " "
|
||||
. shows (toList a)
|
||||
|
||||
instance (a ~ Char) => IsString (Array a) where
|
||||
{-# INLINE fromString #-}
|
||||
fromString = fromList
|
||||
|
||||
-- GHC versions 8.0 and below cannot derive IsList
|
||||
instance Prim a => IsList (Array a) where
|
||||
type (Item (Array a)) = a
|
||||
|
||||
{-# INLINE fromList #-}
|
||||
fromList = fromList
|
||||
|
||||
{-# INLINE fromListN #-}
|
||||
fromListN = fromListN
|
||||
|
||||
{-# INLINE toList #-}
|
||||
toList = toList
|
||||
|
||||
instance (Prim a, Read a, Show a) => Read (Array a) where
|
||||
{-# INLINE readPrec #-}
|
||||
readPrec = fromList <$> readPrec
|
||||
readListPrec = readListPrecDefault
|
||||
|
||||
-- XXX these folds can be made common with mutable arrays by defining a
|
||||
-- unsafeIndex in the specific module?
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Folds
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE foldr #-}
|
||||
foldr ::
|
||||
forall a b. Prim a
|
||||
=> (a -> b -> b)
|
||||
-> b
|
||||
-> Array a
|
||||
-> b
|
||||
foldr f z arr = go 0
|
||||
|
||||
where
|
||||
|
||||
!len = length arr
|
||||
|
||||
go !i
|
||||
| len > i = f (unsafeIndex arr i) (go (i + 1))
|
||||
| otherwise = z
|
||||
|
||||
-- | Strict right-associated fold over the elements of an 'Array'.
|
||||
{-# INLINE foldr' #-}
|
||||
foldr' ::
|
||||
forall a b. Prim a
|
||||
=> (a -> b -> b)
|
||||
-> b
|
||||
-> Array a
|
||||
-> b
|
||||
foldr' f z0 arr = go (length arr - 1) z0
|
||||
|
||||
where
|
||||
|
||||
go !i !acc
|
||||
| i < 0 = acc
|
||||
| otherwise = go (i - 1) (f (unsafeIndex arr i) acc)
|
||||
|
||||
-- | Strict left-associated fold over the elements of an 'Array'.
|
||||
{-# INLINE foldl' #-}
|
||||
foldl' ::
|
||||
forall a b. Prim a
|
||||
=> (b -> a -> b)
|
||||
-> b
|
||||
-> Array a
|
||||
-> b
|
||||
foldl' f z0 arr = go 0 z0
|
||||
|
||||
where
|
||||
|
||||
!len = length arr
|
||||
|
||||
go !i !acc
|
||||
| i < len = go (i + 1) (f acc (unsafeIndex arr i))
|
||||
| otherwise = acc
|
||||
|
||||
-- | Strict left-associated fold over the elements of an 'Array'.
|
||||
{-# INLINE foldlM' #-}
|
||||
foldlM' :: (Prim a, Monad m) => (b -> a -> m b) -> b -> Array a -> m b
|
||||
foldlM' f z0 arr = go 0 z0
|
||||
|
||||
where
|
||||
|
||||
!len = length arr
|
||||
|
||||
go !i !acc1
|
||||
| i < len = do
|
||||
acc2 <- f acc1 (unsafeIndex arr i)
|
||||
go (i + 1) acc2
|
||||
| otherwise = return acc1
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Converting to streams
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
{-# INLINE_NORMAL toStreamD #-}
|
||||
toStreamD :: (Prim a, Monad m) => Array a -> D.Stream m a
|
||||
toStreamD arr = D.Stream step 0
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step #-}
|
||||
step _ i
|
||||
| i == length arr = return D.Stop
|
||||
step _ i = return $ D.Yield (unsafeIndex arr i) (i + 1)
|
||||
|
||||
{-# INLINE toStreamK #-}
|
||||
toStreamK ::
|
||||
forall t m a. (K.IsStream t, Prim a)
|
||||
=> Array a
|
||||
-> t m a
|
||||
toStreamK arr = go 0
|
||||
|
||||
where
|
||||
|
||||
len = length arr
|
||||
|
||||
go p
|
||||
| p == len = K.nil
|
||||
| otherwise =
|
||||
let !x = unsafeIndex arr p
|
||||
in x `K.cons` go (p + 1)
|
||||
|
||||
{-# INLINE_NORMAL toStreamDRev #-}
|
||||
toStreamDRev :: (Prim a, Monad m) => Array a -> D.Stream m a
|
||||
toStreamDRev arr = D.Stream step (length arr - 1)
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step #-}
|
||||
step _ i
|
||||
| i < 0 = return D.Stop
|
||||
step _ i = return $ D.Yield (unsafeIndex arr i) (i - 1)
|
||||
|
||||
{-# INLINE toStreamKRev #-}
|
||||
toStreamKRev ::
|
||||
forall t m a. (K.IsStream t, Prim a)
|
||||
=> Array a
|
||||
-> t m a
|
||||
toStreamKRev arr = go (length arr - 1)
|
||||
|
||||
where
|
||||
|
||||
go p | p == -1 = K.nil
|
||||
| otherwise =
|
||||
let !x = unsafeIndex arr p
|
||||
in x `K.cons` go (p - 1)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Stream of Arrays
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
data FlattenState s a =
|
||||
OuterLoop s
|
||||
| InnerLoop s !(Array a) !Int !Int
|
||||
|
||||
{-# INLINE_NORMAL flattenArrays #-}
|
||||
flattenArrays :: (PrimMonad m, Prim a) => D.Stream m (Array a) -> D.Stream m a
|
||||
flattenArrays (D.Stream step state) = D.Stream step' (OuterLoop state)
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step' #-}
|
||||
step' gst (OuterLoop st) = do
|
||||
r <- step (adaptState gst) st
|
||||
return $ case r of
|
||||
D.Yield arr s ->
|
||||
let len = length arr
|
||||
in if len == 0
|
||||
then D.Skip (OuterLoop s)
|
||||
else D.Skip (InnerLoop s arr len 0)
|
||||
D.Skip s -> D.Skip (OuterLoop s)
|
||||
D.Stop -> D.Stop
|
||||
|
||||
step' _ (InnerLoop st _ len i) | i == len =
|
||||
return $ D.Skip $ OuterLoop st
|
||||
|
||||
step' _ (InnerLoop st arr len i) = do
|
||||
let x = unsafeIndex arr i
|
||||
return $ D.Yield x (InnerLoop st arr len (i + 1))
|
||||
|
||||
{-# INLINE_NORMAL flattenArraysRev #-}
|
||||
flattenArraysRev ::
|
||||
(PrimMonad m, Prim a)
|
||||
=> D.Stream m (Array a)
|
||||
-> D.Stream m a
|
||||
flattenArraysRev (D.Stream step state) = D.Stream step' (OuterLoop state)
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step' #-}
|
||||
step' gst (OuterLoop st) = do
|
||||
r <- step (adaptState gst) st
|
||||
return $ case r of
|
||||
D.Yield arr s ->
|
||||
let len = length arr
|
||||
in if len == 0
|
||||
then D.Skip (OuterLoop s)
|
||||
else D.Skip (InnerLoop s arr len (len - 1))
|
||||
D.Skip s -> D.Skip (OuterLoop s)
|
||||
D.Stop -> D.Stop
|
||||
|
||||
step' _ (InnerLoop st _ _ i) | i == -1 =
|
||||
return $ D.Skip $ OuterLoop st
|
||||
|
||||
step' _ (InnerLoop st arr len i) = do
|
||||
let x = unsafeIndex arr i
|
||||
return $ D.Yield x (InnerLoop st arr len (i - 1))
|
||||
|
||||
{-# INLINE_NORMAL unlines #-}
|
||||
unlines ::
|
||||
(PrimMonad m, Prim a)
|
||||
=> a
|
||||
-> D.Stream m (Array a)
|
||||
-> D.Stream m a
|
||||
unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step' #-}
|
||||
step' gst (OuterLoop st) = do
|
||||
r <- step (adaptState gst) st
|
||||
return $ case r of
|
||||
D.Yield arr s ->
|
||||
let len = length arr
|
||||
in D.Skip (InnerLoop s arr len 0)
|
||||
D.Skip s -> D.Skip (OuterLoop s)
|
||||
D.Stop -> D.Stop
|
||||
|
||||
step' _ (InnerLoop st _ len i) | i == len =
|
||||
return $ D.Yield sep $ OuterLoop st
|
||||
|
||||
step' _ (InnerLoop st arr len i) = do
|
||||
let x = unsafeIndex arr i
|
||||
return $ D.Yield x (InnerLoop st arr len (i + 1))
|
||||
|
||||
-- Splice an array into a pre-reserved mutable array. The user must ensure
|
||||
-- that there is enough space in the mutable array.
|
||||
{-# INLINE spliceInto #-}
|
||||
spliceInto ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> MA.Array (PrimState m) a
|
||||
-> Int
|
||||
-> Array a
|
||||
-> m Int
|
||||
spliceInto dst doff src@(Array _ _ len) = do
|
||||
unsafeCopy dst doff src 0 len
|
||||
return $ doff + len
|
||||
|
||||
data SpliceState s arr1 arr2
|
||||
= SpliceInitial s
|
||||
| SpliceBuffering s arr2
|
||||
| SpliceYielding arr1 (SpliceState s arr1 arr2)
|
||||
| SpliceFinish
|
||||
|
||||
-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
|
||||
-- maximum specified size in bytes. Note that if a single array is bigger than
|
||||
-- the specified size we do not split it to fit. When we coalesce multiple
|
||||
-- arrays if the size would exceed the specified size we do not coalesce
|
||||
-- therefore the actual array size may be less than the specified chunk size.
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE_NORMAL packArraysChunksOf #-}
|
||||
packArraysChunksOf ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> D.Stream m (Array a)
|
||||
-> D.Stream m (Array a)
|
||||
packArraysChunksOf n (D.Stream step state) =
|
||||
D.Stream step' (SpliceInitial state)
|
||||
|
||||
where
|
||||
|
||||
nElem = n `quot` sizeOf (undefined :: a)
|
||||
|
||||
{-# INLINE_LATE step' #-}
|
||||
step' gst (SpliceInitial st) = do
|
||||
when (n <= 0) $
|
||||
-- XXX we can pass the module string from the higher level API
|
||||
error $ "Streamly.Internal.Memory.Array.Types.packArraysChunksOf: the size of "
|
||||
++ "arrays [" ++ show n ++ "] must be a natural number"
|
||||
r <- step gst st
|
||||
case r of
|
||||
D.Yield arr s ->
|
||||
if length arr >= nElem
|
||||
then return $ D.Skip (SpliceYielding arr (SpliceInitial s))
|
||||
else do
|
||||
buf <- MA.newArray nElem
|
||||
noff <- spliceInto buf 0 arr
|
||||
return $ D.Skip (SpliceBuffering s (buf, noff))
|
||||
D.Skip s -> return $ D.Skip (SpliceInitial s)
|
||||
D.Stop -> return $ D.Stop
|
||||
|
||||
step' gst (SpliceBuffering st arr2@(buf, boff)) = do
|
||||
r <- step gst st
|
||||
case r of
|
||||
D.Yield arr s -> do
|
||||
if boff + length arr > nElem
|
||||
then do
|
||||
nArr <- unsafeFreeze buf
|
||||
return $ D.Skip (SpliceYielding (slice nArr 0 boff) (SpliceBuffering s arr2))
|
||||
else do
|
||||
noff <- spliceInto buf boff arr
|
||||
return $ D.Skip (SpliceBuffering s (buf, noff))
|
||||
D.Skip s -> return $ D.Skip (SpliceBuffering s arr2)
|
||||
D.Stop -> do
|
||||
nArr <- unsafeFreeze buf
|
||||
return $ D.Skip (SpliceYielding (slice nArr 0 boff) SpliceFinish)
|
||||
|
||||
step' _ SpliceFinish = return D.Stop
|
||||
|
||||
step' _ (SpliceYielding arr next) = return $ D.Yield arr next
|
||||
|
||||
{-# INLINE_NORMAL lpackArraysChunksOf #-}
|
||||
lpackArraysChunksOf ::
|
||||
forall m a. (PrimMonad m, Prim a)
|
||||
=> Int
|
||||
-> Fold m (Array a) ()
|
||||
-> Fold m (Array a) ()
|
||||
lpackArraysChunksOf n (Fold step1 initial1 extract1) =
|
||||
Fold step initial extract
|
||||
|
||||
where
|
||||
|
||||
nElem = n `quot` sizeOf (undefined :: a)
|
||||
|
||||
initial = do
|
||||
when (n <= 0) $
|
||||
-- XXX we can pass the module string from the higher level API
|
||||
error $ "Streamly.Internal.Memory.Array.Types.packArraysChunksOf: the size of "
|
||||
++ "arrays [" ++ show n ++ "] must be a natural number"
|
||||
r1 <- initial1
|
||||
return (Tuple3' Nothing' 0 r1)
|
||||
|
||||
extract (Tuple3' Nothing' _ r1) = extract1 r1
|
||||
extract (Tuple3' (Just' buf) boff r1) = do
|
||||
nArr <- unsafeFreeze buf
|
||||
r <- step1 r1 (slice nArr 0 boff)
|
||||
extract1 r
|
||||
|
||||
step (Tuple3' Nothing' _ r1) arr =
|
||||
|
||||
if length arr >= nElem
|
||||
then do
|
||||
r <- step1 r1 arr
|
||||
extract1 r
|
||||
r1' <- initial1
|
||||
return (Tuple3' Nothing' 0 r1')
|
||||
else do
|
||||
buf <- MA.newArray nElem
|
||||
noff <- spliceInto buf 0 arr
|
||||
return (Tuple3' (Just' buf) noff r1)
|
||||
|
||||
step (Tuple3' (Just' buf) boff r1) arr = do
|
||||
noff <- spliceInto buf boff arr
|
||||
|
||||
if noff >= nElem
|
||||
then do
|
||||
nArr <- unsafeFreeze buf
|
||||
r <- step1 r1 (slice nArr 0 noff)
|
||||
extract1 r
|
||||
r1' <- initial1
|
||||
return (Tuple3' Nothing' 0 r1')
|
||||
else return (Tuple3' (Just' buf) noff r1)
|
||||
|
||||
data SplitState s arr
|
||||
= Initial s
|
||||
| Buffering s arr
|
||||
| Splitting s arr
|
||||
| Yielding arr (SplitState s arr)
|
||||
| Finishing
|
||||
|
||||
-- | Split a stream of arrays on a given separator byte, dropping the separator
|
||||
-- and coalescing all the arrays between two separators into a single array.
|
||||
--
|
||||
-- @since VERSION
|
||||
{-# INLINE_NORMAL splitOn #-}
|
||||
splitOn
|
||||
:: PrimMonad m
|
||||
=> Word8
|
||||
-> D.Stream m (Array Word8)
|
||||
-> D.Stream m (Array Word8)
|
||||
splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
|
||||
|
||||
where
|
||||
|
||||
{-# INLINE_LATE step' #-}
|
||||
step' gst (Initial st) = do
|
||||
r <- step gst st
|
||||
case r of
|
||||
D.Yield arr s -> do
|
||||
(arr1, marr2) <- breakOn byte arr
|
||||
return $ case marr2 of
|
||||
Nothing -> D.Skip (Buffering s arr1)
|
||||
Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
|
||||
D.Skip s -> return $ D.Skip (Initial s)
|
||||
D.Stop -> return $ D.Stop
|
||||
|
||||
step' gst (Buffering st buf) = do
|
||||
r <- step gst st
|
||||
case r of
|
||||
D.Yield arr s -> do
|
||||
(arr1, marr2) <- breakOn byte arr
|
||||
buf' <- spliceTwo buf arr1
|
||||
return $ case marr2 of
|
||||
Nothing -> D.Skip (Buffering s buf')
|
||||
Just x -> D.Skip (Yielding buf' (Splitting s x))
|
||||
D.Skip s -> return $ D.Skip (Buffering s buf)
|
||||
D.Stop -> return $
|
||||
if byteLength buf == 0
|
||||
then D.Stop
|
||||
else D.Skip (Yielding buf Finishing)
|
||||
|
||||
step' _ (Splitting st buf) = do
|
||||
(arr1, marr2) <- breakOn byte buf
|
||||
return $ case marr2 of
|
||||
Nothing -> D.Skip $ Buffering st arr1
|
||||
Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
|
||||
|
||||
step' _ (Yielding arr next) = return $ D.Yield arr next
|
||||
step' _ Finishing = return $ D.Stop
|
||||
180
src/prim-array.hs
Normal file
180
src/prim-array.hs
Normal file
@ -0,0 +1,180 @@
|
||||
import Control.Monad (when)
|
||||
import Control.Monad.Primitive
|
||||
import Data.Primitive.Types
|
||||
|
||||
import Prelude hiding (length, null, last, map, (!!), read, concat)
|
||||
|
||||
import Streamly.Internal.Data.Fold.Types (Fold(..))
|
||||
import Streamly.Internal.Data.Unfold.Types (Unfold(..))
|
||||
import Streamly.Internal.Data.Stream.Serial (SerialT)
|
||||
import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
|
||||
|
||||
import qualified Streamly.Internal.Data.Stream.Prelude as P
|
||||
import qualified Streamly.Internal.Data.Stream.Serial as Serial
|
||||
import qualified Streamly.Internal.Data.Stream.StreamD as D
|
||||
import qualified Streamly.Internal.Data.Stream.StreamK as K
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Construction
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Create an 'Array' from the first N elements of a stream. The array is
|
||||
-- allocated to size N, if the stream terminates before N elements then the
|
||||
-- array may hold less than N elements.
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE fromStreamN #-}
|
||||
fromStreamN :: (PrimMonad m, Prim a) => Int -> SerialT m a -> m (Array a)
|
||||
fromStreamN n m = do
|
||||
when (n < 0) $ error "writeN: negative write count specified"
|
||||
A.fromStreamDN n $ D.toStreamD m
|
||||
|
||||
-- | Create an 'Array' from a stream. This is useful when we want to create a
|
||||
-- single array from a stream of unknown size. 'writeN' is at least twice
|
||||
-- as efficient when the size is already known.
|
||||
--
|
||||
-- Note that if the input stream is too large memory allocation for the array
|
||||
-- may fail. When the stream size is not known, `arraysOf` followed by
|
||||
-- processing of indvidual arrays in the resulting stream should be preferred.
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE fromStream #-}
|
||||
fromStream :: (PrimMonad m, Prim a) => SerialT m a -> m (Array a)
|
||||
fromStream = P.runFold A.write
|
||||
-- write m = A.fromStreamD $ D.toStreamD m
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Elimination
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Convert an 'Array' into a stream.
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE_EARLY toStream #-}
|
||||
toStream :: (PrimMonad m, K.IsStream t, Prim a) => Array a -> t m a
|
||||
toStream = D.fromStreamD . A.toStreamD
|
||||
-- XXX add fallback to StreamK rule
|
||||
-- {-# RULES "Streamly.Array.read fallback to StreamK" [1]
|
||||
-- forall a. S.readK (read a) = K.fromArray a #-}
|
||||
|
||||
-- | Convert an 'Array' into a stream in reverse order.
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE_EARLY toStreamRev #-}
|
||||
toStreamRev :: (PrimMonad m, IsStream t, Prim a) => Array a -> t m a
|
||||
toStreamRev = D.fromStreamD . A.toStreamDRev
|
||||
-- XXX add fallback to StreamK rule
|
||||
-- {-# RULES "Streamly.Array.readRev fallback to StreamK" [1]
|
||||
-- forall a. S.toStreamK (readRev a) = K.revFromArray a #-}
|
||||
|
||||
-- | Unfold an array into a stream.
|
||||
--
|
||||
-- @since 0.7.0
|
||||
{-# INLINE_NORMAL read #-}
|
||||
read :: forall m a. (PrimMonad m, Prim a) => Unfold m (Array a) a
|
||||
read = Unfold step inject
|
||||
where
|
||||
|
||||
inject = return
|
||||
|
||||
{-# INLINE_LATE step #-}
|
||||
step (Array _ _ len) | len == 0 = return D.Stop
|
||||
step arr@(Array arr# off len) =
|
||||
let !x = A.unsafeIndex arr 0
|
||||
in return $ D.Yield x (Array arr# (off + 1) (len - 1))
|
||||
|
||||
-- | Unfold an array into a stream, does not check the end of the array, the
|
||||
-- user is responsible for terminating the stream within the array bounds. For
|
||||
-- high performance application where the end condition can be determined by
|
||||
-- a terminating fold.
|
||||
--
|
||||
-- The following might not be true, not that the representation changed.
|
||||
-- Written in the hope that it may be faster than "read", however, in the case
|
||||
-- for which this was written, "read" proves to be faster even though the core
|
||||
-- generated with unsafeRead looks simpler.
|
||||
--
|
||||
-- /Internal/
|
||||
--
|
||||
{-# INLINE_NORMAL unsafeRead #-}
|
||||
unsafeRead :: forall m a. (PrimMonad m, Prim a) => Unfold m (Array a) a
|
||||
unsafeRead = Unfold step inject
|
||||
where
|
||||
|
||||
inject = return
|
||||
|
||||
{-# INLINE_LATE step #-}
|
||||
step arr@(Array arr# off len) =
|
||||
let !x = A.unsafeIndex arr 0
|
||||
in return $ D.Yield x (Array arr# (off + 1) (len - 1))
|
||||
|
||||
-- | > null arr = length arr == 0
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE null #-}
|
||||
null :: Array a -> Bool
|
||||
null arr = length arr == 0
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Folds
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Fold an array using a 'Fold'.
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE fold #-}
|
||||
fold :: forall m a b. (PrimMonad m, Prim a) => Fold m a b -> Array a -> m b
|
||||
fold f arr = P.runFold f (toStream arr :: Serial.SerialT m a)
|
||||
|
||||
-- | Fold an array using a stream fold operation.
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE streamFold #-}
|
||||
streamFold :: (PrimMonad m, Prim a) => (SerialT m a -> m b) -> Array a -> m b
|
||||
streamFold f arr = f (toStream arr)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Random reads
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | /O(1)/ Lookup the element at the given index, starting from 0.
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE readIndex #-}
|
||||
readIndex :: Prim a => Array a -> Int -> Maybe a
|
||||
readIndex arr i =
|
||||
if i < 0 || i > length arr - 1
|
||||
then Nothing
|
||||
else Just $ A.unsafeIndex arr i
|
||||
|
||||
-- | > last arr = readIndex arr (length arr - 1)
|
||||
--
|
||||
-- /Internal/
|
||||
{-# INLINE last #-}
|
||||
last :: Prim a => Array a -> Maybe a
|
||||
last arr = readIndex arr (length arr - 1)
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
-- Array stream operations
|
||||
-------------------------------------------------------------------------------
|
||||
|
||||
-- | Convert a stream of arrays into a stream of their elements.
|
||||
--
|
||||
-- Same as the following but more efficient:
|
||||
--
|
||||
-- > concat = S.concatMap A.read
|
||||
--
|
||||
-- @since 0.7.0
|
||||
{-# INLINE concat #-}
|
||||
concat :: (IsStream t, PrimMonad m, Prim a) => t m (Array a) -> t m a
|
||||
-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
|
||||
-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
|
||||
concat m = D.fromStreamD $ D.concatMapU read (D.toStreamD m)
|
||||
|
||||
-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
|
||||
-- maximum specified size in bytes.
|
||||
--
|
||||
-- @since 0.7.0
|
||||
{-# INLINE compact #-}
|
||||
compact :: (PrimMonad m, Prim a)
|
||||
=> Int -> SerialT m (Array a) -> SerialT m (Array a)
|
||||
compact n xs = D.fromStreamD $ A.packArraysChunksOf n (D.toStreamD xs)
|
||||
@ -113,6 +113,7 @@ extra-source-files:
|
||||
benchmark/Streamly/Benchmark/Data/*.hs
|
||||
benchmark/Streamly/Benchmark/Data/Parser/*.hs
|
||||
benchmark/Streamly/Benchmark/Data/Prim/*.hs
|
||||
benchmark/Streamly/Benchmark/Data/Prim/Pinned/*.hs
|
||||
benchmark/Streamly/Benchmark/Data/Stream/*.hs
|
||||
benchmark/Streamly/Benchmark/FileSystem/*.hs
|
||||
benchmark/Streamly/Benchmark/FileSystem/Handle/*.hs
|
||||
@ -145,6 +146,9 @@ extra-source-files:
|
||||
examples/README.md
|
||||
src/Streamly/Internal/Data/Stream/Instances.hs
|
||||
src/Streamly/Internal/Data/Time/config-clock.h
|
||||
src/prim-array.hs
|
||||
src/prim-array-types.hs
|
||||
src/mutable-prim-array-types.hs
|
||||
src/Streamly/Internal/Data/Time/config.h.in
|
||||
src/inline.hs
|
||||
stack.yaml
|
||||
@ -382,6 +386,10 @@ library
|
||||
, Streamly.Internal.Data.Array
|
||||
, Streamly.Internal.Data.Prim.Array.Types
|
||||
, Streamly.Internal.Data.Prim.Array
|
||||
, Streamly.Internal.Data.Prim.Mutable.Array.Types
|
||||
, Streamly.Internal.Data.Prim.Pinned.Array.Types
|
||||
, Streamly.Internal.Data.Prim.Pinned.Array
|
||||
, Streamly.Internal.Data.Prim.Pinned.Mutable.Array.Types
|
||||
, Streamly.Internal.Data.SmallArray.Types
|
||||
, Streamly.Internal.Data.SmallArray
|
||||
, Streamly.Internal.Memory.Array.Types
|
||||
@ -440,6 +448,7 @@ library
|
||||
, Streamly.Internal.Data.Unicode.Stream
|
||||
, Streamly.Internal.Data.Unicode.Char
|
||||
, Streamly.Internal.Memory.Unicode.Array
|
||||
, Streamly.Internal.Data.Prim.Pinned.Unicode.Array
|
||||
|
||||
if !impl(ghcjs)
|
||||
exposed-modules:
|
||||
@ -672,6 +681,23 @@ test-suite primarray-test
|
||||
transformers >= 0.4 && < 0.6
|
||||
default-language: Haskell2010
|
||||
|
||||
test-suite primpinnedarray-test
|
||||
import: test-options
|
||||
type: exitcode-stdio-1.0
|
||||
main-is: Streamly/Test/Array.hs
|
||||
js-sources: jsbits/clock.js
|
||||
hs-source-dirs: test
|
||||
cpp-options: -DTEST_PRIM_PINNED_ARRAY
|
||||
build-depends:
|
||||
streamly
|
||||
, base >= 4.8 && < 5
|
||||
, QuickCheck >= 2.10 && < 2.15
|
||||
, hspec >= 2.0 && < 3
|
||||
if impl(ghc < 8.0)
|
||||
build-depends:
|
||||
transformers >= 0.4 && < 0.6
|
||||
default-language: Haskell2010
|
||||
|
||||
test-suite string-test
|
||||
import: test-options
|
||||
type: exitcode-stdio-1.0
|
||||
|
||||
@ -25,13 +25,20 @@ import qualified Streamly.Prelude as S
|
||||
import qualified Streamly.Internal.Data.SmallArray as A
|
||||
type Array = A.SmallArray
|
||||
#elif defined(TEST_ARRAY)
|
||||
import qualified Streamly.Memory.Array as A
|
||||
import qualified Streamly.Internal.Memory.Array as A
|
||||
import qualified Streamly.Internal.Memory.Array.Types as A
|
||||
import qualified Streamly.Internal.Prelude as IP
|
||||
type Array = A.Array
|
||||
#elif defined(TEST_PRIM_PINNED_ARRAY)
|
||||
import qualified Streamly.Internal.Data.Prim.Pinned.Array as A
|
||||
import qualified Streamly.Internal.Data.Prim.Pinned.Array.Types as A
|
||||
import qualified Streamly.Internal.Prelude as IP
|
||||
type Array = A.Array
|
||||
#elif defined(TEST_PRIM_ARRAY)
|
||||
import qualified Streamly.Internal.Data.Prim.Array as A
|
||||
type Array = A.PrimArray
|
||||
import qualified Streamly.Internal.Data.Prim.Array.Types as A
|
||||
import qualified Streamly.Internal.Prelude as IP
|
||||
type Array = A.Array
|
||||
#else
|
||||
import qualified Streamly.Internal.Data.Array as A
|
||||
type Array = A.Array
|
||||
@ -122,17 +129,27 @@ testFoldUnfold :: Property
|
||||
testFoldUnfold = genericTestFromTo (const (S.fold A.write)) (S.unfold A.read) (==)
|
||||
#endif
|
||||
|
||||
#ifdef TEST_ARRAY
|
||||
#if defined(TEST_ARRAY) ||\
|
||||
defined(TEST_PRIM_ARRAY) ||\
|
||||
defined(TEST_PRIM_PINNED_ARRAY)
|
||||
|
||||
testArraysOf :: Property
|
||||
testArraysOf =
|
||||
forAll (choose (0, maxArrLen)) $ \len ->
|
||||
forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
|
||||
monadicIO $ do
|
||||
xs <- S.toList
|
||||
xs <- run
|
||||
$ S.toList
|
||||
$ S.concatUnfold A.read
|
||||
$ IP.arraysOf 240
|
||||
$ arraysOf 240
|
||||
$ S.fromList list
|
||||
assert (xs == list)
|
||||
where
|
||||
arraysOf n = IP.chunksOf n (A.writeNUnsafe n)
|
||||
|
||||
#endif
|
||||
|
||||
#ifdef TEST_ARRAY
|
||||
|
||||
lastN :: Int -> [a] -> [a]
|
||||
lastN n l = drop (length l - n) l
|
||||
@ -143,7 +160,8 @@ testLastN =
|
||||
forAll (choose (0, len)) $ \n ->
|
||||
forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
|
||||
monadicIO $ do
|
||||
xs <- fmap A.toList
|
||||
xs <- run
|
||||
$ fmap A.toList
|
||||
$ S.fold (A.lastN n)
|
||||
$ S.fromList list
|
||||
assert (xs == lastN n list)
|
||||
@ -174,9 +192,13 @@ main =
|
||||
prop "toStream . fromStream === id" testFromStreamToStream
|
||||
prop "read . write === id" testFoldUnfold
|
||||
#endif
|
||||
#ifdef TEST_ARRAY
|
||||
|
||||
#if defined(TEST_ARRAY) ||\
|
||||
defined(TEST_PRIM_ARRAY) ||\
|
||||
defined(TEST_PRIM_PINNED_ARRAY)
|
||||
prop "arraysOf concats to original" testArraysOf
|
||||
#endif
|
||||
|
||||
#ifdef TEST_ARRAY
|
||||
describe "Fold" $ do
|
||||
prop "lastN : 0 <= n <= len" $ testLastN
|
||||
|
||||
Loading…
Reference in New Issue
Block a user