From 6e245576e8e99bad922e4316b91c3365630e5f76 Mon Sep 17 00:00:00 2001
From: Edsko de Vries <edsko@well-typed.com>
Date: Wed, 29 Mar 2023 13:14:11 +0200
Subject: [PATCH] Improve external buffer API

---
 cbits/wrap-rust-haskell-ffi.c         | 10 ++++++
 foreign-rust.cabal                    |  2 ++
 src/Foreign/Rust/Marshall/External.hs | 51 ++++++++++++++++++---------
 3 files changed, 47 insertions(+), 16 deletions(-)
 create mode 100644 cbits/wrap-rust-haskell-ffi.c

diff --git a/cbits/wrap-rust-haskell-ffi.c b/cbits/wrap-rust-haskell-ffi.c
new file mode 100644
index 0000000..b03bb25
--- /dev/null
+++ b/cbits/wrap-rust-haskell-ffi.c
@@ -0,0 +1,10 @@
+// Forward-declare the Rust-exported function
+void haskell_ffi_external_free(void* vec);
+
+// Wrapper around the Rust function that takes an additional (unused) argument,
+// which makes it match the Haskell `FinalizerEnvPtr` type. The wrapper also
+// avoids linker errors when the Rust library is not available (of course,
+// the Rust library must be linked into the final application).
+void haskell_ffi_external_free_env(void* vec, void* ptr) {
+    haskell_ffi_external_free(vec);
+}
\ No newline at end of file
diff --git a/foreign-rust.cabal b/foreign-rust.cabal
index ff91c58..f0e3e7d 100644
--- a/foreign-rust.cabal
+++ b/foreign-rust.cabal
@@ -93,6 +93,8 @@ library
     , th-abstraction
     , vector
     , wide-word
+  c-sources:
+    cbits/wrap-rust-haskell-ffi.c
 
 test-suite test-foreign-rust
   import:
diff --git a/src/Foreign/Rust/Marshall/External.hs b/src/Foreign/Rust/Marshall/External.hs
index 697d832..b73ed35 100644
--- a/src/Foreign/Rust/Marshall/External.hs
+++ b/src/Foreign/Rust/Marshall/External.hs
@@ -6,48 +6,67 @@ module Foreign.Rust.Marshall.External (
   ) where
 
 import Codec.Borsh
-import Control.Exception
 import Data.Typeable
 import Data.Word
 import Foreign.C
-import Foreign.Concurrent
 import Foreign.Ptr
 
 import qualified Data.ByteString.Internal as Strict
 
 import Foreign.Rust.Marshall.Util
+import Foreign.ForeignPtr
 
-data External
+data ExternalBuffer
 
-foreign import ccall "haskell_ffi_external_ptr"
+{-------------------------------------------------------------------------------
+  Foreign imports
+
+  Although 'externalPtr' and 'externalLen' are morally pure, we make
+  them live in IO to make reasoning about order of operations easier in
+  'fromExternalBorsh'.
+
+  These C functions are defined in the companion Rust @haskell-ffi@ library.
+-------------------------------------------------------------------------------}
+
+foreign import ccall unsafe "haskell_ffi_external_ptr"
      externalPtr
-  :: Ptr External -> Ptr Word8
+  :: Ptr ExternalBuffer -> IO (Ptr Word8)
 
-foreign import ccall "haskell_ffi_external_len"
+foreign import ccall unsafe "haskell_ffi_external_len"
      externalLen
-  :: Ptr External -> CSize
+  :: Ptr ExternalBuffer -> IO CSize
 
-foreign import ccall "haskell_ffi_external_free"
+foreign import ccall unsafe "&haskell_ffi_external_free_env"
      externalFree
-  :: Ptr External -> IO ()
+  :: FinalizerEnvPtr ExternalBuffer Word8
 
--- | Internal auxiliary: cast pointer
+{-------------------------------------------------------------------------------
+  Internal auxiliary
+-------------------------------------------------------------------------------}
+
+-- | Cast pointer
 --
 -- For ease of integration with c2hs, 'fromExternalBorsh' takes a @Ptr ()@ as
--- input instead of the more accurate @Ptr External@.
-castToExternal :: Ptr () -> Ptr External
+-- input instead of the more accurate @Ptr ExternalBuffer@.
+castToExternal :: Ptr () -> Ptr ExternalBuffer
 castToExternal = castPtr
 
+{-------------------------------------------------------------------------------
+  Public API
+-------------------------------------------------------------------------------}
+
 -- | Output marshaller for values stored in Rust-allocated buffer
 --
 -- Should be used together with the Rust function @marshall_to_haskell_external@
 -- (from @haskell-ffi@).
 fromExternalBorsh :: (FromBorsh a, Typeable a) => Ptr () -> IO a
-fromExternalBorsh (castToExternal -> ptr) = do
-    len  <- evaluate $ fromIntegral $ externalLen ptr
-    fptr <- newForeignPtr (externalPtr ptr) (externalFree ptr)
+fromExternalBorsh (castToExternal -> vec) = do
+    ptr  <- externalPtr vec
+    len  <- fromIntegral <$> externalLen vec
+    fptr <- newForeignPtrEnv externalFree vec ptr
 
     let bs :: Strict.ByteString
-        bs = Strict.PS fptr 0 len
+        bs = Strict.fromForeignPtr fptr 0 len
 
     return $ deserialiseStrictOrPanic bs
+