Use ForeignPtr instead of allocaBytes for direct serialization

IntersectMBO · Apr 29, 2024 · e2d86bb · e2d86bb
1 parent d1e8df9
commit e2d86bb
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 16 deletions.
diff --git a/cardano-crypto-class/src/Cardano/Crypto/KES/CompactSum.hs b/cardano-crypto-class/src/Cardano/Crypto/KES/CompactSum.hs
@@ -87,6 +87,7 @@ module Cardano.Crypto.KES.CompactSum (
 import           Data.Proxy (Proxy(..))
 import           GHC.Generics (Generic)
 import qualified Data.ByteString as BS
+import qualified Data.ByteString.Internal as BS
 import           Control.Monad (guard, (<$!>))
 import           NoThunks.Class (NoThunks, OnlyCheckWhnfNamed (..))
 
@@ -510,7 +511,8 @@ instance (HashAlgorithm h)
   directDeserialise pull = do
     let len :: Num a => a
         len = fromIntegral $ sizeHash (Proxy @h)
-    allocaBytes len $ \ptr -> do
-      pull ptr len
-      bs <- packByteStringCStringLen (ptr, len)
-      maybe (error "Invalid hash") return $! VerKeyCompactSumKES <$!> hashFromBytes bs
+    fptr <- mallocForeignPtrBytes len
+    withForeignPtr fptr $ \ptr -> do
+      pull (castPtr ptr) len
+    let bs = BS.fromForeignPtr0 fptr len
+    maybe (error "Invalid hash") return $! VerKeyCompactSumKES <$!> hashFromBytes bs
diff --git a/cardano-crypto-class/src/Cardano/Crypto/KES/Mock.hs b/cardano-crypto-class/src/Cardano/Crypto/KES/Mock.hs
@@ -24,6 +24,8 @@ import Data.Proxy (Proxy(..))
 import GHC.Generics (Generic)
 import GHC.TypeNats (Nat, KnownNat, natVal)
 import NoThunks.Class (NoThunks)
+import qualified Data.ByteString.Internal as BS
+import Foreign.Ptr (castPtr)
 
 import Control.Exception (assert)
 
@@ -39,8 +41,8 @@ import Cardano.Crypto.Libsodium
   )
 import Cardano.Crypto.Libsodium.Memory
   ( unpackByteStringCStringLen
-  , packByteStringCStringLen
-  , allocaBytes
+  , mallocForeignPtrBytes
+  , withForeignPtr
   )
 import Cardano.Crypto.DirectSerialise
 
@@ -210,9 +212,10 @@ instance (KnownNat t) => DirectSerialise (SignKeyKES (MockKES t)) where
 instance (KnownNat t) => DirectDeserialise (SignKeyKES (MockKES t)) where
   directDeserialise pull = do
     let len = fromIntegral $ sizeSignKeyKES (Proxy @(MockKES t))
-    bs <- allocaBytes len $ \cstr -> do
-        pull cstr (fromIntegral len)
-        packByteStringCStringLen (cstr, len)
+    fptr <- mallocForeignPtrBytes len
+    withForeignPtr fptr $ \ptr ->
+        pull (castPtr ptr) (fromIntegral len)
+    let bs = BS.fromForeignPtr0 fptr len
     maybe (error "directDeserialise @(SignKeyKES (MockKES t))") return $
         rawDeserialiseSignKeyMockKES bs
 
@@ -224,8 +227,9 @@ instance (KnownNat t) => DirectSerialise (VerKeyKES (MockKES t)) where
 instance (KnownNat t) => DirectDeserialise (VerKeyKES (MockKES t)) where
   directDeserialise pull = do
     let len = fromIntegral $ sizeVerKeyKES (Proxy @(MockKES t))
-    bs <- allocaBytes len $ \cstr -> do
-        pull cstr (fromIntegral len)
-        packByteStringCStringLen (cstr, len)
+    fptr <- mallocForeignPtrBytes len
+    withForeignPtr fptr $ \ptr ->
+        pull (castPtr ptr) (fromIntegral len)
+    let bs = BS.fromForeignPtr0 fptr len
     maybe (error "directDeserialise @(VerKeyKES (MockKES t))") return $
         rawDeserialiseVerKeyKES bs
diff --git a/cardano-crypto-class/src/Cardano/Crypto/KES/Sum.hs b/cardano-crypto-class/src/Cardano/Crypto/KES/Sum.hs
@@ -54,6 +54,7 @@ module Cardano.Crypto.KES.Sum (
 import           Data.Proxy (Proxy(..))
 import           GHC.Generics (Generic)
 import qualified Data.ByteString as BS
+import qualified Data.ByteString.Internal as BS
 import           Control.Monad (guard, (<$!>))
 import           NoThunks.Class (NoThunks, OnlyCheckWhnfNamed (..))
 
@@ -430,7 +431,8 @@ instance (HashAlgorithm h)
   directDeserialise pull = do
     let len :: Num a => a
         len = fromIntegral $ sizeHash (Proxy @h)
-    allocaBytes len $ \ptr -> do
-      pull ptr len
-      bs <- packByteStringCStringLen (ptr, len)
-      maybe (error "Invalid hash") return $! VerKeySumKES <$!> hashFromBytes bs
+    fptr <- mallocForeignPtrBytes len
+    withForeignPtr fptr $ \ptr -> do
+      pull (castPtr ptr) len
+    let bs = BS.fromForeignPtr0 fptr len
+    maybe (error "Invalid hash") return $! VerKeySumKES <$!> hashFromBytes bs
diff --git a/cardano-crypto-class/src/Cardano/Crypto/Libsodium/Memory.hs b/cardano-crypto-class/src/Cardano/Crypto/Libsodium/Memory.hs
@@ -25,6 +25,10 @@ module Cardano.Crypto.Libsodium.Memory (
   copyMem,
   allocaBytes,
 
+  -- * 'ForeignPtr' operations, generalized to 'MonadST'
+  mallocForeignPtrBytes,
+  withForeignPtr,
+
   -- * ByteString memory access, generalized to 'MonadST'
   unpackByteStringCStringLen,
   packByteStringCStringLen,