Add fromByteStr# and fromStr#

Routines to generate byte/char streams from Haskell string literals.

core/src/Streamly/Internal/Data/Stream/Generate.hs

@@ -30,16 +30,20 @@ module Streamly.Internal.Data.Stream.Generate
 
     -- * From memory
     , fromPtr
+    , fromPtrN
+    , fromByteStr#
+ -- , fromByteArray#
     )
 where
 
 import Control.Monad.IO.Class (MonadIO)
+import Data.Word (Word8)
 import Foreign.Storable (Storable)
-import Foreign.Ptr (Ptr)
+import GHC.Exts (Addr#, Ptr (Ptr))
 import Streamly.Internal.Data.Stream.Type (Stream, fromStreamK, toStreamK)
 import Streamly.Internal.Data.Unfold.Type (Unfold)
 
-import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D
+import qualified Streamly.Internal.Data.Stream.StreamD as D
 import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
 import qualified Streamly.Internal.Data.Stream.Type as Stream
 
@@ -48,6 +52,7 @@ import qualified Streamly.Internal.Data.Stream.Type as Stream
 -- >>> import qualified Streamly.Data.Fold as Fold
 -- >>> import qualified Streamly.Data.Unfold as Unfold
 -- >>> import qualified Streamly.Internal.Data.Stream as Stream
+-- >>> import GHC.Exts (Ptr (Ptr))
 
 
 ------------------------------------------------------------------------------
@@ -151,6 +156,41 @@ fromFoldableM = Prelude.foldr Stream.consM Stream.nil
 -- From pointers
 ------------------------------------------------------------------------------
 
+-- | Keep reading 'Storable' elements from 'Ptr' onwards.
+--
+-- /Unsafe:/ The caller is responsible for safe addressing.
+--
+-- /Pre-release/
 {-# INLINE fromPtr #-}
-fromPtr :: (MonadIO m, Storable a) => Int -> Ptr a -> Stream m a
-fromPtr n = Stream.fromStreamD . D.fromPtr n
+fromPtr :: (MonadIO m, Storable a) => Ptr a -> Stream m a
+fromPtr = Stream.fromStreamD . D.fromPtr
+
+-- | Take @n@ 'Storable' elements starting from 'Ptr' onwards.
+--
+-- >>> fromPtrN n = Stream.take n . Stream.fromPtr
+--
+-- /Unsafe:/ The caller is responsible for safe addressing.
+--
+-- /Pre-release/
+{-# INLINE fromPtrN #-}
+fromPtrN :: (MonadIO m, Storable a) => Int -> Ptr a -> Stream m a
+fromPtrN n = Stream.fromStreamD . D.take n . D.fromPtr
+
+-- | Read bytes from an 'Addr#' until a 0 byte is encountered, the 0 byte is
+-- not included in the stream.
+--
+-- >>> fromByteStr# addr = Stream.takeWhile (/= 0) $ Stream.fromPtr $ Ptr addr
+--
+-- /Unsafe:/ The caller is responsible for safe addressing.
+--
+-- Note that this is completely safe when reading from Haskell string
+-- literals because they are guaranteed to be NULL terminated:
+--
+-- >>> Stream.fold Fold.toList $ Stream.fromByteStr# "\1\2\3\0"#
+-- [1,2,3]
+--
+{-# INLINE fromByteStr# #-}
+fromByteStr# :: MonadIO m => Addr# -> Stream m Word8
+fromByteStr# addr =
+    Stream.fromStreamD $ D.takeWhile (/= 0) $ D.fromPtr $ Ptr addr
+

core/src/Streamly/Internal/Data/Stream/StreamD/Generate.hs

@@ -470,14 +470,15 @@ fromListM = Stream step
 -- From pointers
 -------------------------------------------------------------------------------
 
+-- | Read an infinite stream from a pointer, advancing the pointer as needed.
+-- The caller is responsible to end the stream safely.
 {-# INLINE fromPtr #-}
-fromPtr :: forall m a. (MonadIO m, Storable a) => Int -> Ptr a -> Stream m a
-fromPtr count ptr = Stream step (count, ptr)
+fromPtr :: forall m a. (MonadIO m, Storable a) => Ptr a -> Stream m a
+fromPtr = Stream step
 
     where
 
     {-# INLINE_LATE step #-}
-    step _ (0, _) = return Stop
-    step _ (n, p) = do
+    step _ p = do
         x <- liftIO $ peek p
-        return $ Yield x (n - 1, PTR_NEXT(p, a))
+        return $ Yield x (PTR_NEXT(p, a))

core/src/Streamly/Internal/Data/Unfold.hs

@@ -177,7 +177,7 @@ module Streamly.Internal.Data.Unfold
     , fromList
     , fromListM
 
-    -- ** From Pointer
+    -- ** From Memory
     , fromPtr
 
     , fromStreamK
@@ -636,16 +636,15 @@ fromListM = Unfold step pure
     step [] = pure Stop
 
 {-# INLINE fromPtr #-}
-fromPtr :: forall m a. (MonadIO m, Storable a) => Unfold m (Int, Ptr a) a
+fromPtr :: forall m a. (MonadIO m, Storable a) => Unfold m (Ptr a) a
 fromPtr = Unfold step return
 
     where
 
     {-# INLINE_LATE step #-}
-    step (0, _) = return Stop
-    step (n, p) = do
+    step p = do
         x <- liftIO $ peek p
-        return $ Yield x (n - 1, PTR_NEXT(p, a))
+        return $ Yield x (PTR_NEXT(p, a))
 
 ------------------------------------------------------------------------------
 -- Specialized Generation

src/Streamly/Internal/Unicode/Stream.hs

@@ -72,6 +72,9 @@ module Streamly.Internal.Unicode.Stream
     , decodeUtf8ArraysD'
     , decodeUtf8ArraysD_
 
+    -- * Decoding String Literals
+    , fromStr#
+
     -- * Deprecations
     , decodeUtf8Lax
     , encodeLatin1Lax
@@ -94,6 +97,7 @@ import Foreign.Marshal.Alloc (mallocBytes)
 import Foreign.Storable (Storable(..))
 import Fusion.Plugin.Types (Fuse(..))
 import GHC.Base (assert, unsafeChr)
+import GHC.Exts (Addr#)
 import GHC.IO.Encoding.Failure (isSurrogate)
 import GHC.Ptr (Ptr (..), plusPtr)
 import System.IO.Unsafe (unsafePerformIO)
@@ -114,6 +118,7 @@ import qualified Streamly.Internal.Data.Fold as Fold
 import qualified Streamly.Internal.Data.Unfold as Unfold
 import qualified Streamly.Internal.Data.Parser as Parser
 import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
+import qualified Streamly.Internal.Data.Stream as Stream
 import qualified Streamly.Internal.Data.Stream.Serial as Serial
 import qualified Streamly.Internal.Data.Array.Unboxed as Array
 import qualified Streamly.Internal.Data.Array.Unboxed.Type as A
@@ -124,9 +129,11 @@ import Prelude hiding (lines, words, unlines, unwords)
 
 -- $setup
 -- >>> :m
+-- >>> :set -XMagicHash
 -- >>> import Prelude hiding (lines, words, unlines, unwords)
 -- >>> import qualified Streamly.Prelude as Stream
 -- >>> import qualified Streamly.Data.Fold as Fold
+-- >>> import qualified Streamly.Internal.Unicode.Stream as Unicode
 -- >>> import Streamly.Internal.Unicode.Stream
 
 -------------------------------------------------------------------------------
@@ -980,6 +987,25 @@ encodeUtf8_ = fromStreamD . encodeUtf8D_ . toStreamD
 encodeUtf8Lax :: (IsStream t, Monad m) => t m Char -> t m Word8
 encodeUtf8Lax = encodeUtf8
 
+-------------------------------------------------------------------------------
+-- Decoding string literals
+-------------------------------------------------------------------------------
+
+-- | Read UTF-8 encoded bytes as chars from an 'Addr#' until a 0 byte is
+-- encountered, the 0 byte is not included in the stream.
+--
+-- /Unsafe:/ The caller is responsible for safe addressing.
+--
+-- Note that this is completely safe when reading from Haskell string
+-- literals because they are guaranteed to be NULL terminated:
+--
+-- >>> Stream.fold Fold.toList (Unicode.fromStr# "Haskell"#)
+-- "Haskell"
+--
+{-# INLINE fromStr# #-}
+fromStr# :: MonadIO m => Addr# -> SerialT m Char
+fromStr# addr = decodeUtf8 $ Stream.fromByteStr# addr
+
 -------------------------------------------------------------------------------
 -- Encode streams of containers
 -------------------------------------------------------------------------------