Hack in support for Maybe ByteString

Author: vkleen

src/Language/Rust/Inline.hs

@@ -103,7 +103,7 @@ import Foreign.Marshal.Alloc (alloca, free)
 import Foreign.Marshal.Array (newArray, withArrayLen)
 import Foreign.Marshal.Unsafe (unsafeLocalState)
 import Foreign.Marshal.Utils (new, with)
-import Foreign.Ptr (FunPtr, Ptr, freeHaskellFunPtr)
+import Foreign.Ptr (FunPtr, Ptr, freeHaskellFunPtr, nullPtr)
 
 import Control.Monad (void)
 import Data.List (intercalate)
@@ -315,6 +315,7 @@ processQQ safety isPure (QQParse rustRet rustBody rustNamedArgs) = do
     (returnFfi, haskRet') <- do
         marshalForm <- ghcMarshallable haskRet
         let fptrRet haskRet' = [t|Ptr (Ptr $(pure haskRet'), FunPtr (Ptr $(pure haskRet') -> IO ())) -> IO ()|]
+        let bsRet = [t|Ptr (Ptr Word8, Word, FunPtr (Ptr Word8 -> Word -> IO ())) -> IO ()|]
         ret <- case marshalForm of
             BoxedDirect -> [t|IO $(pure haskRet)|]
             BoxedIndirect -> [t|Ptr $(pure haskRet) -> IO ()|]
@@ -323,7 +324,8 @@ processQQ safety isPure (QQParse rustRet rustBody rustNamedArgs) = do
                 | otherwise ->
                     let retTy = showTy haskRet
                      in fail ("Cannot put unlifted type ‘" ++ retTy ++ "’ in IO")
-            ByteString -> [t|Ptr (Ptr Word8, Word, FunPtr (Ptr Word8 -> Word -> IO ())) -> IO ()|]
+            ByteString -> bsRet
+            OptionalByteString -> bsRet
             ForeignPtr
                 | AppT _ haskRet' <- haskRet -> fptrRet haskRet'
                 | otherwise -> fail ("Cannot marshal " ++ showTy haskRet ++ " using the ForeignPtr calling convention")
@@ -356,6 +358,9 @@ processQQ safety isPure (QQParse rustRet rustBody rustNamedArgs) = do
                 ByteString -> do
                     rbsT <- [t|Ptr (Ptr Word8, Word)|]
                     pure (ByteString, rbsT)
+                OptionalByteString -> do
+                    rbsT <- [t|Ptr (Ptr Word8, Word)|]
+                    pure (OptionalByteString, rbsT)
                 ForeignPtr
                     | AppT _ haskArg' <- haskArg -> do
                         ptr <- [t|Ptr $(pure haskArg')|]
@@ -431,6 +436,26 @@ processQQ safety isPure (QQParse rustRet rustBody rustNamedArgs) = do
                                 else Just <$> newForeignPtr $(varE finalizer) $(varE ptr)
                         )
                     |]
+            | returnFfi == OptionalByteString = do
+                ret <- newName "ret"
+                ptr <- newName "ptr"
+                len <- newName "len"
+                finalizer <- newName "finalizer"
+                [e|
+                    alloca
+                        ( \($(varP ret)) -> do
+                            $(appsE (varE qqName : reverse (varE ret : acc)))
+                            ($(varP ptr), $(varP len), $(varP finalizer)) <- peek $(varE ret)
+                            if $(varE ptr) == nullPtr
+                                then pure Nothing
+                                else
+                                    Just
+                                        <$> ByteString.unsafePackCStringFinalizer
+                                            $(varE ptr)
+                                            (fromIntegral $(varE len))
+                                            ($(varE bsFree) $(varE finalizer) $(varE ptr) $(varE len))
+                        )
+                    |]
             | returnByValue returnFfi = appsE (varE qqName : reverse acc)
             | otherwise = do
                 ret <- newName "ret"
@@ -475,6 +500,7 @@ processQQ safety isPure (QQParse rustRet rustBody rustNamedArgs) = do
                                 Just $(varP fptr) ->
                                     withForeignPtr $(varE fptr) (\($(varP ptr)) -> $(goArgs (varE ptr : acc) args))
                             |]
+                    | marshalForm == OptionalByteString -> fail "Don't"
                     | passByValue marshalForm -> goArgs (varE argName : acc) args
                     | otherwise -> do
                         x <- newName "x"

src/Language/Rust/Inline/Context/ByteString.hs

@@ -29,13 +29,16 @@ import Data.Functor (void)
 import Data.Word (Word8)
 import Foreign.Ptr (Ptr)
 
+import Debug.Trace (traceM)
+
 bytestrings :: Q Context
 bytestrings =
     pure $ Context ([rule], [], [rustByteString, impl])
   where
     rule rty _
         | rty == void [ty| &[u8] |] = pure ([t|ByteString|], pure . pure $ void [ty| RustByteString |])
         | rty == void [ty| Vec<u8> |] = pure ([t|ByteString|], pure . pure $ void [ty| RustOwnedByteString |])
+        | rty == void [ty| Option<Vec<u8>> |] = pure ([t|Maybe ByteString|], pure . pure $ void [ty| RustOwnedByteString |])
     rule _ _ = mempty
 
     rustByteString =
@@ -70,4 +73,13 @@ bytestrings =
             , "    RustOwnedByteString(bytes.as_mut_ptr(), len, free)"
             , "  }"
             , "}"
+            , ""
+            , "impl MarshalInto<RustOwnedByteString> for Option<Vec<u8>> {"
+            , "  fn marshal(self) -> RustOwnedByteString {"
+            , "    extern fn panic(ptr: *mut u8, len: usize) {"
+            , "      panic!(\"Attempted to free a null ByteString\");"
+            , "    }"
+            , "    self.map(|bs| bs.marshal()).unwrap_or(RustOwnedByteString(std::ptr::null_mut(), 0, panic))"
+            , "  }"
+            , "}"
             ]

