module MultiMVar where

import Control.Concurrent

-- An abstraction around MVar that supports multiple nondeterministic take.
--   ToDo: we need Ord on MVar to do this properly.
--   ToDo: exception safety.

newtype MultiMVar a = MultiMVar (MVar (Either (a,[MVar a]) [Reader a]))
	-- Left indicates MultiMVar is full
	-- Right indicates it is empty

data Reader a = 
  Reader 
   (MVar ())	-- if this is empty, this reader is already satisfied
   (MVar a)	-- the box to put the result in for this reader

newEmptyMultiMVar :: IO (MultiMVar a)
newEmptyMultiMVar = newMVar (Right []) >>= return . MultiMVar

putMultiMVar :: MultiMVar a -> a -> IO ()
putMultiMVar (MultiMVar mm) a = do
  r <- takeMVar mm
  case r of
     Left (a,writers) -> do -- already full, so we need to wait
	wait <- newMVar a
	putMVar mm (Left (a,wait:writers))
	putMVar wait undefined -- this blocks until the put is performed
     Right rdrs -> loop rdrs
      where loop [] = putMVar mm (Left (a, []))
	    loop (Reader ctrl put : rdrs) = do
		k <- tryTakeMVar ctrl
		case k of
		   Nothing -> loop rdrs -- somebody got here first
		   Just _  -> do putMVar put a
				 putMVar mm (Right rdrs)

multiTake :: [MultiMVar a] -> IO a
multiTake mms = loop mms []
  -- we need to take all the MultiMVars, in order (to avoid deadlock)
  -- if we find one that is full, then return that, and unlock all the others
 where
  loop [] taken = do
    ctrl <- newMVar ()
    box  <- newEmptyMVar 
    sequence_ [ putMVar m (Right (Reader ctrl box : rdrs)) | (m,rdrs) <- taken ]
    takeMVar box
  loop (MultiMVar mm : mms) taken = do
    r <- takeMVar mm
    case r of
	Left (a,writers) -> do
	    -- got a full one, so unlock all the others
	    sequence_ [ putMVar m (Right rdrs) | (m,rdrs) <- taken ]
	    case writers of
		[] -> do
		   putMVar mm (Right [])
		   return a
		(wr:wrs) -> do
		   new_a <- takeMVar wr
		   putMVar mm (Left (new_a,wrs))
		   return a
	Right rdrs ->
	    loop mms ((mm,rdrs):taken)