-- Communicating Haskell Processes.
-- Copyright (c) 2008, University of Kent.
-- All rights reserved.
-- 
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are
-- met:
--
--  * Redistributions of source code must retain the above copyright
--    notice, this list of conditions and the following disclaimer.
--  * Redistributions in binary form must reproduce the above copyright
--    notice, this list of conditions and the following disclaimer in the
--    documentation and/or other materials provided with the distribution.
--  * Neither the name of the University of Kent nor the names of its
--    contributors may be used to endorse or promote products derived from
--    this software without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
-- IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
-- THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
-- PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
-- CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
-- EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
-- PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
-- PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
-- LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
-- NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-- SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

--TODO document this (for internal purposes)
module Control.Concurrent.CHP.Event (RecordedEventType(..), Event, getEventUnique,
  SignalVar, SignalValue(..), enableEvents, disableEvents,
  newEvent, newEventUnique, enrollEvent, resignEvent, poisonEvent, checkEventForPoison,
  getEventTypeVal
#ifdef CHP_TEST
  , testAll
#endif
  ) where

import Control.Arrow
import Control.Concurrent.STM
import Control.Monad
import Data.Function
import Data.List
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Set as Set
import qualified Data.Traversable as T
import Data.Unique
import Prelude hiding (seq)
#ifdef CHP_TEST
import Test.HUnit hiding (test)
#endif

import Control.Concurrent.CHP.Poison
import Control.Concurrent.CHP.ProcessId

-- | The type of an event in the CSP and VCR traces.
--
-- ClockSync was added in version 1.2.0.
--
-- The extra parameter on ChannelComm and BarrierSync (which are the result of
-- showing the value sent and phase ended respectively) was added in version 1.5.0.
data RecordedEventType
  = ChannelComm String
  | BarrierSync String
  | ClockSync String deriving (Eq, Ord, Show)

getEventTypeVal :: RecordedEventType -> String
getEventTypeVal (ChannelComm s) = s
getEventTypeVal (BarrierSync s) = s
getEventTypeVal (ClockSync s) = s

-- Not really a CSP event, more like an enrollable poisonable alting barrier!
newtype Event = Event (
  Unique, -- Event identifier
  STM RecordedEventType, -- Event type for trace recording
  TVar (WithPoison
    (Int, -- Enrolled count
     Integer, -- Event sequence count
    [OfferSet]) -- A list of offer sets
 ))

instance Eq Event where
  (==) = (==) `on` getEventUnique

instance Ord Event where
  compare = compare `on` getEventUnique

-- For testing:
instance Show Event where
  show (Event (u, _t, _tv)) = "Event " ++ show (hashUnique u)

getEventUnique :: Event -> Unique
getEventUnique (Event (u,_,_)) = u

getEventTVar :: Event -> TVar (WithPoison (Int, Integer, [OfferSet]))
getEventTVar (Event (_,_,tv)) = tv

getEventType :: Event -> STM RecordedEventType
getEventType (Event (_,t,_)) = t

-- The value used to pass information to a waiting process once one of their events
-- has fired (and they have been committed to it).  The Int is an index into their
-- list of guards
newtype SignalValue = Signal (WithPoison Int)
  deriving (Eq, Show)

type SignalVar = TVar (Maybe (SignalValue, Map.Map Unique (Integer, RecordedEventType)))

addPoison :: SignalValue -> SignalValue
addPoison = const $ Signal PoisonItem

nullSignalValue :: SignalValue
nullSignalValue = Signal $ NoPoison (-1)

isNullSignal :: SignalValue -> Bool
isNullSignal (Signal n) = n == NoPoison (-1)

newtype OfferSet = OfferSet (SignalVar -- Variable to use to signal when committed
                , ProcessId -- Id of the process making the offer
                , [((SignalValue, STM ()), Map.Map Event ())]) -- Value to send when committed
                                    -- A list of all sets of events currently offered

instance Eq OfferSet where
  (==) = (==) `on` (\(OfferSet (tv,_,_)) -> tv)

instance Show OfferSet where
  show (OfferSet (_, pid, vs)) = "OfferSet " ++ show (pid, map (first fst) vs)

    -- Each event in the map can have three possible values:
    -- PoisonItem; event is poisoned, can always be completed
    -- NoPoison True; event has been chosen by previous process, you must choose
    -- it too
    -- NoPoison False; event has been rejected by previous process, you cannot
    -- choose it

unionAll :: Ord k => [Map.Map k a] -> Map.Map k a
unionAll [] = Map.empty
unionAll ms = foldl1 Map.union ms

allEventsInOffer :: OfferSet -> Map.Map Event ()
allEventsInOffer (OfferSet (_, _, [(_,es)])) = es
allEventsInOffer (OfferSet (_, _, eventSets)) = unionAll (map snd eventSets)

getAndIncCounter :: Event -> (a, b) -> STM (WithPoison (Integer, a))
getAndIncCounter e (r, _)
  = do x <- readTVar (getEventTVar e)
       case x of
         PoisonItem -> return PoisonItem
         NoPoison (a, !n, c) -> do writeTVar (getEventTVar e) $
                                     NoPoison (a, succ n, c)
                                   return $ NoPoison (n, r)

