Nesting support for StatusCache

1 files changed, 21 insertions, 5 deletions

diff --git a/Control/Concurrent/STM/StatusCache.hs b/Control/Concurrent/STM/StatusCache.hs
index 601de14c..e08be33c 100644
--- a/Control/Concurrent/STM/StatusCache.hs
+++ b/Control/Concurrent/STM/StatusCache.hs
@@ -76,6 +76,8 @@ data StatusCache x =
     StatusCache { feed :: TVar (TChan x)
                 , cache :: TVar (TChan x)
                 , feedFlag :: TVar Bool
+                , pushDepth :: TVar Int
+                , pullDepth :: TVar Int
                 , isStarter :: x -> Bool -- ^ True if the given chunk begins a message.
                 , isStopper :: x -> Bool -- ^ True if the given chunk ends a message.
                 }
@@ -89,9 +91,13 @@ new isStart isStop = do
     feed <- newTChan >>= newTVar 
     cache <- newTChan >>= newTVar
     flag <- newTVar True
+    pushd <- newTVar 0
+    pulld <- newTVar 0
     return StatusCache { feed = feed
                        , cache = cache
                        , feedFlag = flag
+                       , pushDepth = pushd
+                       , pullDepth = pulld
                        , isStarter = isStart
                        , isStopper = isStop
                        }
@@ -114,23 +120,33 @@ pull q = do
     if exhausted then retry
     else do
         v <- readTChan chan
-        when (isStarter q v)
-            $ writeTVar (feedFlag q) False
+        when (isStarter q v) $ do
+            depth <- readTVar (pullDepth q)
+            modifyTVar' (pullDepth q) (+1)
+            when (depth==0)
+                $ writeTVar (feedFlag q) False
+        when (isStopper q v) 
+            $ modifyTVar' (pullDepth q) (subtract 1)
         return v
 
 -- | Enqueue a chunk into the 'StatusCache'.
 push :: StatusCache a -> a -> STM ()
 push q v = do
     shouldCache <- readTVar (feedFlag q)
+    depth <- readTVar (pushDepth q)
+    when (isStopper q v) $ do
+        modifyTVar' (pushDepth q) (subtract 1)
+        when (depth==0)
+            $ writeTVar (feedFlag q) True
+    when (isStarter q v) 
+        $ modifyTVar' (pushDepth q) (+1)
     chan <-
         if shouldCache then do
-            when (isStarter q v)
+            when (depth==0 && isStarter q v)
                 $ newTChan 
                   >>= writeTVar (cache q) 
             readTVar $ cache q
         else do
-            when (isStopper q v)
-              $ writeTVar (feedFlag q) True
             readTVar $ feed q
     writeTChan chan v