variant using text internals in place of text streamDecodeUtf8

2 files changed, 216 insertions, 77 deletions

diff --git a/Pipes/Text.hs b/Pipes/Text.hs
index a5859a3..6845dd3 100644
--- a/Pipes/Text.hs
+++ b/Pipes/Text.hs
@@ -81,10 +81,6 @@ module Pipes.Text  (
     filter,
     scan,
     encodeUtf8,
-#if MIN_VERSION_text(0,11,4)
-    pipeDecodeUtf8,
-    pipeDecodeUtf8With,
-#endif
     pack,
     unpack,
     toCaseFold,
@@ -119,10 +115,8 @@ module Pipes.Text  (
     group,
     lines,
     words,
-#if MIN_VERSION_text(0,11,4)
     decodeUtf8,
     decodeUtf8With,
-#endif
     -- * Transformations
     intersperse,
     
@@ -167,6 +161,7 @@ import qualified GHC.IO.Exception as G
 import Pipes
 import qualified Pipes.ByteString as PB
 import qualified Pipes.ByteString.Parse as PBP
+import qualified Pipes.Text.Internal as PE
 import Pipes.Text.Parse (
     nextChar, drawChar, unDrawChar, peekChar, isEndOfChars )
 import Pipes.Core (respond, Server')
@@ -214,43 +209,60 @@ fromLazy  = foldrChunks (\e a -> yield e >> a) (return ())
 {-# INLINABLE fromLazy #-}
 
 -- | Stream text from 'stdin'
-stdin :: MonadIO m => Producer' Text m ()
+stdin :: MonadIO m => Producer' Text m (Producer ByteString m ())
 stdin = fromHandle IO.stdin
 {-# INLINABLE stdin #-}
 
 {-| Convert a 'IO.Handle' into a text stream using a text size 
     determined by the good sense of the text library. 
-
 -}
 
-fromHandle :: MonadIO m => IO.Handle -> Producer' Text m ()
-#if MIN_VERSION_text(0,11,4)
-fromHandle h = go TE.streamDecodeUtf8 where
+fromHandle :: MonadIO m => IO.Handle -> Producer' Text m (Producer ByteString m ())
+-- TODO: this should perhaps just be `decodeUtf8 (PB.fromHandle h)`
+-- if only so that mistakes can be concentrated in one place.
+-- This modifies something that was faster on an earlier iteration.
+-- Note also that the `text` replacement system is being ignored;
+-- with a replacement scheme one could have `Producer Text m ()`
+-- the relation to the replacement business needs to be thought out.
+-- The complicated type seems overmuch for the toy stdin above
+fromHandle h = go PE.streamDecodeUtf8 B.empty where
   act = B.hGetSome h defaultChunkSize
-  go dec = do chunk <- liftIO act
-              case dec chunk of 
-                TE.Some text _ dec' -> do yield text
-                                          unless (B.null chunk) (go dec')
+  go dec old = do chunk <- liftIO act
+                  if B.null chunk 
+                    then if B.null old then return (return ())
+                                       else return (yield old >> return ())
+                    else case dec chunk of 
+                           PE.Some text bs dec' -> 
+                              if T.null text then go dec' (B.append old bs) 
+                                             else do yield text
+                                                     go dec' B.empty
+                           PE.Other text bs ->
+                              if T.null text then return (do yield old
+                                                             yield bs
+                                                             PB.fromHandle h)
+                                             else do yield text
+                                                     return (do yield bs
+                                                                PB.fromHandle h)
 {-# INLINE fromHandle#-}
 -- bytestring fromHandle + streamDecodeUtf8 is 3 times as fast as
 -- the dedicated Text IO function 'hGetChunk' ;
 -- this way "runEffect $ PT.fromHandle hIn  >->  PT.toHandle hOut"
 -- runs the same as the conduit equivalent, only slightly slower 
 -- than "runEffect $ PB.fromHandle hIn  >->  PB.toHandle hOut"
-#else
-fromHandle h = go where
-    go = do txt <- liftIO (T.hGetChunk h)
-            unless (T.null txt) $ do yield txt
-                                     go
-{-# INLINABLE fromHandle#-}
-#endif
+-- #else
+-- fromHandle h = go where
+--     go = do txt <- liftIO (T.hGetChunk h)
+--             unless (T.null txt) $ do yield txt
+--                                      go
+-- {-# INLINABLE fromHandle#-}
+-- #endif
 {-| Stream text from a file using Pipes.Safe
 
 >>> runSafeT $ runEffect $ Text.readFile "hello.hs" >-> Text.map toUpper >-> hoist lift Text.stdout
 MAIN = PUTSTRLN "HELLO WORLD"
 -}
 
-readFile :: (MonadSafe m, Base m ~ IO) => FilePath -> Producer' Text m ()
+readFile :: (MonadSafe m, Base m ~ IO) => FilePath -> Producer' Text m (Producer ByteString m ())
 readFile file = Safe.withFile file IO.ReadMode fromHandle
 {-# INLINABLE readFile #-}
 
@@ -610,74 +622,44 @@ count :: (Monad m, Num n) => Text -> Producer Text m () -> m n
 count c p = P.fold (+) 0 id (p >-> P.map (fromIntegral . T.count c))
 {-# INLINABLE count #-}
 
-#if MIN_VERSION_text(0,11,4)
 -- | Transform a Pipe of 'ByteString's expected to be UTF-8 encoded
 -- into a Pipe of Text
 decodeUtf8
   :: Monad m
   => Producer ByteString m r -> Producer Text m (Producer ByteString m r)
-decodeUtf8 = go TE.streamDecodeUtf8
-  where go dec p = do
-            x <- lift (next p)
-            case x of
-                Left r -> return (return r)
-                Right (chunk, p') -> do
-                    let TE.Some text l dec' = dec chunk
-                    if B.null l
-                      then do
-                          yield text
-                          go dec' p'
-                      else return $ do
-                          yield l
-                          p'
+decodeUtf8 = decodeUtf8With Nothing
 {-# INLINEABLE decodeUtf8 #-}
 
 -- | Transform a Pipe of 'ByteString's expected to be UTF-8 encoded
 -- into a Pipe of Text with a replacement function of type @String -> Maybe Word8 -> Maybe Char@
 -- E.g. 'Data.Text.Encoding.Error.lenientDecode', which simply replaces bad bytes with \"�\"
-decodeUtf8With 
+decodeUtf8With
   :: Monad m  
-  => TE.OnDecodeError 
+  => Maybe TE.OnDecodeError 
   -> Producer ByteString m r -> Producer Text m (Producer ByteString m r)
-decodeUtf8With onErr = go (TE.streamDecodeUtf8With onErr)
-  where go dec p = do
-            x <- lift (next p)
-            case x of
-                Left r -> return (return r)
-                Right (chunk, p') -> do
-                    let TE.Some text l dec' = dec chunk
-                    if B.null l
-                      then do
-                          yield text
-                          go dec' p'
-                      else return $ do
-                          yield l
-                          p'
+decodeUtf8With onErr = go (PE.streamDecodeUtf8With onErr) B.empty where 
+  go dec old p = do
+    x <- lift (next p)
+    case x of
+      Left r -> if B.null old then return (return r)
+                              else return (do yield old 
+                                              return r)
+      Right (chunk, p') -> 
+        case dec chunk of 
+          PE.Some text l dec' -> 
+            if T.null text then go dec' (B.append old l) p'
+                           else do yield text
+                                   go dec' B.empty p'
+          PE.Other text bs ->
+            if T.null text then return (do yield old 
+                                           yield bs
+                                           p')
+                           else do yield text
+                                   return (do yield bs
+                                              p')
 {-# INLINEABLE decodeUtf8With #-}
 
--- | A simple pipe from 'ByteString' to 'Text'; a decoding error will arise
--- with any chunk that contains a sequence of bytes that is unreadable. Otherwise
--- only few bytes will only be moved from one chunk to the next before decoding.
-pipeDecodeUtf8 :: Monad m => Pipe ByteString Text m r
-pipeDecodeUtf8 = go TE.streamDecodeUtf8
-  where go dec = do chunk <- await
-                    case dec chunk of 
-                      TE.Some text l dec' -> do yield text
-                                                go dec'
-{-# INLINEABLE pipeDecodeUtf8 #-}
-
--- | A simple pipe from 'ByteString' to 'Text' using a replacement function.
-pipeDecodeUtf8With 
-  :: Monad m  
-  => TE.OnDecodeError 
-  -> Pipe ByteString Text m r 
-pipeDecodeUtf8With onErr = go (TE.streamDecodeUtf8With onErr)
-  where go dec = do chunk <- await
-                    case dec chunk of 
-                      TE.Some text l dec' -> do yield text
-                                                go dec'
-{-# INLINEABLE pipeDecodeUtf8With #-}
-#endif
+
 
 -- | Splits a 'Producer' after the given number of characters
 splitAt
diff --git a/Pipes/Text/Internal.hs b/Pipes/Text/Internal.hs
new file mode 100644
index 0000000..05d9887
--- /dev/null
+++ b/Pipes/Text/Internal.hs
@@ -0,0 +1,157 @@
+{-# LANGUAGE BangPatterns, CPP, ForeignFunctionInterface, GeneralizedNewtypeDeriving, MagicHash,
+    UnliftedFFITypes #-}
+-- This module lifts material from Brian O'Sullivan's text package 
+-- especially Data.Text.Encoding in order to define a pipes-appropriate
+-- streamDecodeUtf8
+module Pipes.Text.Internal 
+    ( Decoding(..)
+    , streamDecodeUtf8With
+    , streamDecodeUtf8
+    ) where
+
+import Control.Exception (evaluate, try)
+#if __GLASGOW_HASKELL__ >= 702
+import Control.Monad.ST.Unsafe (unsafeIOToST, unsafeSTToIO)
+import Control.Monad.ST (ST, runST)
+#else
+import Control.Monad.ST (unsafeIOToST, unsafeSTToIO, ST, runST)
+#endif
+import Data.Bits ((.&.))
+import Data.ByteString as B
+import Data.ByteString.Internal as B
+import Data.Text ()
+import Data.Text.Encoding.Error (OnDecodeError, UnicodeException, strictDecode)
+import Data.Text.Internal (Text(..), safe, textP)
+import Data.Word (Word8, Word32)
+import Foreign.C.Types (CSize)
+import Foreign.ForeignPtr (withForeignPtr)
+import Foreign.Marshal.Utils (with)
+import Foreign.Ptr (Ptr, minusPtr, nullPtr, plusPtr)
+import Foreign.Storable (Storable, peek, poke)
+import GHC.Base hiding (ord)
+import GHC.Word
+import qualified Data.Text.Array as A
+import GHC.Exts (Char(..), Int(..), chr#, ord#, word2Int#)
+import GHC.Word (Word8(..), Word16(..), Word32(..))
+
+import Data.Text.Unsafe (unsafeDupablePerformIO)
+
+#include "pipes_text_cbits.h"
+
+-- | A stream oriented decoding result.
+data Decoding = Some Text ByteString (ByteString -> Decoding)
+              | Other Text ByteString
+instance Show Decoding where
+    showsPrec d (Some t bs _) = showParen (d > prec) $
+                                showString "Some " . showsPrec prec' t .
+                                showChar ' ' . showsPrec prec' bs .
+                                showString " _"
+      where prec = 10; prec' = prec + 1
+    showsPrec d (Other t bs)  = showParen (d > prec) $
+                                showString "Other " . showsPrec prec' t .
+                                showChar ' ' . showsPrec prec' bs .
+                                showString " _"
+      where prec = 10; prec' = prec + 1
+      
+newtype CodePoint = CodePoint Word32 deriving (Eq, Show, Num, Storable)
+newtype DecoderState = DecoderState Word32 deriving (Eq, Show, Num, Storable)
+
+-- | Decode, in a stream oriented way, a 'ByteString' containing UTF-8
+-- encoded text that is known to be valid.
+--
+-- If the input contains any invalid UTF-8 data, an exception will be
+-- thrown (either by this function or a continuation) that cannot be
+-- caught in pure code.  For more control over the handling of invalid
+-- data, use 'streamDecodeUtf8With'.
+streamDecodeUtf8 :: ByteString -> Decoding
+streamDecodeUtf8 = streamDecodeUtf8With (Just strictDecode)
+
+-- | Decode, in a stream oriented way, a 'ByteString' containing UTF-8
+-- encoded text.
+streamDecodeUtf8With :: Maybe OnDecodeError -> ByteString -> Decoding
+streamDecodeUtf8With mErr = case mErr of 
+    Nothing    -> decodeWith False strictDecode 
+    Just onErr -> decodeWith True onErr 
+ where
+  -- We create a slightly larger than necessary buffer to accommodate a
+  -- potential surrogate pair started in the last buffer
+ decodeWith replace onErr = decodeChunk 0 0
+  where
+  decodeChunk :: CodePoint -> DecoderState -> ByteString -> Decoding
+  decodeChunk codepoint0 state0 bs@(PS fp off len) =
+    runST $ (unsafeIOToST . decodeChunkToBuffer) =<< A.new (len+1)
+   where
+    decodeChunkToBuffer :: A.MArray s -> IO Decoding
+    decodeChunkToBuffer dest = withForeignPtr fp $ \ptr ->
+      with (0::CSize) $ \destOffPtr ->
+      with codepoint0 $ \codepointPtr ->
+      with state0 $ \statePtr ->
+      with nullPtr $ \curPtrPtr ->
+        let end = ptr `plusPtr` (off + len)
+            loop curPtr = do
+              poke curPtrPtr curPtr
+              curPtr' <- c_decode_utf8_with_state (A.maBA dest) destOffPtr
+                         curPtrPtr end codepointPtr statePtr
+              state <- peek statePtr
+              case state of
+                UTF8_REJECT ->  
+                  -- We encountered an encoding error
+                 if replace 
+                 then do 
+                  x <- peek curPtr'
+                  case onErr desc (Just x) of
+                    Nothing -> loop $ curPtr' `plusPtr` 1
+                    Just c -> do
+                      destOff <- peek destOffPtr
+                      w <- unsafeSTToIO $
+                           unsafeWrite dest (fromIntegral destOff) (safe c)
+                      poke destOffPtr (destOff + fromIntegral w)
+                      poke statePtr 0
+                      loop $ curPtr' `plusPtr` 1
+                 else do 
+                  n <- peek destOffPtr 
+                  chunkText <- unsafeSTToIO $ do
+                      arr <- A.unsafeFreeze dest
+                      return $! textP arr 0 (fromIntegral n)
+                  lastPtr <- peek curPtrPtr
+                  let left = lastPtr `minusPtr` curPtr
+                  return $ Other chunkText (B.drop left bs)
+                _ -> do
+                  -- We encountered the end of the buffer while decoding
+                  n <- peek destOffPtr
+                  codepoint <- peek codepointPtr
+                  chunkText <- unsafeSTToIO $ do
+                      arr <- A.unsafeFreeze dest
+                      return $! textP arr 0 (fromIntegral n)
+                  lastPtr <- peek curPtrPtr
+                  let left = lastPtr `minusPtr` curPtr
+                  return $ Some chunkText (B.drop left bs)
+                           (decodeChunk codepoint state)
+        in loop (ptr `plusPtr` off)
+  desc = "Data.Text.Encoding.streamDecodeUtf8With: Invalid UTF-8 stream"
+
+ord :: Char -> Int
+ord (C# c#) = I# (ord# c#)
+{-# INLINE ord #-}
+
+
+unsafeWrite :: A.MArray s -> Int -> Char -> ST s Int
+unsafeWrite marr i c
+    | n < 0x10000 = do 
+        A.unsafeWrite marr i (fromIntegral n)
+        return 1
+    | otherwise = do
+        A.unsafeWrite marr i lo
+        A.unsafeWrite marr (i+1) hi
+        return 2
+    where n = ord c
+          m = n - 0x10000
+          lo = fromIntegral $ (m `shiftR` 10) + 0xD800
+          hi = fromIntegral $ (m .&. 0x3FF) + 0xDC00
+          shiftR (I# x#) (I# i#) = I# (x# `iShiftRA#` i#)
+{-# INLINE unsafeWrite #-}
+
+foreign import ccall unsafe "_hs_pipes_text_decode_utf8_state" c_decode_utf8_with_state
+    :: MutableByteArray# s -> Ptr CSize
+    -> Ptr (Ptr Word8) -> Ptr Word8
+    -> Ptr CodePoint -> Ptr DecoderState -> IO (Ptr Word8)
\ No newline at end of file