src/Language/Rust/Inline/Marshal.hs

@@ -39,7 +39,8 @@ data MarshalForm
   | ByteString
   | ForeignPtr
   | OptionalForeignPtr
-  deriving (Eq)
+  | OptionalByteString
+  deriving (Eq, Show)
 
 passByValue :: MarshalForm -> Bool
 passByValue = (`elem` [UnboxedDirect, BoxedDirect, ForeignPtr])
@@ -71,9 +72,10 @@ ghcMarshallable ty = do
      AppT con _ | con `elem` tyconsU -> pure UnboxedDirect
                 | con `elem` tyconsB -> pure BoxedDirect
                 | con == fptrCons    -> pure ForeignPtr
-     AppT mb (AppT fptr _)
-                | mb == maybeCons && fptr == fptrCons -> pure OptionalForeignPtr
-     _                               -> pure BoxedIndirect
+     AppT mb (AppT c _)
+                | mb == maybeCons && c == fptrCons -> pure OptionalForeignPtr
+     AppT mb c | mb == maybeCons && c == bytestring -> pure OptionalByteString
+     _ -> pure BoxedIndirect
   where
   qSimpleUnboxed = [ [t| Char#   |]
                    , [t| Int#    |]

tests/ByteStrings.hs

@@ -9,6 +9,7 @@ import Test.Hspec
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Unsafe as ByteString
+import Data.Maybe (fromJust)
 import Data.String
 
 extendContext basic
@@ -34,3 +35,10 @@ bytestringSpec = describe "ByteStrings" $ do
         ByteString.pack [0, 1, 2, 3]
             `shouldBe` rustBs
         ByteString.unsafeFinalize rustBs
+
+    it "can marshal optional ByteString return values" $ do
+        let noRustBs = [rust| Option<Vec<u8>> { None } |]
+        noRustBs `shouldBe` Nothing
+
+        let rustBs = [rust| Option<Vec<u8>> { Some(vec![0, 1, 2, 3]) } |]
+        fromJust rustBs `shouldBe` ByteString.pack [0, 1, 2, 3]