Clean up buffer reads

Author: runarorama

parser-typechecker/src/Unison/Builtin.hs

@@ -855,9 +855,9 @@ ioBuiltins =
     ("IO.getSomeBytes.impl.v1", handle --> nat --> iof bytes),
     ("IO.putBytes.impl.v3", handle --> bytes --> iof unit),
     ("IO.getLine.impl.v1", handle --> iof text),
-    ("IO.fillBuf.impl.v1", forall1 "g" \g -> handle --> pinnedByteArrayt g --> iof nat),
+    ("IO.fillBuf.impl.v1", forall1 "g" \g -> handle --> pinnedByteArrayt g --> nat --> iof nat),
     ("IO.putBuf.impl.v1", forall1 "g" \g -> handle --> pinnedByteArrayt g --> nat --> iof nat),
-    ("IO.getBufSome.impl.v1", forall1 "g" \g -> handle --> pinnedByteArrayt g --> iof nat),
+    ("IO.getBufSome.impl.v1", forall1 "g" \g -> handle --> pinnedByteArrayt g --> nat --> iof nat),
     ("IO.systemTime.impl.v3", unit --> iof nat),
     ("IO.systemTimeMicroseconds.v1", unit --> io int),
     ("IO.getTempDirectory.impl.v3", unit --> iof text),

unison-runtime/src/Unison/Runtime/Foreign/Function.hs

