Add hGetBufSome, like hGetBuf but can return short reads

1 files changed, 83 insertions, 32 deletions

diff --git a/libraries/base/GHC/IO/Handle/Text.hs b/libraries/base/GHC/IO/Handle/Text.hs
index 9e12283679..39482bc3a3 100644
--- a/libraries/base/GHC/IO/Handle/Text.hs
+++ b/libraries/base/GHC/IO/Handle/Text.hs
@@ -22,7 +22,7 @@
 module GHC.IO.Handle.Text ( 
    hWaitForInput, hGetChar, hGetLine, hGetContents, hPutChar, hPutStr,
    commitBuffer',       -- hack, see below
-   hGetBuf, hGetBufNonBlocking, hPutBuf, hPutBufNonBlocking,
+   hGetBuf, hGetBufSome, hGetBufNonBlocking, hPutBuf, hPutBufNonBlocking,
    memcpy,
  ) where
 
@@ -63,18 +63,24 @@ import GHC.List
 -- | Computation 'hWaitForInput' @hdl t@
 -- waits until input is available on handle @hdl@.
 -- It returns 'True' as soon as input is available on @hdl@,
--- or 'False' if no input is available within @t@ milliseconds.
+-- or 'False' if no input is available within @t@ milliseconds.  Note that
+-- 'hWaitForInput' waits until one or more full /characters/ are available,
+-- which means that it needs to do decoding, and hence may fail
+-- with a decoding error.
 --
 -- If @t@ is less than zero, then @hWaitForInput@ waits indefinitely.
 --
 -- This operation may fail with:
 --
 --  * 'isEOFError' if the end of file has been reached.
+--  * a decoding error, if the input begins with an invalid byte sequence
+--    in this Handle's encoding.
 --
 -- NOTE for GHC users: unless you use the @-threaded@ flag,
 -- @hWaitForInput t@ where @t >= 0@ will block all other Haskell
 -- threads for the duration of the call.  It behaves like a
 -- @safe@ foreign call in this respect.
+--
 
 hWaitForInput :: Handle -> Int -> IO Bool
 hWaitForInput h msecs = do
@@ -790,9 +796,6 @@ writeChunkNonBlocking h_@Handle__{..} ptr bytes
 -- It returns the number of bytes actually read.  This may be zero if
 -- EOF was reached before any data was read (or if @count@ is zero).
 --
--- 'hGetBuf' ignores whatever 'TextEncoding' the 'Handle' is currently
--- using, and reads bytes directly from the underlying IO device.
---
 -- 'hGetBuf' never raises an EOF exception, instead it returns a value
 -- smaller than @count@.
 --
@@ -868,6 +871,51 @@ readChunk h_@Handle__{..} ptr bytes
         then return off
         else loop fd (off + r) (bytes - r)
 
