import Control.Concurrent.Bag.Safe
import Control.Concurrent.STM (TChan)
import Control.Concurrent.STM.TStack
import Control.Concurrent (threadDelay)

data Tree a = Leaf a | Branch (Tree a) (Tree a)

genTree 0 i = Leaf i
genTree n i = Branch (genTree (n-1) (i*2)) (genTree (n-1) (i*2+1))

dfs :: Tree a -> Task TStack a (Maybe a)
dfs (Leaf a)     = return $ Just a
dfs (Branch l r) = do
  addTask $ dfs r
  dfs l

bfs :: Tree a -> Task TChan a (Maybe a)
bfs (Leaf a)     = return $ Just a
bfs (Branch l r) = do
  addTask $ bfs l
  addTask $ bfs r
  return Nothing

main = do
  putStrLn "First bag."
  newTaskBag (Just splitVertical) [dfs $ genTree 20 1] (readResults 0) >>= putStrLn . show
  putStrLn "Second bag."
  newTaskBag (Just splitHalf) [bfs $ genTree 19 1] (readResults 0) >>= putStrLn . show
  putStrLn "Third bag."
  newTaskBag Nothing [dfs $ genTree 21 1] (getResult >>= (liftIO . putStrLn . show))
  putStrLn "Waiting."
  threadDelay 10000000
 where
  readResults :: TaskBufferSTM b => Int -> BagT b Integer IO Int
  readResults n = do
    result <- getResult
    case result of
      Nothing -> return n
      Just v  ->
        readResults (n+1)