@@ -207,8 +207,8 @@ import UnliftIO qualified
 
 withMutableByteArrayContents :: (PA.PrimBase m) => PA.MutableByteArray (PA.PrimState m) -> (Ptr Word8 -> m a) -> m a
 {-# INLINE withMutableByteArrayContents #-}
-withMutableByteArrayContents mba =
-  PA.keepAlive (PA.mutableByteArrayContents mba)
+withMutableByteArrayContents mba f =
+  PA.keepAlive mba (f . PA.mutableByteArrayContents)
 
 -- foreignCall is explicitly NOINLINE'd because it's a _huge_ chunk of code and negatively affects code caching.
 -- Because we're not inlining it, we need a wrapper using an explicitly unboxed Stack so we don't block the
@@ -286,13 +286,21 @@ foreignCallHelper = \case
     \(h, n) -> Bytes.fromArray <$> hGetSome h n
   IO_putBytes_impl_v3 -> mkForeignIOF $ \(h, bs) -> hPut h (Bytes.toArray bs)
   -- TODO: Use `PA.withMutableByteArrayContents` here once we have Data.Primitive v9.
-  IO_fillBuf_impl_v1 -> mkForeignIOF $ \(h, arr) -> let !sz = PA.sizeofMutableByteArray arr in withMutableByteArrayContents arr (\ptr -> hGetBuf h ptr sz)
+  IO_fillBuf_impl_v1 -> mkForeignIOF $ \(h, arr, n) -> do
+    r <- checkBoundsPrim "IO.fillBuf.impl.v1" (PA.sizeofMutableByteArray arr) n 0 . pure $ Right ()
+    case r of
+      Left (F.Failure _ err _) -> ioError (userError (Util.Text.unpack err))
+      Right _ -> withMutableByteArrayContents arr (\ptr -> hGetBuf h ptr (fromIntegral n))
   IO_putBuf_impl_v1 -> mkForeignIOF $ \(h, arr, n) -> do
     r <- checkBoundsPrim "IO.putBuf.impl.v1" (PA.sizeofMutableByteArray arr) n 0 . pure $ Right ()
     case r of
       Left (F.Failure _ err _) -> ioError (userError (Util.Text.unpack err))
-      Right _ -> hPutBuf h (PA.mutableByteArrayContents arr) (fromIntegral n)
-  IO_getBufSome_impl_v1 -> mkForeignIOF $ \(h, arr) -> let !sz = PA.sizeofMutableByteArray arr in withMutableByteArrayContents arr (\ptr -> hGetBufSome h ptr sz)
+      Right _ -> withMutableByteArrayContents arr (\ptr -> hPutBuf h ptr (fromIntegral n))
+  IO_getBufSome_impl_v1 -> mkForeignIOF $ \(h, arr, n) -> do
+    r <- checkBoundsPrim "IO.getBufSome.impl.v1" (PA.sizeofMutableByteArray arr) n 0 . pure $ Right ()
+    case r of
+      Left (F.Failure _ err _) -> ioError (userError (Util.Text.unpack err))
+      Right _ -> withMutableByteArrayContents arr (\ptr -> hGetBufSome h ptr (fromIntegral n))
   IO_systemTime_impl_v3 -> mkForeignIOF $
     \() -> getPOSIXTime
   IO_systemTimeMicroseconds_v1 -> mkForeign $

unison-src/transcripts-using-base/all-base-hashes.output.md

@@ -1725,6 +1725,7 @@ This transcript is intended to make visible accidental changes to the hashing al
   477.  -- ##IO.fillBuf.impl.v1
         builtin.io2.IO.fillBuf.impl : Handle
         -> PinnedByteArray g
+        -> Nat
         ->{IO} Either Failure Nat
         
   478.  -- ##IO.forkComp.v2
@@ -1741,6 +1742,7 @@ This transcript is intended to make visible accidental changes to the hashing al
   481.  -- ##IO.getBufSome.impl.v1
         builtin.io2.IO.getBufSome.impl : Handle
         -> PinnedByteArray g
+        -> Nat
         ->{IO} Either Failure Nat
         
   482.  -- ##IO.getBytes.impl.v3

unison-src/transcripts/idempotent/io.md

@@ -261,6 +261,80 @@ testGetSomeBytes _ =
   Tip: Use view 1 to view the source of a test.
 ```
 
+### Reading and writing via buffers
+
+Tests:
+
+  - fillBuf
+  - putBuf
+  - getBufSome
+
+``` unison
+testFillBuf : '{io2.IO} [Result]
+testFillBuf = do
+  test = 'let
+    tempDir = (newTempDir "getSomeBytes")
+    fooFile = tempDir ++ "/foo"
+
+    testData = "0123456789"
+    testSize = size testData
+
+    wbuf = IO.pinnedByteArray 10
+    arr = PinnedByteArray.cast wbuf
+    dataBytes = ImmutableByteArray.fromBytes (toUtf8 testData)
+    ImmutableByteArray.copyTo! arr 0 dataBytes 0 10
+
+    -- write testData to a temporary file
+    fooWrite = openFile fooFile Write
+    _ = putBuf.impl fooWrite wbuf 10
+    closeFile fooWrite
+    check "file should be closed" (not (isFileOpen fooWrite))
+
+    -- reopen for reading back the data in chunks
+    fooRead = openFile fooFile Read
+
+    rbuf = IO.pinnedByteArray 10
+    rarr = PinnedByteArray.cast rbuf
+    _ = fillBuf.impl fooRead rbuf 4
+    vec = ImmutableByteArray.toBytes (MutableByteArray.freeze rarr 0 4) 0 4
+    check "should be able to read 4 bytes" (vec == (toUtf8 testData |> take 4))
+
+    _ = getBufSome.impl fooRead rbuf 4
+    vec2 = ImmutableByteArray.toBytes (MutableByteArray.freeze rarr 0 4) 0 4
+    check "should be able to read next 4 bytes" (vec2 == (toUtf8 testData |> drop 4 |> take 4))
+
+  runTest test
+```
+
+``` ucm :added-by-ucm
+  Loading changes detected in scratch.u.
+
+  + testFillBuf : '{IO} [Result]
+
+  Run `update` to apply these changes to your codebase.
+```
+
+``` ucm
+> add
+
+  Okay, I'm searching the branch for code that needs to be
+  updated...
+
+  Done.
+
+> io.test testFillBuf
+
+    New test results:
+
+    1. testFillBuf   ◉ file should be closed
+                     ◉ should be able to read 4 bytes
+                     ◉ should be able to read next 4 bytes
+
+  ✅ 3 test(s) passing
+
+  Tip: Use view 1 to view the source of a test.
+```
+
 ### Seeking in open files
 
 Tests: