allocate arguments on the haskell heap

1 files changed, 22 insertions, 23 deletions

diff --git a/wayland/Wayland/Server.hs b/wayland/Wayland/Server.hs
index 77dea63..89a9780 100644
--- a/wayland/Wayland/Server.hs
+++ b/wayland/Wayland/Server.hs
@@ -43,11 +43,11 @@ import Lens.Micro.TH (makeLenses)
 import Lens.Micro.Mtl ((%=), use, (<<%=))
 import qualified Wayland.Server.System as Sys
 import Control.Monad.Trans (MonadTrans(..))
-import Foreign.C (peekCString, newCString)
+import Foreign.C (peekCString, newCString, withCString)
 import Foreign.C.Types (CInt)
-import Foreign (Ptr, nullPtr, new, newStablePtr, castStablePtrToPtr, castPtrToStablePtr, deRefStablePtr, freeStablePtr, Storable (..), Int32, castPtr, free, mallocArray0)
+import Foreign (Ptr, nullPtr, new, newStablePtr, castStablePtrToPtr, castPtrToStablePtr, deRefStablePtr, freeStablePtr, Storable(..), Int32, castPtr, withArray, with, allocaArray0)
 import Data.IORef (IORef, newIORef, writeIORef, readIORef)
-import Control.Monad.Reader (ReaderT, MonadReader (..), runReaderT)
+import Control.Monad.Reader (ReaderT, MonadReader(..), runReaderT)
 import Control.Exception (bracket)
 import Data.Functor (($>))
 import Control.Monad (when)
@@ -58,6 +58,7 @@ import Data.Word (Word32)
 import System.Mem.StableName (StableName, makeStableName)
 import Data.Hashable (Hashable(..))
 import Foreign.Ptr (plusPtr)
+import Control.Monad.Cont (ContT(..))
 
 data Global = Global
 
@@ -342,16 +343,15 @@ getInterface p = do
 data ArgumentIteratorState
   = ArgumentIteratorState { _argNext          :: !(Ptr Sys.WLArgument)
                           , _argRemaining     :: !Int
-                          , _allocatedBuffers :: ![Ptr ()]
                           }
   deriving (Show)
 
 makeLenses ''ArgumentIteratorState
 
-newtype ArgumentIterator a = ArgumentIterator { runArgumentIterator :: StateT ArgumentIteratorState IO a }
+newtype ArgumentIterator r a = ArgumentIterator { runArgumentIterator :: StateT ArgumentIteratorState (ContT r IO) a }
   deriving (Functor, Applicative, Monad, MonadFail, MonadState ArgumentIteratorState, MonadIO)
 
-advanceIterator :: ArgumentIterator (Ptr Sys.WLArgument)
+advanceIterator :: ArgumentIterator r (Ptr Sys.WLArgument)
 advanceIterator = do
   next      <- argNext      <<%= \p -> plusPtr p (sizeOf (undefined :: Sys.WLArgument))
   remaining <- argRemaining <<%= (+(-1))
@@ -364,15 +364,17 @@ getObjectForHandle r = fetch (objectID r)
         check (Sys.WLResource ptr) | ptr == nullPtr = fail "Invalid object handle."
                                    | otherwise      = pure (castPtr ptr)
 
-instance ArgumentMarshaller ArgumentIterator where
+liftCont :: ((a -> IO r) -> IO r) -> ArgumentIterator r a
+liftCont = ArgumentIterator . lift . ContT
+
+instance ArgumentMarshaller (ArgumentIterator r) where
   marshalInt    v = advanceIterator >>= (liftIO . Sys.wlArgumentSetInt   (fromIntegral v))
   marshalUInt   v = advanceIterator >>= (liftIO . Sys.wlArgumentSetUInt  (fromIntegral v))
   marshalFixed  v = advanceIterator >>= (liftIO . Sys.wlArgumentSetFixed (Sys.WLFixed (round (v * 256))))
   marshalFD     v = advanceIterator >>= (liftIO . Sys.wlArgumentSetFD    (fromIntegral v))
   marshalString v = do
     arg <- advanceIterator
-    str <- liftIO $ newCString v
-    allocatedBuffers %= (castPtr str:)
+    str <- liftCont $ withCString v
     liftIO $ Sys.wlArgumentSetString str arg
   marshalObject v = do
     arg <- advanceIterator
@@ -385,26 +387,23 @@ instance ArgumentMarshaller ArgumentIterator where
   marshalArray v = do
     arg <- advanceIterator
     let len = fromIntegral $ length v * sizeOf (head v)
-    buf <- liftIO $ newArray v
-    arr <- liftIO $ new $ Sys.WLArray { Sys.arrSize  = len
-                                      , Sys.arrAlloc = len
-                                      , Sys.arrData  = castPtr buf
-                                      }
-    allocatedBuffers %= \p -> castPtr buf:castPtr arr:p
+    buf <- liftCont $ withArray v
+    arr <- liftCont $ with $ Sys.WLArray { Sys.arrSize  = len
+                                         , Sys.arrAlloc = len
+                                         , Sys.arrData  = castPtr buf
+                                         }
     liftIO $ Sys.wlArgumentSetArray arr arg
 
 sendEvent :: forall r s. ResourceClass r => r -> ResourceEvent r -> DisplayM s ()
 sendEvent res e = liftIO $ do
   let params = msgSignature (eventSignature (Proxy :: Proxy r) e)
   ptr <- getObjectForHandle res
-  buffer <- mallocArray0 (length params)
-  let s = ArgumentIteratorState { _argNext          = buffer
-                                , _argRemaining     = length params
-                                , _allocatedBuffers = [castPtr buffer]
-                                }
-  (opcode, s') <- runStateT (runArgumentIterator (marshalEvent (Proxy :: Proxy r) e)) s
-  Sys.wlResourcePostEventArray (Sys.WLResource ptr) (fromIntegral opcode) buffer
-  mapM_ free (_allocatedBuffers s')
+  allocaArray0 (length params) $ \buffer ->
+    let s = ArgumentIteratorState { _argNext          = buffer
+                                  , _argRemaining     = length params
+                                  }
+        handler (opcode, _) = Sys.wlResourcePostEventArray (Sys.WLResource ptr) (fromIntegral opcode) buffer
+    in runContT (runStateT (runArgumentIterator (marshalEvent (Proxy :: Proxy r) e)) s) handler
 
 flushEvents :: DisplayM s ()
 flushEvents = do