The sad state of property-based testing libraries

gLOBAL_COUNTER :: IORef Int
gLOBAL_COUNTER = unsafePerformIO (newIORef 0)
{-# NOINLINE gLOBAL_COUNTER #-}

incr :: IO ()
incr = do
  n <- readIORef gLOBAL_COUNTER
  writeIORef gLOBAL_COUNTER (n + 1)

get :: IO Int
get = readIORef gLOBAL_COUNTER

newtype Counter = Counter Int
  deriving (Eq, Show)

-- We'll come back to the definition of the `StateModel` type class after this
-- example.
instance StateModel Counter where

  -- We start counting from zero.
  initialState :: Counter
  initialState = Counter 0

  -- The commands correspond to the names of the functions that operate on the
  -- global counter.
  data Command Counter r
    = Incr
    | Get
    deriving (Show, Functor)

  -- The responses correspond to the return types of each function. By
  -- convention we'll add a underscore suffix to a response of the corresponding
  -- command.
  data Response Counter r
    = Incr_ ()
    | Get_ Int
    deriving (Eq, Show, Functor, Foldable)

  -- The state machine takes a command and the model of the counter and returns
  -- a new model and a response. We'll come back to the role of the `Either` later.
  runFake :: Command Counter r -> Counter -> Either Void (Counter, Response Counter r)
  runFake Incr  (Counter n) = return (Counter (n + 1), Incr_ ())
  runFake Get m@(Counter n) = return (m, Get_ n)

  -- We also need to explain which part of the counter API each command
  -- corresponds to.
  runReal :: Command Counter r -> IO (Response Counter r)
  runReal Get  = Get_  <$> get
  runReal Incr = Incr_ <$> incr

  -- We'll generate increments and reads of the counter with equal probability.
  -- Notice that we only need to explain how to generate a single command, the
  -- library will use this to generate sequences of commands as we'll see later.
  generateCommand :: Counter -> Gen (Command Counter r)
  generateCommand _s = elements [Incr, Get]

prop_counter :: Commands Counter -> Property
prop_counter cmds = monadicIO $ do
  runCommands cmds
  run reset
  assert True

reset :: IO ()
reset = writeIORef gLOBAL_COUNTER 0

incr42Bug :: IO ()
incr42Bug = do
  n <- readIORef gLOBAL_COUNTER
  let n' = if n == 42
           then n -- BUG
           else n + 1
  writeIORef gLOBAL_COUNTER n'

- runReal Incr = Incr_ <$> incr
+ runReal Incr = Incr_ <$> incr42Bug

class ( ...

      ) => StateModel state where

  -- If we think of the system under test as a black box, then commands are the
  -- inputs and responses the outputs to the black box.
  data Command  state :: Type -> Type
  data Response state :: Type -> Type

  -- Sometimes a command needs to refer to a previous response, e.g. when a file
  -- is opened we get a handle which is later refered to when writing or reading
  -- form the file. File handles, and similar constructs, are called references
  -- and can be part of commands and responses.
  type Reference state :: Type
  type Reference state = Void

  -- Not all commands are valid in all states. Pre-conditions allow the user to
  -- specify when a command is safe to execute, for example we cannot write or
  -- read to or from an unopened file. The `PreconditionFailure` data type
  -- allows the user to create custom pre-condition failures. By default now
  -- pre-condition failures are allowed, thus the `Void` (empty) type.
  type PreconditionFailure state :: Type
  type PreconditionFailure state = Void


  generateCommand :: state -> Gen (Command state (Var (Reference state)))

  shrinkCommand :: state -> Command state (Var (Reference state))
                -> [Command state (Var (Reference state))]
  shrinkCommand _state _cmd = []

  initialState :: state

  runFake :: Command state (Var (Reference state)) -> state
          -> Either (PreconditionFailure state)
                    (state, Response state (Var (Reference state)))

  runReal :: Command state (Reference state)
          -> CommandMonad state (Response state (Reference state))

  monitoring :: (state, state)
             -> Command state (Reference state)
             -> Response state (Reference state)
             -> Property -> Property
  monitoring _states _cmd _resp = id

  commandName :: (Show (Command state ref), Show ref)
              => Command state ref -> String
  commandName = head . words . show

  -- Most often the result of executing a command against the system under test
  -- will live in the IO monad, but sometimes it can be useful to be able a
  -- choose another monad.
  type CommandMonad state :: Type -> Type
  type CommandMonad state = IO

newtype Commands state = Commands
  { unCommands :: [Command state (Var (Reference state))] }

precondition :: StateModel state
             => state -> Command state (Var (Reference state)) -> Bool
precondition s cmd = case runFake cmd s of
  Left _  -> False
  Right _ -> True

nextState :: StateModel state
          => state -> Command state (Var (Reference state)) -> state
nextState s cmd = case runFake cmd s of
  Right (s', _) -> s'
  Left _err -> error "nextState: impossible, we checked for success in precondition"

instance StateModel state => Arbitrary (Commands state) where

  arbitrary :: Gen (Commands state)
  arbitrary = Commands <$> genCommands initialState
    where
      genCommands :: StateModel state
                  => state -> Gen [Command state (Var (Reference state))]
      genCommands s = sized $ \n ->
        let
          w = n `div` 2 + 1
        in
          frequency
            [ (1, return [])
            , (w, do mcmd <- generateCommand s `suchThatMaybe` precondition s
                     case mcmd of
                       Nothing  -> return []
                       Just cmd -> (cmd :) <$> genCommands (nextState s cmd))
            ]

  shrink :: Commands state -> [Commands state]
  shrink = pruneShrinks . possibleShrinks
    where
      possibleShrinks :: Commands state -> [Commands state]
      possibleShrinks = map (Commands . map fst) . shrinkList shrinker
                      . withStates initialState . unCommands
        where
          shrinker (cmd, s) = [ (cmd', s) | cmd' <- shrinkCommand s cmd ]

          withStates :: StateModel state
                     => state -> [Command state (Var (Reference state))]
                     -> [(Command state (Var (Reference state)), state)]
          withStates s0 = go s0 []
            where
              go _s acc []           = reverse acc
              go  s acc (cmd : cmds) = go (nextState s cmd) ((cmd, s) : acc) cmds

      pruneShrinks :: [Commands state] -> [Commands state]
      pruneShrinks = coerce . filter (not . null)
                   . map (go initialState Set.empty [] . unCommands)
        where
          go _s _vars acc [] = reverse acc
          go  s  vars acc (cmd : cmds)
            | not (scopeCheck vars cmd) = go s vars acc cmds
            | otherwise = case runFake cmd s of
                Left _preconditionFailure -> go s vars acc cmds
                Right (s', resp) ->
                  let
                    returnedVars = Set.fromList (toList resp)
                    vars' = returnedVars `Set.union` vars
                  in
                    go s' vars' (cmd : acc) cmds

scopeCheck :: Foldable (Command state)
           => Set (Var a) -> Command state (Var a) -> Bool
scopeCheck varsInScope cmd = usedVars `Set.isSubsetOf` varsInScope
  where
    usedVars = Set.fromList (toList cmd)

runCommands :: forall state. StateModel state
            => Commands state -> PropertyM (CommandMonad state) ()
runCommands (Commands cmds0) = go initialState emptyEnv cmds0
  where
    go :: state -> Env state -> [Command state (Var (Reference state))]
       -> PropertyM (CommandMonad state) ()
    go _state _env [] = return ()
    go  state  env (cmd : cmds) = do
      case runFake cmd state of
        Left err -> do
          monitor (counterexample ("Preconditon failed: " ++ show err))
          assert False
        Right (state', resp) -> do
          let name = commandName cmd
          monitor (tabulate "Commands" [name] . classify True name)
          -- Here we substitute all symbolic references for real ones:
          let ccmd = fmap (lookupEnv env) cmd
          cresp <- run (runReal ccmd)
          monitor (counterexample (show cmd ++ " --> " ++ show cresp))
          monitor (monitoring (state, state') ccmd cresp)
          -- Here we collect all references from the response and store it in
          -- our environment, so that subsequence commands can be substituted.
          let refs   = toList cresp
              env'   = extendEnv env (zip [sizeEnv env..] refs)
              cresp' = fmap (lookupEnv env') resp
              ok     = cresp == cresp'
          unless ok $
            monitor (counterexample ("Expected: " ++ show cresp' ++ "\nGot: " ++ show cresp))
          -- And finally here's where we assert that the model and the real
          -- implementation agree.
          assert ok
          go state' env' cmds

newtype Env state = Env { unEnv :: IntMap (Reference state) }

typedef struct queue {
  int *buf;
  int inp, outp, size;
} Queue;

Queue *new(int n) {
  int *buff = malloc(n*sizeof(int));
  Queue q = {buff,0,0,n};
  Queue *qptr = malloc(sizeof(Queue));
  *qptr = q;
  return qptr;
}

void put(Queue *q, int n) {
  q->buf[q->inp] = n;
  q->inp = (q->inp + 1) % q->size;
}

int get(Queue *q) {
  int ans = q->buf[q->outp];
  q->outp = (q->outp + 1) % q->size;
  return ans;
}

int size(Queue *q) {
  return (q->inp - q->outp) % q->size;
}

data FQueue = FQueue
  { fqElems :: [Int]
  , fqSize  :: Int
  }

type State = Map (Var Queue) FQueue

emptyState :: State
emptyState = Map.empty

type FakeOp a = State -> Either Err (State, a)

fNew :: Int -> FakeOp (Var Queue)
fNew sz s =
  let
    v  = Var (Map.size s)
    s' = Map.insert v (FQueue [] sz) s
  in
    return (s', v)

data Err
  = QueueDoesNotExist
  | QueueIsEmpty
  deriving (Eq, Show)


fPut :: Var Queue -> Int -> FakeOp ()
fPut q i s
  | q `Map.notMember` s = Left QueueDoesNotExist
  | otherwise = return (Map.adjust (\fq -> fq { fqElems = fqElems fq ++ [i] }) q s, ())

fGet :: Var Queue -> FakeOp Int
fGet q s
  | q `Map.notMember` s = Left QueueDoesNotExist
  | otherwise = case fqElems (s Map.! q) of
      []     -> Left QueueIsEmpty
      i : is -> return (Map.adjust (\fq -> fq { fqElems = is }) q s, i)

fSize :: Var Queue -> FakeOp Int
fSize q s
  | q `Map.notMember` s = Left QueueDoesNotExist
  | otherwise           = return (s, length (fqElems (s Map.! q)))

instance StateModel State where

  initialState = Map.empty

  type Reference State = Queue

  type PreconditionFailure State = Err

  data Command State q
    = New Int
    | Put q Int
    | Get q
    | Size q
    deriving (Show, Functor, Foldable)

  data Response State q
    = New_ q
    | Put_ ()
    | Get_ Int
    | Size_ Int
    deriving (Eq, Show, Functor, Foldable)

  generateCommand s
    | Map.null s = New . getPositive <$> arbitrary
    | otherwise  = oneof
      [ New . getPositive <$> arbitrary
      , Put  <$> arbitraryQueue <*> arbitrary
      , Get  <$> arbitraryQueue
      ]
    where
      arbitraryQueue :: Gen (Var Queue)
      arbitraryQueue = elements (Map.keys s)

  shrinkCommand _s (New i)   = [ New i'   | Positive i' <- shrink (Positive i) ]
  shrinkCommand _s (Put q i) = [ Put q i' | i' <- shrink i ]
  shrinkCommand _s _cmd = []

  runFake (New sz)  s = fmap New_  <$> fNew sz s
  runFake (Put q i) s = fmap Put_  <$> fPut q i s
  runFake (Get q)   s = fmap Get_  <$> fGet q s
  runFake (Size q)  s = fmap Size_ <$> fSize q s

  -- Here `new`, `put`, `get` and `size` are FFI wrappers for their respective C
  -- functions.
  runReal (New sz)  = New_  <$> new sz
  runReal (Put q i) = Put_  <$> put q i
  runReal (Get q)   = Get_  <$> get q
  runReal (Size q)  = Size_ <$> size q

prop_queue :: Commands State -> Property
prop_queue cmds = monadicIO $ do
  runCommands cmds
  assert True

 data Err
   = QueueDoesNotExist
   | QueueIsEmpty
+  | QueueIsFull

fPut :: Var Queue -> Int -> State -> Either Err (State, ())
fPut q i s
  | q `Map.notMember` s = Left QueueDoesNotExist
+ | length (fqElems (s Map.! q)) >= fqSize (s Map.! q) = Left QueueIsFull
  | otherwise = return (Map.adjust (\fq -> fq { fqElems = fqElems fq ++ [i] }) q s, ())

unit_queueFull :: IO ()
unit_queueFull = quickCheck (withMaxSuccess 1 (expectFailure (prop_queue cmds)))
  where
    cmds = Commands
      [ New 1
      , Put (Var 0) 1
      , Put (Var 0) 0
      , Get (Var 0)
      ]

  generateCommand s
    | Map.null s = New . getPositive <$> arbitrary
    | otherwise  = oneof
      [ New . getPositive <$> arbitrary
      , Put  <$> arbitraryQueue <*> arbitrary
      , Get  <$> arbitraryQueue
+     , Size <$> arbitraryQueue
      ]

void put(Queue *q, int n) {
  q->buf[q->inp] = n;
  q->inp = (q->inp + 1) % q->size;
}

int size(Queue *q) {
  return (q->inp - q->outp) % q->size;
}

  Queue *new(int n) {
-   int *buff = malloc(n*sizeof(int));
-   Queue q = {buff,0,0,n};
+   int *buff = malloc((n + 1)*sizeof(int));
+   Queue q = {buff,0,0,n + 1};
    Queue *qptr = malloc(sizeof(Queue));
    *qptr = q;
    return qptr;
}

unit_queueSize :: IO ()
unit_queueSize = quickCheck (withMaxSuccess 1 (prop_queue cmds))
  where
    cmds = Commands
      [ New 1
      , Put (Var 0) 0
      , Size (Var 0)
      ]

  int size(Queue *q) {
-   return (q->inp - q->outp) % q->size;
+   return abs(q->inp - q->outp) % q->size;
  }

  int size(Queue *q) {
-   return abs(q->inp - q->outp) % q->size;
+   return (q->inp - q->outp + q->size) % q->size;
  }

newtype ParallelCommands state = ParallelCommands [Fork state]

newtype Fork state = Fork [Command state (Var (Reference state))]

class (StateModel state, Ord state) => ParallelModel state where

  generateCommandParallel :: [state] -> Gen (Command state (Var (Reference state)))
  generateCommandParallel ss = do
    s <- elements ss
    generateCommand s

  shrinkCommandParallel :: [state] -> Command state (Var (Reference state))
                        -> [Command state (Var (Reference state))]
  shrinkCommandParallel ss cmd = shrinkCommand (maximum ss) cmd

  arbitrary :: Gen (ParallelCommands state)
  arbitrary = ParallelCommands <$> go [initialState]
    where
      go :: [state] -> Gen [Fork state]
      go ss = sized $ \n ->
        let
          w = n `div` 2 + 1
        in
          frequency
            [ (1, return [])
            , (w, do k <- frequency [ (50, return 1) -- 50% single threaded
                                    , (30, return 2) -- 30% double threaded
                                    , (20, return 3) -- 20% triple threaded
                                    ]
                     mCmds <- vectorOf k (generateCommandParallel ss)
                                `suchThatMaybe` (parallelSafe ss . Fork)
                     case mCmds of
                       Nothing   -> return []
                       Just cmds ->
                         (Fork cmds :) <$> go (nextStates ss cmds))
            ]

nextStates :: (StateModel state, Ord state)
           => [state] -> [Command state (Var (Reference state))] -> [state]
nextStates ss cmds = nubOrd [ foldl' nextState s cmds | s <- ss ]

parallelSafe :: ParallelModel state => [state] -> Fork state -> Bool
parallelSafe ss (Fork cmds0) = and
  [ preconditionsHold s cmds | s <- toList ss, cmds <- permutations cmds0 ]
  where
    preconditionsHold s0 = all (go s0) . permutations
      where
        go _s [] = True
        go  s (cmd : cmds)
          | precondition s cmd = go (nextState s cmd) cmds
          | otherwise          = False

  shrink :: ParallelCommands state -> [ParallelCommands state]
  shrink = pruneShrinks . possibleShrinks
    where
      possibleShrinks :: ParallelCommands state -> [ParallelCommands state]
      possibleShrinks
        = map (coerce . map (map fst))
        . shrinkList (shrinkList shrinker) . withParStates . unParallelCommands
        where
          withParStates :: (StateModel state, Ord state)
                        => [Fork state]
                        -> [[(Command state (Var (Reference state)), [state])]]
          withParStates = go [initialState] [] . coerce
            where
              go _ss acc []             = reverse acc
              go  ss acc (cmds : cmdss) =
                go (nextStates ss cmds) (map (\cmd -> (cmd, ss)) cmds : acc) cmdss

          shrinker :: (Command state (Var (Reference state)), [state])
                   -> [(Command state (Var (Reference state)), [state])]
          shrinker (cmd, ss) = [ (cmd', ss) | cmd' <- shrinkCommandParallel ss cmd ]

      pruneShrinks :: [ParallelCommands state] -> [ParallelCommands state]
      pruneShrinks = coerce . filter (not . null)
                   . map (go [initialState] Set.empty [] . unParallelCommands)
        where
          go :: [state] -> Set (Var (Reference state)) -> [Fork state] -> [Fork state] -> [Fork state]
          go _ss _vars acc [] = reverse acc
          go  ss  vars acc (fork@(Fork cmds) : forks)
            | all (scopeCheck vars) cmds
            , parallelSafe ss fork =
              let
                ss'   = nextStates ss cmds
                vars' = getReturnedVars (head ss) vars cmds -- NOTE: head is safe
              in
                go ss' vars' (fork : acc) forks
            | otherwise            = go ss vars acc forks

          -- It doesn't matter which of the possible states we start in, as all
          -- commands in a fork pass their preconditions in all states. It also
          -- doesn't matter in which interleaving we gather the responses, as
          -- all we do is collect the `Var`s that get returned into an unordered
          -- `Set`.
          getReturnedVars _s vars [] = vars
          getReturnedVars s vars (cmd : cmds) = case runFake cmd s of
            Left _preconditionFailed ->
              error "getReturnedVars: impossible, parallelSafe checks that all preconditions hold"
            Right (_s', resp) ->
              getReturnedVars s (vars `Set.union` Set.fromList (toList resp)) cmds

  -- If another command monad is used we need to provide a way run it inside the
  -- IO monad. This is only needed for parallel testing, because IO is the only
  -- monad we can execute on different threads.
  runCommandMonad :: proxy state -> CommandMonad state a -> IO a

newtype History state = History [Event state]
deriving stock instance
   (Show (Command state (Var (Reference state))),
    Show (Response state (Reference state))) => Show (History state)

data Event state
  = Invoke Pid (Command state (Var (Reference state)))
  | Ok     Pid (Response state (Reference state))
deriving stock instance
  (Show (Command state (Var (Reference state))),
   Show (Response state (Reference state))) => Show (Event state)

newtype Pid = Pid Int
  deriving stock (Eq, Ord, Show)
  deriving newtype Enum

runParallelCommands :: forall state. ParallelModel state
                    => ParallelCommands state -> PropertyM IO ()
runParallelCommands cmds0@(ParallelCommands forks0) = do
  forM_ (parallelCommands cmds0) $ \cmd -> do
    let name = commandName cmd
    monitor (tabulate "Commands" [name] . classify True name)
  monitor (tabulate "Concurrency" (map (show . length . unFork) forks0))
  q   <- liftIO newTQueueIO
  c   <- liftIO newAtomicCounter
  env <- liftIO (runForks q c emptyEnv forks0)
  hist <- History <$> liftIO (atomically (flushTQueue q))
  let ok = linearisable env (interleavings hist)
  unless ok (monitor (counterexample (show hist)))
  assert ok
  where
    runForks :: TQueue (Event state) -> AtomicCounter -> Env state -> [Fork state]
             -> IO (Env state)
    runForks _q _c env [] = return env
    runForks  q  c env (Fork cmds : forks) = do
      envs <- liftIO $
        mapConcurrently (runParallelReal q c env) (zip [Pid 0..] cmds)
      let env' = combineEnvs (env : envs)
      runForks q c env' forks

    runParallelReal :: TQueue (Event state) -> AtomicCounter -> Env state
                    -> (Pid, Command state (Var (Reference state))) -> IO (Env state)
    runParallelReal q c env (pid, cmd) = do
      atomically (writeTQueue q (Invoke pid cmd))
      eResp <- try (runCommandMonad (Proxy :: Proxy state) (runReal (fmap (lookupEnv env) cmd)))
      case eResp of
        Left (err :: SomeException) ->
          error ("runParallelReal: " ++ displayException err)
        Right resp -> do
          -- NOTE: It's important that we extend the environment before writing `Ok`
          -- to the history, otherwise we might get scope issues.
          env' <- extendEnvParallel env c (toList resp)
          atomically (writeTQueue q (Ok pid resp))
          return env'

newtype AtomicCounter = AtomicCounter (IORef Int)

newAtomicCounter :: IO AtomicCounter
newAtomicCounter = AtomicCounter <$> newIORef 0

-- Returns old value.
incrAtomicCounter :: AtomicCounter -> Int -> IO Int
incrAtomicCounter (AtomicCounter ioRef) n =
  atomicModifyIORef' ioRef (\old -> (old + n, old))

extendEnvParallel :: Env state -> AtomicCounter -> [Reference state] -> IO (Env state)
extendEnvParallel env c refs = do
  i <- incrAtomicCounter c (length refs)
  return (extendEnv env (zip [i..] refs))

combineEnvs :: [Env state] -> Env state
combineEnvs = Env . IntMap.unions . map unEnv

data Op state = Op (Command state (Var (Reference state)))
                   (Response state (Reference state))

interleavings :: History state -> Forest (Op state)
interleavings (History [])  = []
interleavings (History evs0) =
  [ Node (Op cmd resp) (interleavings (History evs'))
  | (tid, cmd)   <- takeInvocations evs0
  , (resp, evs') <- findResponse tid
                      (filter1 (not . matchInvocation tid) evs0)
  ]
  where
    takeInvocations :: [Event state] -> [(Pid, Command state (Var (Reference state)))]
    takeInvocations []                         = []
    takeInvocations ((Invoke pid cmd)   : evs) = (pid, cmd) : takeInvocations evs
    takeInvocations ((Ok    _pid _resp) : _)   = []

    findResponse :: Pid -> [Event state] -> [(Response state (Reference state), [Event state])]
    findResponse _pid []                                   = []
    findResponse  pid ((Ok pid' resp) : evs) | pid == pid' = [(resp, evs)]
    findResponse  pid (ev             : evs)               =
      [ (resp, ev : evs') | (resp, evs') <- findResponse pid evs ]

    matchInvocation :: Pid -> Event state -> Bool
    matchInvocation pid (Invoke pid' _cmd) = pid == pid'
    matchInvocation _   _                  = False

    filter1 :: (a -> Bool) -> [a] -> [a]
    filter1 _ []                   = []
    filter1 p (x : xs) | p x       = x : filter1 p xs
                       | otherwise = xs

linearisable :: forall state. StateModel state
             => Env state -> Forest (Op state) -> Bool
linearisable env = any' (go initialState)
  where
    go :: state -> Tree (Op state) -> Bool
    go s (Node (Op cmd cresp) ts) =
      case runFake cmd s of
        Left _preconditionFailure ->
          error "linearisable: impossible, all precondtions are satisifed during generation"
        Right (s', resp) ->
          cresp == fmap (lookupEnv env) resp && any' (go s') ts

    any' :: (a -> Bool) -> [a] -> Bool
    any' _p [] = True
    any'  p xs = any p xs

instance ParallelModel Counter where

  -- The command monad is IO, so we don't need to do anything here.
  runCommandMonad _ = id

prop_parallelCounter :: ParallelCommands Counter -> Property
prop_parallelCounter cmds = monadicIO $ do
  replicateM_ 10 $ do
    run reset
    runParallelCommands cmds
  assert True

  runReal Incr = Incr_ <$> incrRaceCondition

incrRaceCondition :: IO ()
incrRaceCondition = do
  n <- readIORef gLOBAL_COUNTER
  writeIORef gLOBAL_COUNTER (n + 1)

incrRaceCondition :: IO ()
incrRaceCondition = do
  n <- readIORef gLOBAL_COUNTER
  threadDelay 100
  writeIORef gLOBAL_COUNTER (n + 1)
  threadDelay 100

module SleepyIORef (module SleepyIORef, IORef) where

import Control.Concurrent (threadDelay)
import Data.IORef (IORef)
import qualified Data.IORef as IORef

readIORef :: IORef a -> IO a
readIORef ref = do
  threadDelay 1000
  IORef.readIORef ref

{-# NOINLINE registry #-}
registry :: IORef [(String,ThreadId)]
registry = unsafePerformIO (newIORef [])

alive :: ThreadId -> IO Bool
alive tid = do
  s <- threadStatus tid
  return $ s /= ThreadFinished && s /= ThreadDied

spawn :: IO ThreadId
spawn = forkIO (threadDelay 100000000)

whereis :: String -> IO (Maybe ThreadId)
whereis name = do
  reg <- readRegistry
  return $ lookup name reg

register :: String -> ThreadId -> IO ()
register name tid = do
  ok <- alive tid
  reg <- readRegistry
  if ok && name `notElem` map fst reg && tid `notElem` map snd reg
    then do
      atomicModifyIORef registry $ \reg' ->
           if name `notElem` map fst reg' && tid `notElem` map snd reg'
             then ((name,tid):reg',())
             else (reg',badarg)
    else badarg

unregister :: String -> IO ()
unregister name = do
  reg <- readRegistry
  if name `elem` map fst reg
    then atomicModifyIORef registry $ \reg' ->
           (filter ((/=name).fst) reg',
            ())
    else badarg

readRegistry :: IO [(String, ThreadId)]
readRegistry = do
  reg <- readIORef registry
  garbage <- filterM (fmap not.alive) (map snd reg)
  atomicModifyIORef' registry $ \reg' ->
    let reg'' = filter ((`notElem` garbage).snd) reg' in (reg'',reg'')

badarg :: a
badarg = error "bad argument"

kill :: ThreadId -> IO ()
kill tid = do
  killThread tid
  waitUntilDead 1000
  where
    waitUntilDead :: Int -> IO ()
    waitUntilDead 0 = error "kill: thread didn't die"
    waitUntilDead n = do
      b <- alive tid
      if b
      then do
        threadDelay 1000
        waitUntilDead (n - 1)
      else return ()

data RegState = RegState
  { tids   :: [Var ThreadId]
  , regs   :: [(String, Var (ThreadId))]
  , killed :: [Var ThreadId]
  }
  deriving (Eq, Ord, Show)

instance StateModel RegState where

  initialState :: RegState
  initialState = RegState [] [] []

  type Reference RegState = ThreadId

  data Command RegState tid
    = Spawn
    | WhereIs String
    | Register String tid
    | Unregister String
    | Kill tid
    deriving (Show, Functor, Foldable)

  data Response RegState tid
    = Spawn_ tid
    | WhereIs_ (NonFoldable (Maybe tid))
    | Register_ (Either ErrorCall ())
    | Unregister_ (Either ErrorCall ())
    | Kill_ ()
    deriving (Eq, Show, Functor, Foldable)

  generateCommand :: RegState -> Gen (Command RegState (Var ThreadId))
  generateCommand s = oneof $
    [ return Spawn ] ++
    [ Register <$> arbitraryName <*> elements (tids s) | not (null (tids s)) ] ++
    [ Unregister <$> arbitraryName
    , WhereIs <$> arbitraryName
    ] ++
    [ Kill <$> elements (tids s) | not (null (tids s)) ]
    where
      arbitraryName :: Gen String
      arbitraryName = elements allNames

  runFake :: Command RegState (Var ThreadId)-> RegState
          -> Either void (RegState, Response RegState (Var ThreadId))
  runFake Spawn               s = let tid = Var (length (tids s)) in
                                  return (s { tids = tids s ++ [tid] }, Spawn_ tid)
  runFake (WhereIs name)      s = return (s, WhereIs_ (NonFoldable (lookup name (regs s))))
  runFake (Register name tid) s
    | tid `elem` tids s
    , name `notElem` map fst (regs s)
    , tid `notElem` map snd (regs s)
    , tid `notElem` killed s
    = return (s { regs = (name, tid) : regs s }, Register_ (Right ()))

    | otherwise
    = return (s, Register_ (Left (ErrorCall "bad argument")))

  runFake (Unregister name)   s
    | name `elem` map fst (regs s) =
        return (s { regs = remove name (regs s) }, Unregister_ (Right ()))
    | otherwise = return (s, Unregister_ (Left (ErrorCall "bad argument")))
    where
      remove x = filter ((/= x) . fst)
  runFake (Kill tid) s = return (s { killed = tid : killed s
                                   , regs   = remove tid (regs s)}, Kill_ ())
    where
      remove x = filter ((/= x) . snd)

  runReal :: Command RegState ThreadId -> IO (Response RegState ThreadId)
  runReal Spawn               = Spawn_      <$> spawn
  runReal (WhereIs name)      = WhereIs_ . NonFoldable <$> whereis name
  runReal (Register name tid) = Register_   <$> fmap (left abstractError) (try (register name tid))
  runReal (Unregister name)   = Unregister_ <$> fmap (left abstractError) (try (unregister name))
  runReal (Kill tid)          = Kill_       <$> kill tid

  monitoring :: (RegState, RegState) -> Command RegState ThreadId
             -> Response RegState ThreadId -> Property -> Property
  monitoring (_s, s') cmd resp =
    let
      aux tag = classify True (show tag)
              . counterexample ("\n    State: " ++ show s' ++ "\n")
    in
      case (cmd, resp) of
        (Register   {}, Register_   (Left _))  -> aux RegisterFailed
        (Register   {}, Register_   (Right _)) -> aux RegisterSucceeded
        (Unregister {}, Unregister_ (Left _))  -> aux UnregisterFailed
        (Unregister {}, Unregister_ (Right _)) -> aux UnregisterSucceeded
        _otherwise -> counterexample $ "\n    State: " ++ show s' ++ "\n"

-- Throws away the location information from the error, so that it matches up
-- with the fake.
abstractError :: ErrorCall -> ErrorCall
abstractError (ErrorCallWithLocation msg _loc) = ErrorCall msg

allNames :: [String]
allNames = ["a", "b", "c", "d", "e"]

data Tag = RegisterFailed | RegisterSucceeded | UnregisterFailed | UnregisterSucceeded
  deriving Show

newtype NonFoldable a = NonFoldable a
  deriving stock (Eq, Show)

instance Functor NonFoldable where
  fmap f (NonFoldable x) = NonFoldable (f x)

instance Foldable NonFoldable where
  foldMap _f (NonFoldable _x) = mempty

prop_registry :: Commands RegState -> Property
prop_registry cmds = monadicIO $ do
  void (run cleanUp)
  runCommands cmds
  assert True

cleanUp :: IO [Either ErrorCall ()]
cleanUp = sequence
  [ try (unregister name) :: IO (Either ErrorCall ())
  | name <- allNames
  ]

  register :: String -> ThreadId -> IO ()
  register name tid = do
    ok <- alive tid
    reg <- readRegistry
    if ok && name `notElem` map fst reg && tid `notElem` map snd reg
      then atomicModifyIORef' registry $ \reg' ->
             if name `notElem` map fst reg' && tid `notElem` map snd reg'
-              then ((name,tid):reg',())
+              then ([(name,tid)],())
               else (reg',badarg)
      else badarg

instance ParallelModel RegState where
  runCommandMonad _ = id

prop_parallelRegistry :: ParallelCommands RegState -> Property
prop_parallelRegistry cmds = monadicIO $ do
  replicateM_ 10 $ do
    void (run cleanUp)
    runParallelCommands cmds
  assert True

- import Data.IORef
+ import SleepyIORef

{-# NOINLINE lock #-}
lock :: MVar ()
lock = unsafePerformIO (newMVar ())

registerNoRace :: String -> ThreadId -> IO ()
registerNoRace name tid = withMVar lock $ \_ -> register name tid