-- | search is /not/ used for discovering offers.  It is used for looking for possible
-- resolutions to a collection of offer sets.  It is pure; it performs no STM actions,
-- it just searches the offer-sets (which will have been discovered through STM)
-- for completions.
--
-- search performs a 2-dimensional traversal of the offers.  The search function
-- is called with a list of offer-sets.  For the offer-set at the head, it calls
-- tryAll.  tryAll searches through each offer in the offer-set, seeing if it can
-- be completed.  If it can, it calls search on the remaining offer-sets.  If this
-- fails, it reverts to trying the other offers in the list. The map of events passed through
-- relates to the previous things found in the search.
search :: [OfferSet]
          -- ^ The collection of all the related offer-sets
          -> Map.Map Event Bool
          -- ^ This contains the events already decided upon in the search.  If
          -- an event maps to True, it means it was chosen by an earlier part of
          -- the search, and thus future parts of the search /must/ have this event
          -- in the chosen offer (if the process offers it at all -- if it doesn't,
          -- it can be ignored).  If an event maps to False, it was already ruled
          -- out by not being chosen in another part of the search, and it cannot
          -- be chosen by any future parts of the search.  Should be empty when first called from the outside.
          -> Maybe ( [(SignalVar, SignalValue, STM ())]
                   , Map.Map Event (STM RecordedEventType, Set.Set ProcessId)
                   )
             -- ^ The list of tvars involved with the completion and the signal
             -- value for them, and the map with information about the completed events.
search [] _ = Just ([], Map.empty)
search (offer@(OfferSet (tv, pid, eventSets)) : offers) eventMap
      | Map.null mustChooseFromEventSets = tryAll eventSets
      | otherwise = tryAll filteredEventSets
      where
        allEventsInOfferMappedToFalse :: Map.Map Event Bool
        allEventsInOfferMappedToFalse = Map.map (const False) (allEventsInOffer offer)

        mustChooseFromEventSets :: Map.Map Event Bool
        mustChooseFromEventSets
          = (Map.filter id {- Keep all True events -} eventMap)
             `Map.intersection` allEventsInOfferMappedToFalse

        -- Only the offers containing all of the mustChooseFromEventSets
        filteredEventSets
          = [ off
            | off@(_,es) <- eventSets,
              Map.isSubmapOfBy (\_ _ -> True)
                mustChooseFromEventSets
                es
            ]

        -- Folds across a map, seeing if the given predicate holds for all values
        -- in the map.
        mapdotall :: Ord k => (a -> Bool) -> Map.Map k a -> Bool
        mapdotall f = Map.fold (\x b -> f x && b) True

        and' :: Ord k => Map.Map k Bool -> Bool
        and' = mapdotall id

        tryAll :: [((SignalValue, STM ()), Map.Map Event ())]
          -> Maybe ( [(SignalVar, SignalValue, STM ())]
                   , Map.Map Event (STM RecordedEventType, Set.Set ProcessId)
                   )
        tryAll [] = Nothing
        tryAll ((ns, es):next)
          | not $ and' (eventMap `Map.intersection` es)
              -- Contains an already-rejected event (one that mapped to False), skip:
              -- Need to reject the other events too though -- well, at least put
              -- them in the appropriate map and pass them through.  They will
              -- only be rejected if they are then not contained in the other chosen
              -- offer
              = tryAll next
          | otherwise = case search offers eventMap' of
            Nothing -> tryAll next
            Just (act, resolved) -> Just
              (if isNullSignal (fst ns) then act else (tv, fst ns, snd ns) : act
              , foldl (\m e -> Map.insertWith add e
                                 (getEventType e, Set.singleton pid) m)
                  resolved (Map.keys es)
              )
              
          where
            -- All events that features in other offers by this process, but not
            -- the current offer
            --
            -- It is very important here that union is left-biased for both unions.  We don't want
            -- to overwrite poison with acceptance, or acceptance with rejection.
            eventMap'
              = (eventMap `Map.union` (Map.map (const True) es)) `Map.union` allEventsInOfferMappedToFalse


            add (tx, pidsx) (_, pidsy) = (tx, pidsx `Set.union` pidsy)
            
-- Given a list of offers that could possibly complete, check if any set
-- of offers can.  If so, complete it (including all retractions and
-- notifications for each process), otherwise leave things untouched.
--
-- Takes an optional tvar identifier for the newest process to make an offer, the
-- list of all offer-sets that need to be considered (they will have come from
-- all events in a connected sub-graph), the map of relevant events to their status,
-- and returns the map of event-identifiers that did complete.
resolveOffers :: Maybe SignalVar -> [OfferSet] -> Set.Set Event
  -> STM (Map.Map Unique (RecordedEventType, Set.Set ProcessId))
