Fixed gatherElement bug

1 files changed, 37 insertions, 1 deletions

diff --git a/Presence/XMPP.hs b/Presence/XMPP.hs
index f81af7c0..bfa5827e 100644
--- a/Presence/XMPP.hs
+++ b/Presence/XMPP.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE OverloadedStrings   #-}
 {-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE CPP                 #-}
 module XMPP 
     ( module XMPPTypes
     , listenForXmppClients
@@ -44,6 +45,9 @@ import GetHostByAddr
 import Data.Monoid
 import qualified Data.Sequence as Seq
 import Data.Foldable (toList)
+#ifdef RENDERFLUSH
+import Data.Conduit.Blaze
+#endif
 
 data Commands = Send [XML.Event] | QuitThread
  deriving Prelude.Show
@@ -134,7 +138,7 @@ gatherElement opentag empty = gatherElement' (empty `mplus` return opentag) 1
                      _ | eventIsEndElement tag -> cnt-1
                      _ | eventIsBeginElement tag -> cnt+1
                      _ -> cnt
-        if (cnt>0) then gatherElement' ts' cnt'
+        if (cnt'>0) then gatherElement' ts' cnt'
                    else return ts' 
 
 voidMaybeT body = (>> return ()) . runMaybeT $ body
@@ -164,6 +168,11 @@ fromClient session cmdChan = voidMaybeT $ do
         tag@(EventBeginElement _ _) -> do
              xs <- gatherElement tag Seq.empty
              prettyPrint "client-in: ignoring..." (toList xs)
+             {-
+             liftIO (putStrLn "client-in: ignoring\n{")
+             liftIO (mapM_ print xs)
+             liftIO (putStrLn "}")
+             -}
              loop
         _ -> loop
 
@@ -202,7 +211,16 @@ handleClient st src snk = do
     pchan <- subscribe session Nothing
     cmdChan <- atomically newTChan
 
+#ifdef RENDERFLUSH
+    writer <- async (   toClient pchan cmdChan 
+                    $$  flushList
+                    =$= renderBuilderFlush def 
+                    =$= builderToByteStringFlush 
+                    =$= discardFlush 
+                     =$ snk )
+#else
     writer <- async ( toClient pchan cmdChan $$ renderChunks =$ snk )
+#endif
     finally         ( src $= parseBytes def $$ fromClient session cmdChan )
         $ do
             atomically $ writeTChan cmdChan QuitThread
@@ -231,8 +249,26 @@ seekRemotePeers config chan = do
     -- TODO
     return ()
 
+#ifdef RENDERFLUSH
+flushList :: Monad m => ConduitM [a] (Flush a) m ()
+flushList = fixMaybeT $ \loop -> do
+    xs <- mawait
+    lift ( CL.sourceList xs $$ CL.mapM_ (yield . Chunk) )
+    lift ( yield Flush )
+    loop
+
+discardFlush :: Monad m => ConduitM (Flush a) a m ()
+discardFlush = fixMaybeT $ \loop -> do
+    x <- mawait
+    let unchunk (Chunk a) = a
+        ischunk (Chunk _) = True
+        ischunk _         = False
+    lift . when (ischunk x) $ yield (unchunk x)
+    loop
+#else
 renderChunks :: (MonadUnsafeIO m, MonadIO m) => ConduitM [Event] ByteString m ()
 renderChunks = fixMaybeT $  \loop -> do
     xs <- mawait
     lift . when (not . null $ xs) $ ( CL.sourceList xs $= renderBytes def $$ CL.mapM_ yield )
     loop
+#endif