Tweak iteratee/enumerator map types

cvb · Dec 27, 2010 · 1d1c0a1 · 1d1c0a1
1 parent 8ba8f06
commit 1d1c0a1
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 25 deletions.
diff --git a/src/Snap/Internal/Http/Types.hs b/src/Snap/Internal/Http/Types.hs
@@ -372,9 +372,10 @@ rspBodyMap f b      = Enum $ f $ rspBodyToEnum b
 ------------------------------------------------------------------------------
 rspBodyToEnum :: ResponseBody -> Enumerator Builder IO a
 rspBodyToEnum (Enum e) = e
-rspBodyToEnum (SendFile fp Nothing)  = I.mapEnum fromByteString $ I.enumFile fp
+rspBodyToEnum (SendFile fp Nothing) =
-rspBodyToEnum (SendFile fp (Just s)) = I.mapEnum fromByteString $
+    I.mapEnum toByteString fromByteString $ I.enumFile fp
-                                       I.enumFilePartial fp s
+rspBodyToEnum (SendFile fp (Just s)) =
+    I.mapEnum toByteString fromByteString $ I.enumFilePartial fp s
 
 
 ------------------------------------------------------------------------------

diff --git a/src/Snap/Internal/Types.hs b/src/Snap/Internal/Types.hs
@@ -251,7 +251,7 @@ transformRequestBody trans = do
     let ioref = rqBody req
     senum <- liftIO $ readIORef ioref
     let (SomeEnumerator enum') = senum
-    let enum = mapEnum fromByteString enum'
+    let enum = mapEnum toByteString fromByteString enum'
     liftIO $ writeIORef ioref (SomeEnumerator enumEOF)
 
     origRsp <- getResponse

diff --git a/src/Snap/Iteratee.hs b/src/Snap/Iteratee.hs
@@ -640,30 +640,56 @@ enumFilePartial fp rng@(start,end) st@(Continue k) = do
 #endif
 
 
+------------------------------------------------------------------------------
+mapEnum :: (Monad m) =>
+           (aOut -> aIn)
+        -> (aIn -> aOut)
+        -> Enumerator aIn m a
+        -> Enumerator aOut m a
+mapEnum f g enum outStep = do
+    let z = I.map g outStep
+    let p = joinI z
+    let q = enum $$ p
+    (I.joinI . I.map f) $$ q
+
+
 ------------------------------------------------------------------------------
 mapIter :: (Monad m) =>
            (aOut -> aIn)
-        -> Iteratee aOut m a
+        -> (aIn -> aOut)
         -> Iteratee aIn m a
-mapIter f iter = iter >>== check
+        -> Iteratee aOut m a
+mapIter f g iter = do
+    step <- lift $ runIteratee iter
+    mapStep step
   where
-    check (Continue k) = k EOF >>== \s -> case s of
+    -- mapStep :: Step aIn m a -> Iteratee aOut m a
-        Continue _ -> error "divergent iteratee"
+    mapStep (Continue k)   = continue $ wrapK k
-        _ -> check s
+    mapStep (Yield x rest) = yield x (fmap g rest)
-    check (Yield x rest) = yield x (fmap f rest)
+    mapStep (Error e)      = throwError e
-    check (Error e) = throwError e
+
+    -- wrapK :: (Stream aIn  -> Iteratee aIn m a)
+    --       -> (Stream aOut -> Iteratee aOut m a)
+    wrapK k streamOut = mapIter f g iterIn
+      where
+        streamIn = fmap f streamOut
+        iterIn   = k streamIn
 
 
 ------------------------------------------------------------------------------
-mapEnum :: (Monad m) =>
+-- testIt :: IO ()
-           (aOut -> aIn)
+-- testIt = do
-        -> Enumerator aOut m a
+--     xs <- run_ (eInt $$ consume)
-        -> Enumerator aIn m a
+--     putStrLn (show xs)
-mapEnum f enum builderStep = do
+
-    -- z :: Iteratee ByteString m (Step Builder m a)
+--     ys <- run_ (eInt $$ replicateM 3 iter1)
-    let z = I.map f builderStep
+--     putStrLn (show ys)
-    -- p :: Iteratee ByteString m a
+
-    let p = joinI z
+--   where
-    -- q :: Iteratee ByteString m a
+--     eBS :: Enumerator ByteString IO a
-    let q = enum $$ p
+--     eBS = enumList 1 ["12345", "300", "200", "400"]
-    mapIter f q
+
+--     eInt :: Enumerator Int IO a
+--     eInt = mapEnum (S.pack . show) (read . S.unpack) eBS
+
+--     iter1 = mapIter (S.pack . show) (read . S.unpack) I.head
diff --git a/src/Snap/Util/GZip.hs b/src/Snap/Util/GZip.hs
@@ -180,9 +180,9 @@ compressEnumerator compFunc enum' origStep = do
     readEnd  <- liftIO $ newChan
     tid      <- liftIO $ forkIO $ threadProc readEnd writeEnd
 
-    let enum = mapEnum toByteString enum'
+    let enum = mapEnum fromByteString toByteString enum'
     let outEnum = enum (f readEnd writeEnd tid step)
-    mapIter fromByteString outEnum
+    mapIter toByteString fromByteString outEnum
 
   where
     --------------------------------------------------------------------------