Wip/davean/safe st

src/Data/Machine/Zlib.hs

 {-# LANGUAGE ScopedTypeVariables #-}
 module Data.Machine.Zlib
- ( gzipper
- , gunzipper
+ ( gunzipper
  , WindowBits(..)
  ) where
 
-import Control.Applicative
-import Control.Exception (throw)
-import Control.Monad
+import           Control.Applicative
+import           Control.Monad
+import           Control.Monad.Catch
+import           Control.Monad.Primitive
+import           Control.Monad.Trans
 import qualified Data.ByteString as BS
-import Data.Machine
-import Data.Machine.Stack
-import Data.Streaming.Zlib
-import System.IO.Unsafe
+import           Data.Machine
+import           Data.Machine.Stack
+import           Data.Streaming.Zlib
 
-unsafeM :: Monad n => IO a -> n a
-unsafeM ma = pure . unsafePerformIO $! ma
-
-handleRes :: Applicative m => PopperRes -> m (Maybe BS.ByteString)
-handleRes PRDone      = pure $ Nothing
-handleRes (PRNext bs) = pure $ Just bs
-handleRes (PRError e) = throw e
-
-gunzipper :: forall m . (Monad m) => MachineT m (Stack BS.ByteString) BS.ByteString
-gunzipper = construct $
-    go =<< (unsafeM $ initInflate (WindowBits 0)) -- Use what the header says.
+gunzipper :: forall m . (MonadThrow m, PrimMonad m, Monad m) => MachineT m (Stack BS.ByteString) BS.ByteString
+gunzipper = construct $ do
+  i <- doPrim $ initInflate (WindowBits 0)
+  start i
   where
-    go :: Inflate -> PlanT (Stack BS.ByteString) BS.ByteString m ()
-    go inflate = do
-        bs <- awaits Pop <|> stop
-        popper <- unsafeM <$> unsafeM (feedInflate inflate bs)
-        exhaust (handleRes =<< popper)
-        complete <- unsafeM $ isCompleteInflate inflate
-        if complete
-        then void $ finish
-        else go'
-       where
-        go' :: PlanT (Stack BS.ByteString) BS.ByteString m ()
-        go' = do
-          bs <- awaits Pop <|> finish
-          popper <- unsafeM <$> unsafeM (feedInflate inflate bs)
-          exhaust (handleRes =<< popper)
-          complete <- unsafeM $ isCompleteInflate inflate
-          if complete
-          then void $ finish
-          else go'
-        finish :: PlanT (Stack BS.ByteString) BS.ByteString m BS.ByteString
-        finish = do
-          r <- unsafeM $ finishInflate inflate
-          yield r
-          unused <- unsafeM $ getUnusedInflate inflate
-          awaits (Push unused)
-          stop
-
-gzipper :: forall m. Monad m => Int -> WindowBits -> ProcessT m BS.ByteString BS.ByteString
-gzipper lvl wb = construct $
-    go =<< (unsafeM $ initDeflate lvl wb)
-  where
-    go :: Deflate -> PlanT (Is BS.ByteString) BS.ByteString m ()
-    go deflate = go'
+    doPrim :: IO a -> PlanT (Stack BS.ByteString) BS.ByteString m a
+    doPrim = lift . unsafePrimToPrim
+    yield' :: BS.ByteString -> PlanT (Stack BS.ByteString) BS.ByteString m ()
+    yield' d = unless (BS.null d) $ yield d
+    push' :: BS.ByteString -> PlanT (Stack BS.ByteString) BS.ByteString m ()
+    push' d = unless (BS.null d) $ push d
+    -- This is all stupid. Its based on the exact behavior of the base API.
+    start :: Inflate -> PlanT (Stack BS.ByteString) BS.ByteString m ()
+    start i = go
       where
-        go' :: PlanT (Is BS.ByteString) BS.ByteString m ()
-        go' = do
-          bs <- await <|> finish
-          popper <- unsafeM <$> unsafeM (feedDeflate deflate bs)
-          exhaust (handleRes =<< popper)
-          go'
-        finish = do
-          let popper = unsafeM (finishDeflate deflate)
-          exhaust (handleRes =<< popper)
-          stop
+        drain :: Popper -> PlanT (Stack BS.ByteString) BS.ByteString m ()
+        drain popper =  do
+          md <- doPrim popper
+          case md of
+            PRNext d -> yield' d >> drain popper
+            PRDone -> do
+              -- This magic tries to detect if we've reached the end of a substream.
+              l <- doPrim (getUnusedInflate i)
+              if BS.null l
+                then go
+                else do
+                  yield' =<< doPrim (flushInflate i)
+                  push l
+            PRError e -> lift $ throwM e
+        getMore :: PlanT (Stack BS.ByteString) BS.ByteString m ()
+        getMore = drain =<< (doPrim . feedInflate i) =<< pop 
+        allDone :: PlanT (Stack BS.ByteString) BS.ByteString m ()
+        allDone = do
+          yield' =<< doPrim (flushInflate i)
+          push' =<< doPrim (getUnusedInflate i)
+        go :: PlanT (Stack BS.ByteString) BS.ByteString m ()
+        go = getMore <|> allDone
+