090d1d8 · streamly · git.iko.soy

git.iko.soy/streamly

Add fromPtrN and fromByteStr

commit090d1d8a82df85c3913962dfa9f519c18763a1e3 authorHarendra Kumar <harendra@composewell.com> dateMon 25 Dec 2023, 09:25 (3 years ago) parent994c442

3 files changed+64−3

compare with master

unifiedsplit

core/src/Streamly/Internal/Data/Array/Type.hs+13−2 core/src/Streamly/Internal/Data/MutArray/Type.hs+40−1 core/src/Streamly/Internal/Data/MutByteArray/Type.hs+11

core/src/Streamly/Internal/Data/Array/Type.hs+13−2

@@ -42,6 +42,8 @@ module Streamly.Internal.Data.Array.Type

4242 , fromPureStreamN

4343 , fromPureStream

4444 , fromByteStr#

45+ , fromByteStr

46+ , fromPtrN

4547

4648 -- ** Split

4749 , breakOn

@@ -121,7 +123,6 @@ import Prelude hiding (Foldable(..), read, unlines, splitAt)

121123import qualified GHC.Exts as Exts

122124import qualified Streamly.Internal.Data.MutArray.Type as MA

123125import qualified Streamly.Internal.Data.Stream.Type as D

124−import qualified Streamly.Internal.Data.Stream.Generate as D

125126import qualified Streamly.Internal.Data.StreamK.Type as K

126127import qualified Streamly.Internal.Data.MutByteArray.Type as Unboxed

127128import qualified Streamly.Internal.Data.Unfold.Type as Unfold

@@ -633,6 +634,13 @@ fromPureStream x = unsafePerformIO $ fmap (unsafeFreeze) (MA.fromPureStream x)

633634-- fromPureStream = runIdentity . D.fold (unsafeMakePure write)

634635-- fromPureStream = fromList . runIdentity . D.toList

635636

637+-- | Copy an immutable 'Ptr Word8' sequence into an array.

638+--

639+-- /Unsafe:/ The caller is responsible for safe addressing.

640+--

641+fromPtrN :: Int -> Ptr Word8 -> Array Word8

642+fromPtrN n addr = unsafePerformIO $ fmap unsafeFreeze (MA.fromPtrN n addr)

643+

636644-- | Copy a null terminated immutable 'Addr#' Word8 sequence into an array.

637645--

638646-- /Unsafe:/ The caller is responsible for safe addressing.

@@ -644,7 +652,10 @@ fromPureStream x = unsafePerformIO $ fmap (unsafeFreeze) (MA.fromPureStream x)

644652-- [1,2,3]

645653--

646654fromByteStr# :: Addr# -> Array Word8

647−fromByteStr# addr = fromPureStream (D.fromByteStr# addr)

655+fromByteStr# addr = unsafePerformIO $ fmap unsafeFreeze (MA.fromByteStr# addr)

656+

657+fromByteStr :: Ptr Word8 -> Array Word8

658+fromByteStr (Ptr addr#) = fromByteStr# addr#

648659

649660-------------------------------------------------------------------------------

650661-- Instances

core/src/Streamly/Internal/Data/MutArray/Type.hs+40−1

@@ -76,6 +76,8 @@ module Streamly.Internal.Data.MutArray.Type

7676 , fromStreamD

7777 , fromPureStreamN

7878 , fromPureStream

79+ , fromByteStr#

80+ , fromPtrN

7981

8082 -- ** Random writes

8183 , putIndex

@@ -173,6 +175,7 @@ module Streamly.Internal.Data.MutArray.Type

173175 , castUnsafe

174176 , asBytes

175177 , asPtrUnsafe

178+ , asUnpinnedPtrUnsafe

176179

177180 -- ** Folding

178181 , foldl'

@@ -255,7 +258,7 @@ import GHC.Base

255258 , copyMutableByteArray#

256259 )

257260import GHC.Base (noinline)