resolveOffers newTvid allOffers events
  = do let (offers', _) = trim (allOffers, events)
           (act, ret) = fromMaybe ([], Map.empty) $
             search (map addNullOffer $ sortOffers offers') Map.empty
       -- The associated event-action must come first as that puts the values in the channels:
       mapM_ (\(_, _, m) -> m) act
       -- These values are then read by these on-completion bits:
       ret' <- T.mapM (\(m,y) -> do x <- m
                                    return (x, y)) ret
       eventCounts <- T.sequence $ Map.mapWithKey getAndIncCounter ret'
       let NoPoison uniqCounts = T.sequence $ Map.mapKeysMonotonic getEventUnique eventCounts
       mapM_ (\(tv, x, _) -> writeTVar tv (Just (x, uniqCounts))) act
       -- do the retractions for all involved processes once the choice is made:
       -- TODO optimise:
       retractOffers $ zip (map fst3 act)
                           (repeat $ unionAll $ map allEventsInOffer allOffers)
       return (Map.mapKeysMonotonic getEventUnique ret')
  where
    fst3 (x, _, _) = x
    -- Don't add the null offer for the newest process, and null offer should be
    -- added to the end:
    addNullOffer :: OfferSet -> OfferSet
    addNullOffer (OfferSet (tv,y,zs)) = OfferSet (tv,y,if Just tv == newTvid then zs else zs++nullOffer)

    nullOffer :: [((SignalValue, STM ()), Map.Map Event ())]
    nullOffer = [((nullSignalValue, return ()) ,Map.empty)]

-- Smallest offers first to minimise backtracking:
sortOffers :: [OfferSet] -> [OfferSet]
sortOffers xs
  | length xs > 2 = sortBy (compare `on` (\(OfferSet (_,_,es)) -> length es)) xs
  | otherwise = xs
-- TODO put the newest process first again

-- Given a list of offer-sets, and a map of events already-looked-at to their status,
-- trims the offer-sets by removing any option in an offer-set that cannot possibly
-- complete.  If this option includes any other events, any other options anywhere
-- that also feature these must be removed too.  The function iterates until it
-- finds a fix-point.
trim :: ([OfferSet], Set.Set Event) -> ([OfferSet], Set.Set Event)
    -- Each iteration, we remove all offersets that reference events that can
    -- never be ready, and if the removing of any of those causes an event
    -- to never become ready, we remove those events too, then we'll go round
    -- again (while finding the fix point)
trim (offers, events) = let ((events', changed), offers') =  mapAccumL trimOffer (events,
                                 False) offers
                            oe = (offers', events')
                        in if changed then trim oe else oe 
  where
    trimOffer :: (Set.Set Event, Bool) -> OfferSet -> ((Set.Set Event, Bool), OfferSet)
    trimOffer (es, changed) o@(OfferSet (tv, pid, eventSets))
            -- An offer is only retained if all the events are in the set of events
            -- that can possibly complete
          = let (eventSetsToRemove, eventSetsTrimmed)
                  | Set.size es == 1 = partition (\(_,x) -> Map.size x /= 1 || fst (Map.findMin x) /= Set.findMin es) eventSets
                  | otherwise = partition (\(_,x) -> not $ (Map.keysSet x) `Set.isSubsetOf` es) eventSets
               -- If any of the events to remove are not also in sets that will
               -- be kept, and the event is not poisoned, that event is no longer completable and should be
               -- removed from the set of events:
                eventsNotCompletable = Map.keysSet $ 
                  (unionAll $ map snd eventSetsToRemove)
                   `Map.difference` (unionAll $ map snd eventSetsTrimmed)
                changed' = changed
                           || not (null eventSetsToRemove)
            in if null eventSetsToRemove then ((es, changed), o)
               else 
               ((es `Set.difference` eventsNotCompletable, changed'),
                OfferSet (tv, pid, eventSetsTrimmed))

-- Semantics of poison with waiting for multiple events is that if /any/ of
-- the events are poisoned, that whole offer will be available immediately
-- even if the other channels are not ready, and the body will throw a poison
-- exception rather than running any of the guards.  This is because in situations
-- where for example you want to wait for two inputs (e.g. invisible process
-- or philosopher's forks) you usually want to forward poison from one onto
-- the other.


-- Finds all the events that could be linked to the given one.
--
-- Given an event, spiders out and discovers all events (connected via mutual offers).
--  Returns the list of offer-sets found.  It also
-- returns a set containing each connected completable event.
-- If any of the events are found to be poisoned, the associated STM action is
-- executed
discoverRelatedOffers :: [(STM (), Event)] -> STM (WithPoison ([OfferSet], Set.Set Event))
discoverRelatedOffers = discoverRelatedOffersAll $ NoPoison ([], Set.empty)
  where
    -- We need the supplied STM () actions for each event to take precedence over
    -- the default ones supplied later in the algorithm.  So if, for example, the
    -- user supplies a,b and c in the list, but our usual depth-first search would
    -- lead a -> d -> c, we do not want to use the default event for c instead
    -- of the supplied one.  Therefore we maintain the work list explicitly.
    
    -- Nothing means that that event is poisoned (and thus always ready)
    discoverRelatedOffersAll :: WithPoison ([OfferSet], Set.Set Event)
      -> [(STM (), Event)]
      -> STM (WithPoison ([OfferSet], Set.Set Event))
    discoverRelatedOffersAll PoisonItem _ = return PoisonItem
    discoverRelatedOffersAll x [] = return x
    discoverRelatedOffersAll a@(NoPoison (accum, events)) ((act,e@(Event (_, _, tv))):next)
        -- Don't process the same event multiple times:
      | e `Set.member` events = discoverRelatedOffersAll a next
      | otherwise
          = do x <- readTVar tv
               case x of
                 PoisonItem -> act >> return PoisonItem
                 NoPoison (count, _, offers) ->
                   let otherEvents = map allEventsInOffer offers in
                   if length offers == count
                     then -- It could be ready
                          discoverRelatedOffersAll
                            (NoPoison (accum ++ offers, Set.insert e events))
                            -- If the offers only have one event, must be this
                            -- one:
                            (if Map.size (unionAll otherEvents) == 1
                               then next
                               else next ++ zip (repeat $ return ())
                                         (Map.keys $ unionAll otherEvents))
                     else -- No way it could be ready, so ignore it:
                       discoverRelatedOffersAll a next

-- Given an event, spiders out, discovers all the offers, then resolves them
-- and returns a map containing all the completed events, mapping the
-- identifier to the event type and the set of process identifiers that
-- participated in the succesfully completed events.  The map will be empty if
-- and only if no events were completed.
discoverAndResolve :: Either OfferSet Event
                        -- ^ Either an OfferSet to spider out from, or a single
                        -- event.  The latter case is for when we are resigning
                        -- from an event and need to check if that completes anything.
                      -> STM (WithPoison (Map.Map Unique (RecordedEventType, Set.Set ProcessId)))
                        -- ^ Gives back either poison, or a map from event identifiers
                        -- to information about the completed event.  The map is
                        -- empty if no events were completed.
discoverAndResolve offOrEvent
  = do r <- discoverRelatedOffers $ case offOrEvent of
              Left off@(OfferSet (tv, _, nes)) ->
                let retract = retractOffers [(tv, allEventsInOffer off)] in
                      concat [zip
                        -- This is the action to execute if an event is found to
                        -- be poisoned:
                        (repeat $ retract >> writeTVar tv (Just (addPoison ns, Map.empty)))
                        (Map.keys es)
                        | ((ns,_), es) <- nes]
              Right e -> [(return (), e)]
       case r of
         PoisonItem -> return PoisonItem
         NoPoison (m, s) -> liftM NoPoison $ resolveOffers tvid (nub m) s
  where
    tvid = case offOrEvent of
             Left (OfferSet (tv, _, _)) -> Just tv
             _ -> Nothing

newEvent :: STM RecordedEventType -> Int -> IO Event
newEvent t n
  = do u <- newUnique
       atomically $ do tv <- newTVar (NoPoison (n, 0, []))
                       return $ Event (u, t, tv)

newEventUnique :: IO Unique
newEventUnique = newUnique

enrollEvent :: Event -> STM (WithPoison ())
enrollEvent e
  = do x <- readTVar $ getEventTVar e
       case x of
         PoisonItem -> return PoisonItem
         NoPoison (count, seq, offers) ->
           do writeTVar (getEventTVar e) $ NoPoison (count + 1, seq, offers)
              return $ NoPoison ()

-- If the event completes, we return details related to it:
resignEvent :: Event -> STM (WithPoison [((RecordedEventType, Unique), Set.Set ProcessId)])
resignEvent e
  = do x <- readTVar $ getEventTVar e
       case x of
         PoisonItem -> return PoisonItem
         NoPoison (count, seq, offers) ->
           do writeTVar (getEventTVar e) $ NoPoison (count - 1, seq, offers)
              if count - 1 == length offers
                then liftM (fmap $ \mu -> [((r,u),pids) | (u,(r,pids)) <- Map.toList mu])
                       $ discoverAndResolve $ Right e
                else return $ NoPoison []

-- Given the list of identifiers paired with all the events that that process might
-- be engaged in, retracts all the offers that are associated with the given TVar;
-- i.e. the TVar is used as an identifier for the process
retractOffers :: [(SignalVar, Map.Map Event ())] -> STM ()
retractOffers = mapM_ retractAll
  where
    retractAll :: (SignalVar, Map.Map Event ()) -> STM ()
    retractAll (tvid, evts) = mapM_ retract (Map.keys evts)
      where
        retract :: Event -> STM ()
        retract e
          = do x <- readTVar $ getEventTVar e
               case x of
                 PoisonItem -> return ()
                 NoPoison (enrolled, seq, offers) ->
                   let reducedOffers = filter (\(OfferSet (tvx,_,_)) -> tvx /= tvid) offers in
                   writeTVar (getEventTVar e) $ NoPoison (enrolled, seq, reducedOffers)

-- Simply adds the offers but doesn't check if that will complete an event:
-- Returns PoisonItem if any of the events were poisoned
makeOffers :: OfferSet -> STM (WithPoison ())
makeOffers offers
  = do let allEvents = Map.keys $ allEventsInOffer offers
       -- No need for nub, as having it come from a map guarantees there are no
       -- duplicates in the list of events
       liftM mergeWithPoison $ mapM makeOffer allEvents
  where
    makeOffer :: Event -> STM (WithPoison ())
    makeOffer e
      = do x <- readTVar $ getEventTVar e
           case x of
             PoisonItem -> return PoisonItem
             NoPoison (count, seq, prevOffers) ->
               do writeTVar (getEventTVar e) $ NoPoison (count, seq, offers : prevOffers)
                  return $ NoPoison ()

-- Returns Nothing if no events were ready.  Returns Just with the signal value
-- if an event was immediately available, followed by the information for each
-- event involved in the synchronisation.  If poison was encounted, this list will
-- be empty.
enableEvents :: SignalVar
                  -- ^ Variable used to signal the process once a choice is made
                -> ProcessId
                  -- ^ The id of the process making the choice
                -> [((SignalValue, STM ()), [Event])]
                  -- ^ The list of options.  Each option has a signalvalue to return
                  -- if chosen, and a list of events (conjoined together).
                  --  So this list is the disjunction of conjunctions, with a little
                  -- more information.
                -> Bool
                  -- ^ True if it can commit to waiting.  If there is an event
                  -- combination ready during the transaction, it will chosen regardless
                  -- of the value of this flag.  However, if there no events ready,
                  -- passing True will leave the offers there, but False will retract
                  -- the offers.
                -> STM (Maybe ((SignalValue, Map.Map Unique (Integer, RecordedEventType)), [((RecordedEventType, Unique), Set.Set ProcessId)]))
enableEvents tvNotify pid events canCommitToWait
  = do let offer = OfferSet (tvNotify, pid, [(nid, Map.fromList (zip es (repeat ()))) | (nid, es) <- events])
       -- First add our offer to all the events:
       -- We don't check the result for poison, as discoverAndResolve will find
       -- it anyway
       makeOffers offer
       -- Then spider out and see if anything can be resolved:
       pmu <- discoverAndResolve (Left offer)
       case (canCommitToWait, pmu) of
         (_, PoisonItem) -> do Just chosen <- readTVar tvNotify
                               return $ Just (chosen, [])
         (True, NoPoison mu) | Map.null mu -> return Nothing
         (False, NoPoison mu) | Map.null mu ->
           do retractOffers [(tvNotify, Map.fromList $ zip es (repeat ())) | (_,es) <- events]
              return Nothing
         (_, NoPoison mu) -> -- Need to turn all the Unique ids back into the custom-typed
                   -- parameter that the user gave in the list.  We assume
                   -- it will be present:
                do {- let y = mapMaybe (\(k,v) -> listToMaybe [(x,v) | (x,_,_,es) <- events,
                              k `elem` map getEventUnique es]) $ Map.toList mu
                                -}
                   Just chosen <- readTVar tvNotify
                   return $ Just (chosen, [((r,u),pids) | (u,(r,pids)) <- Map.toList mu])

-- | Given the variable used to signal the process, and the list of events that
-- were involved in its offers, attempts to disable the events.  If the variable
-- has been signalled (i.e. has a Just value), that is returned and nothing is done, if the variable
-- has not been signalled (i.e. is Nothing), the events are disabled and Nothing
-- is returned.
disableEvents :: SignalVar -> [Event] -> STM (Maybe (SignalValue, Map.Map Unique (Integer,
  RecordedEventType)))
disableEvents tv events
  = do x <- readTVar tv
       -- Since the transaction will be atomic, we know
       -- now that we can disable the barriers and nothing fired:
       when (isNothing x) $
         retractOffers [(tv, Map.fromList $ zip events (repeat ()))]
       return x

checkEventForPoison :: Event -> STM (WithPoison ())
checkEventForPoison e
  = do x <- readTVar $ getEventTVar e
       case x of
         PoisonItem -> return PoisonItem
         _ -> return (NoPoison ())

poisonEvent :: Event -> STM ()
poisonEvent e
  = do x <- readTVar $ getEventTVar e
       case x of
         PoisonItem -> return ()
         NoPoison (_, _, offers) ->
           do retractOffers [(tvw, unionAll $ map snd events)
                            | OfferSet (tvw, _, events) <- offers]
              sequence_ [writeTVar tvw (Just (addPoison $ pickInts events, Map.empty))
                        | OfferSet (tvw, _, events) <- offers]
              writeTVar (getEventTVar e) PoisonItem
  where
    pickInts :: [((SignalValue, STM ()), Map.Map Event ())] -> SignalValue
    pickInts es = case filter ((e `Map.member`) . snd) es of
      [] -> nullSignalValue -- Should never happen
      (((ns,_),_):_) -> ns

--TODO document how if it's poisoned, 0 will be appended to the list

----------------------------------------------------------------------
----------------------------------------------------------------------
-- Testing:
----------------------------------------------------------------------
----------------------------------------------------------------------
#ifdef CHP_TEST

-- Tests if two lists have the same elements, but not necessarily in the same order:
(**==**) :: Eq a => [a] -> [a] -> Bool
a **==** b = (length a == length b) && null (a \\ b)

(**/=**) :: Eq a => [a] -> [a] -> Bool
a **/=** b = not $ a **==** b

testDiscover :: Test
testDiscover = TestCase $
    do test "Empty discover" [(NoPoison 1, False)] [] [0]
       test "Single full event" [(NoPoison 1, True)] [(True, [[0]])] [0]
       test "Two separate events A" [(NoPoison 1, True), (NoPoison 1, False)]
         [ (True, [[0]]), (False, [[1]]) ] [0]
       test "Two separate events B" [(NoPoison 1, False), (NoPoison 1, True)]
         [ (False, [[0]]), (True, [[1]]) ] [1]
       test "Two separate events A, non-completable" [(NoPoison 2, False), (NoPoison 1, False)]
         [ (False, [[0]]), (False, [[1]]) ] [0]
       test "Three channels, linked by two OR-offerers"
         [(NoPoison 2, False), (NoPoison 2, True), (NoPoison
           2, False)]
         (zip (repeat True) [ [[0],[1]] , [[1],[2]] ]) [1,2]
       test "Three channels, linked by two AND-offerers"
         [(NoPoison 2, False), (NoPoison 2, True), (NoPoison
           2, False)]
         (zip (repeat True) [ [[0,1]] , [[1,2]] ]) [0,1]
       test "Three barriers, one process offering all pairs"
         (replicate 3 (NoPoison 2, False))
         [(False,[ [0,1], [0,2], [1,2] ])] [0]
       -- Discovery on a poisoned event will not find offers associated with
       -- that event because they are not stored.  The local offer is added
       -- in discoverAndResolve, not testDiscover, so for poison we expect
       -- to find nothing:
       test_Poison "Single poisoned event" [PoisonItem] [ [[0]] ] [0]
       test_Poison "Two poisoned events"
         [PoisonItem, PoisonItem]
         [ [[0,1]] ] [0,1]
       test_Poison "One poisoned, one non-poisoned event"
         [PoisonItem, NoPoison 1] [ [[0,1]] ] [0,1]
  where
    test :: String ->
            {- Events: -} [(WithPoison Int {-count -}, Bool {- Should be in set -})] ->
            {- Offers: -} [(Bool, [[Int] {- events -}])] -> {-Starting events: -} [Int] -> IO ()
    test testName eventCounts offerSets startEvents
      = do (events, realOffers) <- makeTestEvents (map fst eventCounts) (map snd offerSets)
           let expectedResult
                = ([off | ((yes, _),off) <- zip offerSets realOffers, yes]
                  ,Set.fromList [e
                                | (e,(_count, present)) <- zip events eventCounts,
                                  present])
           act <- atomically $ discoverRelatedOffers
             $ zip (repeat $ return ()) $ map (events!!) startEvents
           case act of
             PoisonItem -> assertFailure $ testName ++ "Unexpected poison"
             NoPoison actualResult -> do
               when (fst expectedResult **/=** fst actualResult)
                 $ assertFailure $ testName ++ " failed offers, exp: "
                   ++ show (length $ fst expectedResult)
                   ++ " got: " ++ show (length $ fst actualResult)
               when (snd expectedResult /= snd actualResult)
                 $ assertFailure $ testName ++ " failed events "
                   ++ "exp: " ++ show (snd expectedResult)
                   ++ "but got: " ++ show (snd actualResult)
    test_Poison :: String ->
            {- Events: -} [WithPoison Int {-count -}] ->
            {- Offers: -} [[[Int] {- events -}]] -> {-Starting events: -} [Int] -> IO ()
    test_Poison testName eventCounts offerSets startEvents
      = do (events, _realOffers) <- makeTestEvents eventCounts offerSets
           act <- atomically $ discoverRelatedOffers
             $ zip (repeat $ return ()) (map (events!!) startEvents)
           case act of
             PoisonItem -> return ()
             NoPoison _ -> assertFailure $ testName ++ " expected poison but none"



testTrim :: Test
testTrim = TestCase $
    do test "Empty trim" [(NoPoison 1, False)] [] [0]
       test "Trim, Three channels, linked by two OR-offerers"
         [(NoPoison 2, False), (NoPoison 2, True), (NoPoison 2, False)]
         [ [(False, [0]), (True, [1])] , [(True, [1]), (False, [2])] ] [1]
       test "Trim, simplified santa not complete"
         (replicate 4 (NoPoison 2, False))
         [ zip (repeat False) [[0,1,2],[0,1,3],[0,2,3],[1,2,3]], [(False, [0])],
           [(False, [1])]] [0]
       test "Trim, simplified santa complete"
         (replicate 3 (NoPoison 2, True) ++ [(NoPoison 2, False)])
         [ [(True,[0,1,2]),(False,[0,1,3]),(False,[0,2,3]),(False,[1,2,3])], [(True, [0])],
           [(True, [1])], [(True, [2])]] [0]
  where
    test :: String -> 
            {- Events: -} [(WithPoison Int {-count -}, Bool {- expected kept -})] ->
            {- Offers: -} [ [(Bool, [Int]) {- events -}]] -> {-Starting events:-} [Int] -> IO ()
    test testName eventCounts offerSets startEvents
      = do (events, realOffers) <- makeTestEvents (map fst eventCounts) (map (map snd) offerSets)
           let expectedResult' = NoPoison
                  ([OfferSet (tv,pid,[off | (m,off) <- zip [0..] offs, fst $ offerSets !! n !! m])
                   | (n,OfferSet (tv,pid,offs)) <- zip [0..] realOffers]
                  ,Set.fromList [events !! n
                                | (n,(_count, present)) <- zip [0..] eventCounts,
                                  present])
           actualResult' <- liftM (fmap $ trim  . (\(xs,y) -> (nub $ maybe id (:) (listToMaybe realOffers) xs, y)))
             $ atomically $ discoverRelatedOffers $ zip (repeat $ return ()) (map (events!!) startEvents)
           case (expectedResult', actualResult') of
             (PoisonItem, PoisonItem) -> return ()
             (PoisonItem, _) -> assertFailure $ testName ++ " expected poison but none found"
             (_, PoisonItem) -> assertFailure $ testName ++ " unexpected poison"
             (NoPoison expectedResult, NoPoison actualResult)
               -> do
             when (fst expectedResult **/=** fst actualResult)
               $ assertFailure $ testName ++ " failed offers, exp: "
               ++ show (length $ fst expectedResult)
               ++ " got: " ++ show (length $ fst actualResult)
             when (snd expectedResult /= snd actualResult)
               $ assertFailure $ testName ++ " failed events, exp: "
               ++ show (snd expectedResult)
               ++ "but got: " ++ show (snd actualResult)

testPoison :: Test
testPoison = TestCase $ do
  test "Poison empty event" [(NoPoison 2, PoisonItem)] [] 0
  test "Poison, single offerer" [(NoPoison 2, PoisonItem)] [[[0]]] 0
  test "Poison, offered on two (AND)" [(NoPoison 2, PoisonItem), (NoPoison 2, NoPoison [])] [[[0,1]]] 0
  test "Poison, offered on two (OR)" [(NoPoison 2, PoisonItem), (NoPoison 2, NoPoison [])] [[[0],[1]]] 0
  where
    test :: String ->
      [(WithPoison Int {-count -}, WithPoison [Int] {- remaining offers -})] ->
      {- Offers: -} [[[Int] {- events -}]] -> Int {-Poison Event-} -> IO ()

    test testName eventCounts offerSets poisoned = do
      (events, realOffers) <- makeTestEvents (map fst eventCounts) offerSets
      atomically $ poisonEvent $ events !! poisoned
      -- Now we must check that the event is poisoned, and that all processes
      -- that were offering on that event have had their offers retracted (by
      -- checking that only the specified offers remain on each event)

      sequence_ [do x <- atomically $ readTVar $ getEventTVar $ events !! n
                    case (expect, x) of
                      (PoisonItem, PoisonItem) -> return ()
                      (NoPoison _, PoisonItem) -> assertFailure $ testName ++
                        " expected no poison but found it"
                      (PoisonItem, NoPoison _) -> assertFailure $ testName ++
                        " expected poison but found none"
                      (NoPoison expOff, NoPoison (_, _, actOff)) ->
                        when (map (realOffers !!) expOff **/=** actOff) $
                          assertFailure $ testName ++ " offers did not match"
                | (n, (_, expect)) <- zip [0..] eventCounts]


    
testAll :: Test
testAll = TestList [testDiscover, testTrim, testResolve, testPoison]

makeTestEvents ::
            {- Events: -} [WithPoison Int {-count -}] ->
            {- Offers: -} [[[Int] {- events -}]] -> IO ([Event], [OfferSet])
            -- Offers is a list of list of list of ints
            -- Outermost list is one-per-process
            -- Middle list is one-per-offer
            -- Inner list is a conjunction of events
makeTestEvents eventCounts offerSets
      = do events <- mapM (\n -> newEvent (return $ ChannelComm "") $ case n of
             NoPoison n' -> n'
             PoisonItem -> 0) eventCounts
           -- Poison all the events marked as poisoned:
           atomically $ sequence_ [writeTVar tv PoisonItem | (n,Event (_,_,tv)) <- zip [0..] events,
             eventCounts !! n == PoisonItem]
           realOffers <- sequence
             [ do tv <- atomically $ newTVar Nothing
                  let pid = testProcessId processN
                      -- TODO test the STM actions too
                      offSub = [ ((Signal $ NoPoison (processN + offerN), return ()),
                                  Map.fromList [ (events !! indivEvent, ())
                                  | indivEvent <- singleOffer])
                               | (offerN, singleOffer) <- zip [0..] processOffers]
                      off = OfferSet (tv, pid, offSub)
                  mapM_ (\e -> atomically $ do
                    x <- readTVar (getEventTVar e)
                    case x of
                      NoPoison (count, s, offs) ->
                        writeTVar (getEventTVar e) $ NoPoison (count, s, off : offs)
                      PoisonItem -> return ()
                    ) (Map.keys $ unionAll $ map snd offSub)
                  return off
             | (processN, processOffers) <- zip (map (*1000) [0..]) offerSets]
           return (events, realOffers)

testResolve :: Test
testResolve = TestCase $
    do test "Empty Resolve" [(NoPoison 0, Right [])] [[]]
       test "Single offer" [(NoPoison 1, Left [(0,0)])] [[[0]]]
       test "Not enough" [(NoPoison 2, Right [0])] [[[0]]]
       test "One channel" [(NoPoison 2, Left [(0,0),(1,0)])] [[[0]],[[0]]]
       test "Two channels, two single offerers and one double"
         [(NoPoison 2, Left [(0,0),(2,0)]), (NoPoison 2, Left [(1,0),(2,0)])]
         [ [[0]], [[1]], [[0,1]] ]
       test "Two channels, two single offerers and one choosing"
         [(NoPoison 2, Left [(0,0),(2,0)]), (NoPoison 2, Right [1])]
         [ [[0]], [[1]], [[0],[1]] ]
       test "Three channels, both offering different pair"
         [(NoPoison 2, Right []), (NoPoison 2, Left [(0,1),(1,0)]), (NoPoison 2, Right [])]
         [ [[0],[1]] , [[1],[2]] ]
       -- This test is a bit hacky, given there are two valid results:
       test "Two channels, both could complete"
         [(NoPoison 2, Left [(0,0),(1,0)]), (NoPoison 2, Right [])]
         [ [[0],[1]] , [[0],[1]] ]
       test "Three channels, any could complete"
         [(NoPoison 2, Left [(0,0),(1,0)]), (NoPoison 2, Right [2]), (NoPoison 2,
           Right [2])]
         [ [[0],[1]] , [[0],[2]], [[1],[2]] ]
       test "Three channels, one guy offering three pairs, two single offerers"
         [(NoPoison 2, Left [(0,1),(1,0)]), (NoPoison 2, Right []), (NoPoison 2,
           Left [(0,1),(2,0)])]
         [ [[0,1],[0,2],[1,2]], [[0]], [[2]] ]
       test "Three channels, one guy offering three pairs, three single offerers"
         [(NoPoison 2, Left [(0,0),(1,0)]), (NoPoison 2, Left [(0,0),(2,0)]), (NoPoison 2,
           Right [3])]
         [ [[0,1],[0,2],[1,2]], [[0]], [[1]], [[2]] ]
       test "Four channels, one guy offering sets of three, three single offerers"
         [(NoPoison 2, Left [(0,0),(1,0)]), (NoPoison 2, Left [(0,0),(2,0)]),
          (NoPoison 2, Left [(0,0),(3,0)]), (NoPoison 2, Right [])]
         [ [[0,1,2],[0,1,3],[0,2,3],[1,2,3]], [[0]], [[1]], [[2]] ]
       test "Four channels, one guy offering sets of three, two single offerers"
         [(NoPoison 2, Right [1,0]), (NoPoison 2, Right [2,0]),
          (NoPoison 2, Right [0]), (NoPoison 2, Right [0])]
         [ [[0,1,2],[0,1,3],[0,2,3],[1,2,3]], [[0]], [[1]] ]
       -- test resolutions with poison:
       --
       test' "One event, poisoned"
         [(PoisonItem, Left [(0,0)])]
         [[[0]]] True
       test' "Two events, one poisoned"
         [(PoisonItem, Left [(0,0)]), (NoPoison 2, Left [(0,0)])]
         [[[0,1]]] True
  where
    test testName eventCounts offerSets = test' testName eventCounts offerSets False
    
    test' :: String ->
      -- List of events:
      [(WithPoison Int {- enrolled count -},
        Either [(Int, Int)] {- success: expected process, offer indexes -}
               [Int] {- remaining offers -})] ->
      {- Offers: -} [[[Int] {- events -}]] -> Bool {-Poisoned-} -> IO ()

    test' testName eventCounts offerSets poisoned = do
           (events, realOffers) <- makeTestEvents (map fst eventCounts) offerSets

           actualResult <- liftM (liftM (fmap snd)) $ atomically $ discoverAndResolve $ Left $ head realOffers
           let expectedResult = if poisoned then PoisonItem else NoPoison $
                                Map.fromList [ (getEventUnique e,
                                               Set.fromList $ map (testProcessId . (*1000) . fst) is)
                                             | (e, Left is) <- zip events (map snd eventCounts)]
           when (expectedResult /= actualResult) $
             assertFailure $ testName ++ " failed on direct result, expected: "
               ++ showStuff expectedResult ++ " got: " ++ showStuff actualResult

           allFired <- liftM concat $ mapM (flip either (const $ return []) $ mapM $ \(pn, en) ->
             let OfferSet (tv,_,_) = realOffers !! pn in
               do x <- atomically $ readTVar tv
                  case x of
                    Nothing -> assertFailure $ "Unexpected no-win for " ++ show (pn,en)
                    Just v -> when (fst v /= (if poisoned then addPoison else id)
                                             (Signal $ NoPoison ((pn*1000)+en))) $
                      assertFailure $ testName ++ " wrong choice: " ++ " exp: " ++ show
                        (pn+en)
                  return pn
             ) $ map snd eventCounts
           -- test the others are unchanged
           sequence_ [ let OfferSet (tv,_,_) = realOffers !! n in
                         do x <- atomically $ readTVar tv
                            case x of
                              Nothing -> return ()
                              Just _ -> assertFailure $ testName ++ " Unexpected win for process: " ++
                                show n
                     | n <- [0 .. length offerSets - 1] \\ allFired]
           -- check events are blanked afterwards:
           c <- sequence
                     [ let e = events !! n
                           expVal = case st of
                             Left _ -> []
                             Right ns -> map (realOffers !!) ns in do
                         x <- atomically $ readTVar $ getEventTVar e
                         case x of
                           NoPoison (c, _, e') -> return $ Just ((count, expVal), (c, e'))
                           _ -> do assertFailure $ testName ++ " unexpected poison"
                                   return Nothing
{-                           NoPoison (c, _, e') | c == count && e' == expVal -> return ()
                           _ ->
                             assertFailure $ testName ++ "Event " ++ show n ++
                             " not as expected after, exp: " ++ show (length expVal)
                             ++ " act: " ++ (let NoPoison (_,_,act) = x in show (length act))-}
                     | (n,(NoPoison count, st)) <- zip [0..] eventCounts]
           uncurry (assertEqual testName) (unzip $ catMaybes c)

    showStuff = show . fmap (map (first hashUnique) . Map.toList)

#endif
-- CHP_TEST