+-- ---------------------------------------------------------------------------
+-- hGetBufSome
+
+-- | 'hGetBufSome' @hdl buf count@ reads data from the handle @hdl@
+-- into the buffer @buf@.  If there is any data available to read,
+-- then 'hGetBufSome' returns it immediately; it only blocks if there
+-- is no data to be read.
+--
+-- It returns the number of bytes actually read.  This may be zero if
+-- EOF was reached before any data was read (or if @count@ is zero).
+--
+-- 'hGetBufSome' never raises an EOF exception, instead it returns a value
+-- smaller than @count@.
+--
+-- If the handle is a pipe or socket, and the writing end
+-- is closed, 'hGetBufSome' will behave as if EOF was reached.
+--
+-- 'hGetBufSome' ignores the prevailing 'TextEncoding' and 'NewlineMode'
+-- on the 'Handle', and reads bytes directly.
+
+hGetBufSome :: Handle -> Ptr a -> Int -> IO Int
+hGetBufSome h ptr count
+  | count == 0 = return 0
+  | count <  0 = illegalBufferSize h "hGetBuf" count
+  | otherwise =
+      wantReadableHandle_ "hGetBuf" h $ \ h_@Handle__{..} -> do
+         flushCharReadBuffer h_
+         buf@Buffer{ bufSize=sz } <- readIORef haByteBuffer
+         if isEmptyBuffer buf
+            then if count > sz  -- large read?
+                    then do RawIO.read (haFD h_) (castPtr ptr) count
+                    else do (r,buf') <- Buffered.fillReadBuffer haDevice buf
+                            if r == 0
+                               then return 0
+                               else do writeIORef haByteBuffer buf'
+                                       bufReadNBNonEmpty h_ buf' (castPtr ptr) 0 count
+            else
+              bufReadNBEmpty h_ buf (castPtr ptr) 0 count
+
+haFD :: Handle__ -> FD
+haFD h_@Handle__{} = 
+   case cast h_ of
+             Nothing -> error "not an FD"
+             Just fd -> fd
+
 -- | 'hGetBufNonBlocking' @hdl buf count@ reads data from the handle @hdl@
 -- into the buffer @buf@ until either EOF is reached, or
 -- @count@ 8-bit bytes have been read, or there is no more data available
@@ -878,10 +926,6 @@ readChunk h_@Handle__{..} ptr bytes
 -- only whatever data is available.  To wait for data to arrive before
 -- calling 'hGetBufNonBlocking', use 'hWaitForInput'.
 --
--- 'hGetBufNonBlocking' ignores whatever 'TextEncoding' the 'Handle'
--- is currently using, and reads bytes directly from the underlying IO
--- device.
---
 -- If the handle is a pipe or socket, and the writing end
 -- is closed, 'hGetBufNonBlocking' will behave as if EOF was reached.
 --
@@ -902,21 +946,33 @@ bufReadNonBlocking h_@Handle__{..} ptr so_far count =
   seq so_far $ seq count $ do -- strictness hack
   buf@Buffer{ bufRaw=raw, bufR=w, bufL=r, bufSize=sz } <- readIORef haByteBuffer
   if isEmptyBuffer buf
-     then if count > sz  -- large read?
-                then do rest <- readChunkNonBlocking h_ ptr count
-                        return (so_far + rest)
-                else do (r,buf') <- Buffered.fillReadBuffer0 haDevice buf
-                        case r of
-                          Nothing -> return so_far
-                          Just 0  -> return so_far
-                          Just r  -> do
-                            writeIORef haByteBuffer buf'
-                            bufReadNonBlocking h_ ptr so_far (min count r)
-                                  -- NOTE: new count is    min count w'
-                                  -- so we will just copy the contents of the
-                                  -- buffer in the recursive call, and not
-                                  -- loop again.
-     else do
+     then bufReadNBEmpty    h_ buf ptr so_far count
+     else bufReadNBNonEmpty h_ buf ptr so_far count
+
+bufReadNBEmpty :: Handle__ -> Buffer Word8 -> Ptr Word8 -> Int -> Int -> IO Int
+bufReadNBEmpty   h_@Handle__{..}
+                 buf@Buffer{ bufRaw=raw, bufR=w, bufL=r, bufSize=sz }
+                 ptr so_far count
+   = if count > sz  -- large read?
+        then do rest <- readChunkNonBlocking h_ ptr count
+                return (so_far + rest)
+        else do (r,buf') <- Buffered.fillReadBuffer0 haDevice buf
+                case r of
+                  Nothing -> return so_far
+                  Just 0  -> return so_far
+                  Just r  -> do
+                    writeIORef haByteBuffer buf'
+                    bufReadNBNonEmpty h_ buf' ptr so_far (min count r)
+                          -- NOTE: new count is    min count w'
+                          -- so we will just copy the contents of the
+                          -- buffer in the recursive call, and not
+                          -- loop again.
+
+bufReadNBNonEmpty :: Handle__ -> Buffer Word8 -> Ptr Word8 -> Int -> Int -> IO Int
+bufReadNBNonEmpty h_@Handle__{..}
+                  buf@Buffer{ bufRaw=raw, bufR=w, bufL=r, bufSize=sz }
+                  ptr so_far count
+  = do
         let avail = w - r
         if (count == avail)
            then do 
@@ -932,18 +988,13 @@ bufReadNonBlocking h_@Handle__{..} ptr so_far count =
            else do
 
         copyFromRawBuffer ptr raw (fromIntegral r) (fromIntegral avail)
-        writeIORef haByteBuffer buf{ bufR=0, bufL=0 }
+        let buf' = buf{ bufR=0, bufL=0 }
+        writeIORef haByteBuffer buf'
         let remaining = count - avail
             so_far' = so_far + avail
             ptr' = ptr `plusPtr` avail
 
-        -- we haven't attempted to read anything yet if we get to here.
-        if remaining < sz
-           then bufReadNonBlocking h_ ptr' so_far' remaining
-           else do 
-
-        rest <- readChunkNonBlocking h_ ptr' remaining
-        return (so_far' + rest)
+        bufReadNBEmpty h_ buf' ptr' so_far' remaining
 
 
 readChunkNonBlocking :: Handle__ -> Ptr Word8 -> Int -> IO Int