258−import GHC.Exts (unsafeCoerce#)

261+import GHC.Exts (unsafeCoerce#, Addr#)

259262import GHC.Ptr (Ptr(..))

260263

261264import Streamly.Internal.Data.Fold.Type (Fold(..))

@@ -283,6 +286,7 @@ import Prelude hiding

283286-- Foreign helpers

284287-------------------------------------------------------------------------------

285288

289+-- NOTE: Have to be "ccall unsafe" so that we can pass unpinned memory to these

286290foreign import ccall unsafe "string.h memcpy" c_memcpy

287291 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO (Ptr Word8)

288292

@@ -292,6 +296,9 @@ foreign import ccall unsafe "string.h memchr" c_memchr

292296foreign import ccall unsafe "string.h memcmp" c_memcmp

293297 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt

294298

299+foreign import ccall unsafe "string.h strlen" c_strlen

300+ :: Ptr Word8 -> IO CSize

301+

295302-- | Given an 'Unboxed' type (unused first arg) and a number of bytes, return

296303-- how many elements of that type will completely fit in those bytes.

297304--

@@ -2042,6 +2049,32 @@ fromPureStream :: (MonadIO m, Unbox a) => Stream Identity a -> m (MutArray a)

20422049fromPureStream xs =

20432050 liftIO $ D.fold write $ D.morphInner (return . runIdentity) xs

20442051

2052+{-# INLINABLE fromPtrN #-}

2053+fromPtrN :: MonadIO m => Int -> Ptr Word8 -> m (MutArray Word8)

2054+fromPtrN len addr = do

2055+ -- memcpy is better than stream copy when the size is known.

2056+ -- XXX We can implement a stream copy in a similar way by streaming Word64

2057+ -- first and then remaining Word8.

2058+ arr <- new len

2059+ _ <- asUnpinnedPtrUnsafe arr

2060+ (\ptr -> liftIO $ c_memcpy ptr addr (fromIntegral len))

2061+ return (arr {arrEnd = len})

2062+

2063+{-# INLINABLE fromByteStr# #-}

2064+fromByteStr# :: MonadIO m => Addr# -> m (MutArray Word8)

2065+fromByteStr# addr = do

2066+ -- It is better to count the size first and allocate exact space.

2067+ -- Also, memcpy is better than stream copy when the size is known.

2068+ -- C strlen compares 4 bytes at a time, so is better than the stream

2069+ -- version. https://github.com/bminor/glibc/blob/master/string/strlen.c

2070+ -- XXX We can possibly use a stream of Word64 to do the same.

2071+ -- fromByteStr# addr = fromPureStream (D.fromByteStr# addr)

2072+ len <- liftIO $ c_strlen (Ptr addr)

2073+ let lenInt = fromIntegral len

2074+ arr <- new lenInt

2075+ _ <- asUnpinnedPtrUnsafe arr (\ptr -> liftIO $ c_memcpy ptr (Ptr addr) len)

2076+ return (arr {arrEnd = lenInt})

2077+

20452078-------------------------------------------------------------------------------

20462079-- convert stream to a single array

20472080-------------------------------------------------------------------------------

@@ -2379,6 +2412,12 @@ asPtrUnsafe arr f =

23792412 Unboxed.asPtrUnsafe

23802413 (arrContents arr) (\ptr -> f (ptr `plusPtr` arrStart arr))

23812414

2415+{-# INLINE asUnpinnedPtrUnsafe #-}

2416+asUnpinnedPtrUnsafe :: MonadIO m => MutArray a -> (Ptr a -> m b) -> m b

2417+asUnpinnedPtrUnsafe arr f =

2418+ Unboxed.asUnpinnedPtrUnsafe

2419+ (arrContents arr) (\ptr -> f (ptr `plusPtr` arrStart arr))

2420+

23822421-------------------------------------------------------------------------------

23832422-- Equality

23842423-------------------------------------------------------------------------------

core/src/Streamly/Internal/Data/MutByteArray/Type.hs+11

@@ -35,6 +35,7 @@ module Streamly.Internal.Data.MutByteArray.Type

3535 , cloneSliceUnsafe

3636 , pinnedCloneSliceUnsafe

3737 , asPtrUnsafe

38+ , asUnpinnedPtrUnsafe

3839 ) where

3940

4041import Control.Monad.IO.Class (MonadIO(..))

@@ -120,6 +121,16 @@ asPtrUnsafe arr f = do

120121 liftIO $ touch contents

121122 return r

122123

124+-- | For use with unsafe FFI functions. Does not force pin the array memory.

125+{-# INLINE asUnpinnedPtrUnsafe #-}

126+asUnpinnedPtrUnsafe :: MonadIO m => MutByteArray -> (Ptr a -> m b) -> m b

127+asUnpinnedPtrUnsafe arr f = do

128+ let !ptr = Ptr (byteArrayContents#

129+ (unsafeCoerce# (getMutableByteArray# arr)))

130+ r <- f ptr

131+ liftIO $ touch arr

132+ return r

133+

123134--------------------------------------------------------------------------------

124135-- Creation

125136--------------------------------------------------------------------------------