Internal.hs

module Test.QuickCheck.DynamicLogic.Internal where

import Control.Applicative
import Control.Arrow (second)
import Control.Monad
import Data.Typeable
import Test.QuickCheck (Gen, Property, Testable)
import Test.QuickCheck qualified as QC
import Test.QuickCheck.DynamicLogic.CanGenerate
import Test.QuickCheck.DynamicLogic.Quantify
import Test.QuickCheck.DynamicLogic.SmartShrinking
import Test.QuickCheck.DynamicLogic.Utils qualified as QC
import Test.QuickCheck.StateModel

-- | A `DynFormula` may depend on the QuickCheck size parameter
newtype DynFormula s = DynFormula {unDynFormula :: Int -> DynLogic s}

-- | Base Dynamic logic formulae language.
-- Formulae are parameterised
-- over the type of state `s` to which they apply. A `DynLogic` value
-- cannot be constructed directly, one has to use the various "smart
-- constructors" provided, see the /Building formulae/ section.
data DynLogic s
  = -- | False
    EmptySpec
  | -- | True
    Stop
  | -- | After any action the predicate should hold
    AfterAny (DynPred s)
  | -- | Choice (angelic or demonic)
    Alt ChoiceType (DynLogic s) (DynLogic s)
  | -- | Prefer this branch if trying to stop.
    Stopping (DynLogic s)
  | -- | After a specific action the predicate should hold
    forall a.
    (Eq (Action s a), Show (Action s a), Typeable a) =>
    After (ActionWithPolarity s a) (Var a -> DynPred s)
  | Error String (DynPred s)
  | -- | Adjust the probability of picking a branch
    Weight Double (DynLogic s)
  | -- | Generating a random value
    forall a.
    QuantifyConstraints a =>
    ForAll (Quantification a) (a -> DynLogic s)
  | -- | Apply a QuickCheck property modifier (like `tabulate` or `collect`)
    Monitor (Property -> Property) (DynLogic s)

data ChoiceType = Angelic | Demonic
  deriving (Eq, Show)

type DynPred s = Annotated s -> DynLogic s

-- * Building formulae

-- | Ignore this formula, i.e. backtrack and try something else. @forAllScripts ignore (const True)@
--   will discard all test cases (equivalent to @False ==> True@).
ignore :: DynFormula s
ignore = DynFormula . const $ EmptySpec

-- | `True` for DL formulae.
passTest :: DynFormula s
passTest = DynFormula . const $ Stop

-- | Given `f` must be `True` given /any/ state.
afterAny :: (Annotated s -> DynFormula s) -> DynFormula s
afterAny f = DynFormula $ \n -> AfterAny $ \s -> unDynFormula (f s) n

afterPolar
  :: (Typeable a, Eq (Action s a), Show (Action s a))
  => ActionWithPolarity s a
  -> (Var a -> Annotated s -> DynFormula s)
  -> DynFormula s
afterPolar act f = DynFormula $ \n -> After act $ \x s -> unDynFormula (f x s) n

-- | Given `f` must be `True` after /some/ action.
-- `f` is passed the state resulting from executing the `Action`.
after
  :: (Typeable a, Eq (Action s a), Show (Action s a))
  => Action s a
  -> (Var a -> Annotated s -> DynFormula s)
  -> DynFormula s
after act f = afterPolar (ActionWithPolarity act PosPolarity) f

-- | Given `f` must be `True` after /some/ negative action.
-- `f` is passed the state resulting from executing the `Action`
-- as a negative action.
afterNegative
  :: (Typeable a, Eq (Action s a), Show (Action s a))
  => Action s a
  -> (Annotated s -> DynFormula s)
  -> DynFormula s
afterNegative act f = afterPolar (ActionWithPolarity act NegPolarity) (const f)

-- | Disjunction for DL formulae.
-- Is `True` if either formula is `True`. The choice is /angelic/, ie. it is
-- always made by the "caller". This is  mostly important in case a test is
-- `Stuck`.
(|||) :: DynFormula s -> DynFormula s -> DynFormula s
-- In formulae, we use only angelic choice. But it becomes demonic
-- after one step (that is, the choice has been made).
DynFormula f ||| DynFormula g = DynFormula $ \n -> Alt Angelic (f n) (g n)

-- | First-order quantification of variables.
-- Formula @f@ is `True` iff. it is `True` /for all/ possible values of `q`. The
-- underlying framework will generate values of `q` and check the formula holds
-- for those values. `Quantifiable` values are thus values that can be generated
-- and checked and the `Test.QuickCheck.DynamicLogic.Quantify` module defines
-- basic combinators to build those from building blocks.
forAllQ
  :: Quantifiable q
  => q
  -> (Quantifies q -> DynFormula s)
  -> DynFormula s
forAllQ q f
  | isEmptyQ q' = ignore
  | otherwise = DynFormula $ \n -> ForAll q' $ ($ n) . unDynFormula . f
  where
    q' = quantify q

-- | Adjust weight for selecting formula.
-- This is mostly useful in relation with `(|||)` combinator, in order to tweak the
-- priority for generating the next step(s) of the test that matches the formula.
weight :: Double -> DynFormula s -> DynFormula s
weight w f = DynFormula $ Weight w . unDynFormula f

-- | Get the current QuickCheck size parameter.
withSize :: (Int -> DynFormula s) -> DynFormula s
withSize f = DynFormula $ \n -> unDynFormula (f n) n

-- | Prioritise doing this if we are
-- trying to stop generation.
toStop :: DynFormula s -> DynFormula s
toStop (DynFormula f) = DynFormula $ Stopping . f

-- | Successfully ends the test.
done :: Annotated s -> DynFormula s
done _ = passTest

-- | Ends test with given error message.
errorDL :: String -> DynFormula s
errorDL s = DynFormula . const $ Error s (const EmptySpec)

-- | Embed QuickCheck's monitoring functions (eg. `label`, `tabulate`) in
-- a formula.
-- This is useful to improve the reporting from test execution, esp. in the
-- case of failures.
monitorDL :: (Property -> Property) -> DynFormula s -> DynFormula s
monitorDL m (DynFormula f) = DynFormula $ Monitor m . f

-- | Formula should hold at any state.
-- In effect this leads to exploring alternatives from a given state `s` and ensuring
-- formula holds in all those states.
always :: (Annotated s -> DynFormula s) -> (Annotated s -> DynFormula s)
always p s = withSize $ \n -> toStop (p s) ||| p s ||| weight (fromIntegral n) (afterAny (always p))

data FailingAction s
  = ErrorFail String
  | forall a. (Typeable a, Eq (Action s a)) => ActionFail (ActionWithPolarity s a)

instance StateModel s => HasVariables (FailingAction s) where
  getAllVariables ErrorFail{} = mempty
  getAllVariables (ActionFail a) = getAllVariables a

instance StateModel s => Eq (FailingAction s) where
  ErrorFail s == ErrorFail s' = s == s'
  ActionFail (a :: ActionWithPolarity s a) == ActionFail (a' :: ActionWithPolarity s' a')
    | Just Refl <- eqT @a @a' = a == a'
  _ == _ = False

instance StateModel s => Show (FailingAction s) where
  show (ErrorFail s) = "Error " ++ show s
  show (ActionFail (ActionWithPolarity a pol)) = show pol ++ " : " ++ show a

data DynLogicTest s
  = BadPrecondition (TestSequence s) (FailingAction s) (Annotated s)
  | Looping (TestSequence s)
  | Stuck (TestSequence s) (Annotated s)
  | DLScript (TestSequence s)

data Witnesses r where
  Do :: r -> Witnesses r
  Witness :: QuantifyConstraints a => a -> Witnesses r -> Witnesses r

discardWitnesses :: Witnesses r -> r
discardWitnesses (Do r) = r
discardWitnesses (Witness _ k) = discardWitnesses k

pattern Witnesses :: Witnesses () -> r -> Witnesses r
pattern Witnesses w r <- ((\wit -> (void wit, discardWitnesses wit)) -> (w, r))
  where
    Witnesses w r = r <$ w

{-# COMPLETE Witnesses #-}

deriving instance Functor Witnesses
deriving instance Foldable Witnesses
deriving instance Traversable Witnesses

instance Eq r => Eq (Witnesses r) where
  Do r == Do r' = r == r'
  Witness (a :: a) k == Witness (a' :: a') k' =
    case eqT @a @a' of
      Just Refl -> a == a' && k == k'
      Nothing -> False
  _ == _ = False

instance Show r => Show (Witnesses r) where
  show (Do r) = "Do $ " ++ show r
  show (Witness a k) = "Witness (" ++ show a ++ " :: " ++ show (typeOf a) ++ ")\n" ++ show k -- TODO

type TestStep s = Witnesses (Step s)

newtype TestSequence s = TestSeq (Witnesses (TestContinuation s))

deriving instance StateModel s => Show (TestSequence s)
deriving instance StateModel s => Eq (TestSequence s)

data TestContinuation s
  = ContStep (Step s) (TestSequence s)
  | ContStop

pattern TestSeqStop :: TestSequence s
pattern TestSeqStop = TestSeq (Do ContStop)

pattern TestSeqStep :: Step s -> TestSequence s -> TestSequence s
pattern TestSeqStep s ss = TestSeq (Do (ContStep s ss))

-- The `()` are the constraints required to use the pattern, and the `(Typeable a, ...)` are the
-- ones provided when you do (including a fresh type variable `a`).
-- See https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/pattern_synonyms.html#typing-of-pattern-synonyms
pattern TestSeqWitness :: () => forall a. QuantifyConstraints a => a -> TestSequence s -> TestSequence s
pattern TestSeqWitness a ss <- TestSeq (Witness a (TestSeq -> ss))
  where
    TestSeqWitness a (TestSeq ss) = TestSeq (Witness a ss)

{-# COMPLETE TestSeqWitness, TestSeqStep, TestSeqStop #-}

deriving instance StateModel s => Show (TestContinuation s)
deriving instance StateModel s => Eq (TestContinuation s)

consSeq :: TestStep s -> TestSequence s -> TestSequence s
consSeq step ss = TestSeq $ flip ContStep ss <$> step

unconsSeq :: TestSequence s -> Maybe (TestStep s, TestSequence s)
unconsSeq (TestSeq ss) =
  case discardWitnesses ss of
    ContStop -> Nothing
    ContStep s rest -> Just (s <$ ss, rest)

unstopSeq :: TestSequence s -> Maybe (Witnesses ())
unstopSeq (TestSeq ss) =
  case discardWitnesses ss of
    ContStop -> Just $ () <$ ss
    ContStep{} -> Nothing

pattern TestSeqStopW :: Witnesses () -> TestSequence s
pattern TestSeqStopW w <- (unstopSeq -> Just w)
  where
    TestSeqStopW w = TestSeq (ContStop <$ w)

pattern (:>) :: TestStep s -> TestSequence s -> TestSequence s
pattern step :> ss <- (unconsSeq -> Just (step, ss))
  where
    step :> ss = consSeq step ss

{-# COMPLETE TestSeqStopW, (:>) #-}

nullSeq :: TestSequence s -> Bool
nullSeq TestSeqStop = True
nullSeq _ = False

dropSeq :: Int -> TestSequence s -> TestSequence s
dropSeq n _ | n < 0 = error "dropSeq: negative number"
dropSeq 0 ss = ss
dropSeq _ TestSeqStop = TestSeqStop
dropSeq n (TestSeqWitness _ ss) = dropSeq (n - 1) ss
dropSeq n (TestSeqStep _ ss) = dropSeq (n - 1) ss

getContinuation :: TestSequence s -> TestSequence s
getContinuation (TestSeq w) = case discardWitnesses w of
  ContStop -> TestSeqStop
  ContStep _ s -> s

unlines' :: [String] -> String
unlines' [] = ""
unlines' xs = init $ unlines xs

prettyTestSequence :: VarContext -> TestSequence s -> String
prettyTestSequence ctx ss = unlines' $ zipWith (++) ("do " : repeat "   ") $ prettySeq ss
  where
    prettySeq (TestSeqStopW w) = prettyWitnesses w
    prettySeq (Witnesses w step :> ss') = prettyWitnesses w ++ show (WithUsedVars ctx step) : prettySeq ss'

prettyWitnesses :: Witnesses () -> [String]
prettyWitnesses (Witness a w) = ("_ <- forAllQ $ exactlyQ $ " ++ show a) : prettyWitnesses w
prettyWitnesses Do{} = []

instance StateModel s => Show (DynLogicTest s) where
  show (BadPrecondition ss bad s) =
    prettyTestSequence (usedVariables ss <> allVariables bad) ss
      ++ "\n   -- In state: "
      ++ show s
      ++ "\n   "
      ++ prettyBad bad
    where
      prettyBad :: FailingAction s -> String
      prettyBad (ErrorFail e) = "assert " ++ show e ++ " False"
      prettyBad (ActionFail (ActionWithPolarity a p)) = f ++ " $ " ++ show a ++ "  -- Failed precondition\n   pure ()"
        where
          f
            | p == PosPolarity = "action"
            | otherwise = "failingAction"
  show (Looping ss) = prettyTestSequence (usedVariables ss) ss ++ "\n   pure ()\n   -- Looping"
  show (Stuck ss s) = prettyTestSequence (usedVariables ss) ss ++ "\n   pure ()\n   -- Stuck in state " ++ show s
  show (DLScript ss) = prettyTestSequence (usedVariables ss) ss ++ "\n   pure ()\n"

usedVariables :: forall s. StateModel s => TestSequence s -> VarContext
usedVariables = go initialAnnotatedState
  where
    go :: Annotated s -> TestSequence s -> VarContext
    go aState TestSeqStop = allVariables (underlyingState aState)
    go aState (TestSeqWitness a ss) = allVariables a <> go aState ss
    go aState (TestSeqStep step@(_ := act) ss) =
      allVariables act
        <> allVariables (underlyingState aState)
        <> go (nextStateStep step aState) ss

-- | Restricted calls are not generated by "AfterAny"; they are included
-- in tests explicitly using "After" in order to check specific
-- properties at controlled times, so they are likely to fail if
-- invoked at other times.
class StateModel s => DynLogicModel s where
  restricted :: Action s a -> Bool
  restricted _ = False

restrictedPolar :: DynLogicModel s => ActionWithPolarity s a -> Bool
restrictedPolar (ActionWithPolarity a _) = restricted a

-- * Generate Properties

-- | Simplest "execution" function for `DynFormula`.
-- Turns a given a `DynFormula` paired with an interpreter function to produce some result from an
-- `Actions` sequence into a proper `Property` that can be run by QuickCheck.
forAllScripts
  :: (DynLogicModel s, Testable a)
  => DynFormula s
  -> (Actions s -> a)
  -> Property
forAllScripts = forAllMappedScripts id id

-- | `Property` function suitable for formulae without choice.
forAllUniqueScripts
  :: (DynLogicModel s, Testable a)
  => Annotated s
  -> DynFormula s
  -> (Actions s -> a)
  -> Property
forAllUniqueScripts s f k =
  QC.withSize $ \sz ->
    let d = unDynFormula f sz
        n = unsafeNextVarIndex $ vars s
     in case generate chooseUniqueNextStep d n s 500 of
          Nothing -> QC.counterexample "Generating Non-unique script in forAllUniqueScripts" False
          Just test -> validDLTest test . applyMonitoring d test . QC.property $ k (scriptFromDL test)

-- | Creates a `Property` from `DynFormula` with some specialised isomorphism for shrinking purpose.
forAllMappedScripts
  :: (DynLogicModel s, Testable a)
  => (rep -> DynLogicTest s)
  -> (DynLogicTest s -> rep)
  -> DynFormula s
  -> (Actions s -> a)
  -> Property
forAllMappedScripts to from f k =
  QC.withSize $ \n ->
    let d = unDynFormula f n
     in QC.forAllShrinkBlind
          (QC.Smart 0 <$> QC.sized ((from <$>) . generateDLTest d))
          (shrinkSmart ((from <$>) . shrinkDLTest d . to))
          $ \(QC.Smart _ script) ->
            withDLScript d k (to script)

withDLScript :: (DynLogicModel s, Testable a) => DynLogic s -> (Actions s -> a) -> DynLogicTest s -> Property
withDLScript d k test =
  validDLTest test . applyMonitoring d test . QC.property $ k (scriptFromDL test)

withDLScriptPrefix :: (DynLogicModel s, Testable a) => DynFormula s -> (Actions s -> a) -> DynLogicTest s -> Property
withDLScriptPrefix f k test =
  QC.withSize $ \n ->
    let d = unDynFormula f n
        test' = unfailDLTest d test
     in validDLTest test' . applyMonitoring d test' . QC.property $ k (scriptFromDL test')

-- | Validate generated test case.
--
-- Test case generation does not always produce a valid test case. In
-- some cases, we did not find a suitable test case matching some
-- `DynFormula` and we are `Stuck`, hence we want to discard the test
-- case and start over ; in other cases we found a genuine issue with
-- the formula leading to the impossibility of producing a valid test
-- case.
validDLTest :: StateModel s => DynLogicTest s -> Property -> Property
validDLTest test prop =
  case test of
    DLScript{} -> QC.counterexample (show test) prop
    Stuck{} -> QC.property QC.Discard
    _other -> QC.counterexample (show test) False

generateDLTest :: DynLogicModel s => DynLogic s -> Int -> Gen (DynLogicTest s)
generateDLTest d size = generate chooseNextStep d 0 (initialStateFor d) size

onDLTestSeq :: (TestSequence s -> TestSequence s) -> DynLogicTest s -> DynLogicTest s
onDLTestSeq f (BadPrecondition ss bad s) = BadPrecondition (f ss) bad s
onDLTestSeq f (Looping ss) = Looping (f ss)
onDLTestSeq f (Stuck ss s) = Stuck (f ss) s
onDLTestSeq f (DLScript ss) = DLScript (f ss)

consDLTest :: TestStep s -> DynLogicTest s -> DynLogicTest s
consDLTest step = onDLTestSeq (step :>)

consDLTestW :: Witnesses () -> DynLogicTest s -> DynLogicTest s
consDLTestW w = onDLTestSeq (addW w)
  where
    addW Do{} ss = ss
    addW (Witness a w') ss = TestSeqWitness a (addW w' ss)

generate
  :: (Monad m, DynLogicModel s)
  => (Annotated s -> Int -> DynLogic s -> m (NextStep s))
  -> DynLogic s
  -> Int
  -> Annotated s
  -> Int
  -> m (DynLogicTest s)
generate chooseNextStepFun d n s size =
  if n > sizeLimit size
    then return $ Looping TestSeqStop
    else do
      let preferred = if n > size then stopping d else noStopping d
          useStep (BadAction (Witnesses ws bad)) _ = return $ BadPrecondition (TestSeqStopW ws) bad s
          useStep StoppingStep _ = return $ DLScript TestSeqStop
          useStep (Stepping step d') _ =
            case discardWitnesses step of
              var := act ->
                consDLTest step
                  <$> generate
                    chooseNextStepFun
                    d'
                    (n + 1)
                    (computeNextState s act var)
                    size
          useStep NoStep alt = alt
      foldr
        (\step k -> do try <- chooseNextStepFun s n step; useStep try k)
        (return $ Stuck TestSeqStop s)
        [preferred, noAny preferred, d, noAny d]

sizeLimit :: Int -> Int
sizeLimit size = 2 * size + 20

initialStateFor :: StateModel s => DynLogic s -> Annotated s
initialStateFor _ = initialAnnotatedState

stopping :: DynLogic s -> DynLogic s
stopping EmptySpec = EmptySpec
stopping Stop = Stop
stopping (After act k) = After act k
stopping (Error m k) = Error m k
stopping (AfterAny _) = EmptySpec
stopping (Alt b d d') = Alt b (stopping d) (stopping d')
stopping (Stopping d) = d
stopping (Weight w d) = Weight w (stopping d)
stopping (ForAll _ _) = EmptySpec -- ???
stopping (Monitor f d) = Monitor f (stopping d)

noStopping :: DynLogic s -> DynLogic s
noStopping EmptySpec = EmptySpec
noStopping Stop = EmptySpec
noStopping (After act k) = After act k
noStopping (Error m k) = Error m k
noStopping (AfterAny k) = AfterAny k
noStopping (Alt b d d') = Alt b (noStopping d) (noStopping d')
noStopping (Stopping _) = EmptySpec
noStopping (Weight w d) = Weight w (noStopping d)
noStopping (ForAll q f) = ForAll q f
noStopping (Monitor f d) = Monitor f (noStopping d)

noAny :: DynLogic s -> DynLogic s
noAny EmptySpec = EmptySpec
noAny Stop = Stop
noAny (After act k) = After act k
noAny (Error m k) = Error m k
noAny (AfterAny _) = EmptySpec
noAny (Alt b d d') = Alt b (noAny d) (noAny d')
noAny (Stopping d) = Stopping (noAny d)
noAny (Weight w d) = Weight w (noAny d)
noAny (ForAll q f) = ForAll q f
noAny (Monitor f d) = Monitor f (noAny d)

nextSteps :: DynLogic s -> Gen [(Double, Witnesses (DynLogic s))]
nextSteps = nextSteps' generateQ

nextSteps' :: Monad m => (forall a. Quantification a -> m a) -> DynLogic s -> m [(Double, Witnesses (DynLogic s))]
nextSteps' _ EmptySpec = pure []
nextSteps' _ Stop = pure [(1, Do $ Stop)]
nextSteps' _ (After act k) = pure [(1, Do $ After act k)]
nextSteps' _ (Error m k) = pure [(1, Do $ Error m k)]
nextSteps' _ (AfterAny k) = pure [(1, Do $ AfterAny k)]
nextSteps' gen (Alt _ d d') = (++) <$> nextSteps' gen d <*> nextSteps' gen d'
nextSteps' gen (Stopping d) = nextSteps' gen d
nextSteps' gen (Weight w d) = do
  steps <- nextSteps' gen d
  pure [(w * w', s) | (w', s) <- steps, w * w' > never]
nextSteps' gen (ForAll q f) = do
  x <- gen q
  map (second $ Witness x) <$> nextSteps' gen (f x)
nextSteps' gen (Monitor _f d) = nextSteps' gen d

chooseOneOf :: [(Double, a)] -> Gen a
chooseOneOf steps = QC.frequency [(round (w / never), return s) | (w, s) <- steps]

never :: Double
never = 1.0e-9

data NextStep s
  = StoppingStep
  | Stepping (Witnesses (Step s)) (DynLogic s)
  | NoStep
  | BadAction (Witnesses (FailingAction s))

chooseNextStep :: DynLogicModel s => Annotated s -> Int -> DynLogic s -> Gen (NextStep s)
chooseNextStep s n d = do
  nextSteps d >>= \case
    [] -> return NoStep
    steps -> do
      let bads = concatMap (findBad . snd) steps
          findBad = traverse $ flip badActions s
      case bads of
        [] -> do
          chosen <- chooseOneOf steps
          let takeStep = \case
                Stop -> return StoppingStep
                After a k ->
                  return $ Stepping (Do $ mkVar n := a) (k (mkVar n) (computeNextState s a (mkVar n)))
                AfterAny k -> do
                  m <- keepTryingUntil 100 (computeArbitraryAction s) $
                    \case
                      Some act -> computePrecondition s act && not (restrictedPolar act)
                  case m of
                    Nothing -> return NoStep
                    Just (Some a@ActionWithPolarity{}) ->
                      return $
                        Stepping
                          (Do $ mkVar n := a)
                          (k (computeNextState s a (mkVar n)))
                EmptySpec -> error "chooseNextStep: EmptySpec"
                ForAll{} -> error "chooseNextStep: ForAll"
                Error{} -> error "chooseNextStep: Error"
                Alt{} -> error "chooseNextStep: Alt"
                Stopping{} -> error "chooseNextStep: Stopping"
                Weight{} -> error "chooseNextStep: Weight"
                Monitor{} -> error "chooseNextStep: Monitor"
              go (Do d') = takeStep d'
              go (Witness a step) =
                go step
                  >>= pure . \case
                    StoppingStep -> StoppingStep -- TODO: This is a bit fishy
                    Stepping step' dl -> Stepping (Witness a step') dl
                    NoStep -> NoStep
                    BadAction bad -> BadAction (Witness a bad)
          go chosen
        b : _ -> return $ BadAction b

chooseUniqueNextStep :: (MonadFail m, DynLogicModel s) => Annotated s -> Int -> DynLogic s -> m (NextStep s)
chooseUniqueNextStep s n d = do
  steps <- map snd <$> nextSteps' (const bad) d
  case steps of
    [] -> return NoStep
    [Do EmptySpec] -> return NoStep
    [Do Stop] -> return StoppingStep
    [Do (After a k)] -> return $ Stepping (Do $ mkVar n := a) (k (mkVar n) (computeNextState s a (mkVar n)))
    _ -> bad
  where
    bad = fail "chooseUniqueNextStep: non-unique action in DynLogic"

keepTryingUntil :: Int -> Gen a -> (a -> Bool) -> Gen (Maybe a)
keepTryingUntil 0 _ _ = return Nothing
keepTryingUntil n g p = do
  x <- g
  if p x then return $ Just x else QC.scale (+ 1) $ keepTryingUntil (n - 1) g p

shrinkDLTest :: DynLogicModel s => DynLogic s -> DynLogicTest s -> [DynLogicTest s]
shrinkDLTest _ (Looping _) = []
shrinkDLTest d tc =
  [ test | as' <- shrinkScript d (getScript tc), let pruned = pruneDLTest d as'
                                                     test = makeTestFromPruned d pruned,
  -- Don't shrink a non-executable test case to an executable one.
  case (tc, test) of
    (DLScript _, _) -> True
    (_, DLScript _) -> False
    _ -> True
  ]

nextStateStep :: StateModel s => Step s -> Annotated s -> Annotated s
nextStateStep (var := act) s = computeNextState s act var

shrinkScript :: forall s. DynLogicModel s => DynLogic s -> TestSequence s -> [TestSequence s]
shrinkScript = shrink' initialAnnotatedState
  where
    shrink' :: Annotated s -> DynLogic s -> TestSequence s -> [TestSequence s]
    shrink' s d ss = structural s d ss ++ nonstructural s d ss

    nonstructural s d (TestSeqWitness a ss) =
      [ TestSeqWitness a' ss'
      | a' <- shrinkWitness @s d a
      , ss' <- ss : shrink' s (stepDLSeq d s $ TestSeqWitness a TestSeqStop) ss
      ]
    nonstructural s d (TestSeqStep step@(var := act) ss) =
      [TestSeqStep (unsafeCoerceVar var := act') ss | Some act'@ActionWithPolarity{} <- computeShrinkAction s act]
        ++ [ TestSeqStep step ss'
           | ss' <-
              shrink'
                (nextStateStep step s)
                (stepDLStep d s step)
                ss
           ]
    nonstructural _ _ TestSeqStop = []

    structural _ _ TestSeqStopW{} = []
    structural s d (step :> rest) =
      TestSeqStop
        : reverse (takeWhile (not . nullSeq) [dropSeq (n - 1) rest | n <- iterate (* 2) 1])
        ++ map (step :>) (shrink' (nextStateStep (discardWitnesses step) s) (stepDLSeq d s (step :> TestSeqStop)) rest)

shrinkWitness :: (StateModel s, Typeable a) => DynLogic s -> a -> [a]
shrinkWitness (ForAll (q :: Quantification a) _) (a :: a') =
  case eqT @a @a' of
    Just Refl | isaQ q a -> shrinkQ q a
    _ -> []
shrinkWitness (Alt _ d d') a = shrinkWitness d a ++ shrinkWitness d' a
shrinkWitness (Stopping d) a = shrinkWitness d a
shrinkWitness (Weight _ d) a = shrinkWitness d a
shrinkWitness (Monitor _ d) a = shrinkWitness d a
shrinkWitness EmptySpec{} _ = []
shrinkWitness Stop{} _ = []
shrinkWitness Error{} _ = []
shrinkWitness After{} _ = []
shrinkWitness AfterAny{} _ = []

-- The result of pruning a list of actions is a prefix of a list of actions that
-- could have been generated by the dynamic logic.
pruneDLTest :: forall s. DynLogicModel s => DynLogic s -> TestSequence s -> TestSequence s
pruneDLTest dl = prune [dl] initialAnnotatedState
  where
    prune [] _ _ = TestSeqStop
    prune _ _ TestSeqStop = TestSeqStop
    prune ds s (TestSeqWitness a ss) =
      case [d' | d <- ds, d' <- stepDLW @s d a] of
        [] -> prune ds s ss
        ds' -> TestSeqWitness a $ prune ds' s ss
    prune ds s (TestSeqStep step@(_ := act) ss)
      | computePrecondition s act =
          case [d' | d <- ds, d' <- stepDL d s (Do step)] of
            [] -> prune ds s ss
            ds' -> TestSeqStep step $ prune ds' (nextStateStep step s) ss
      | otherwise = prune ds s ss

stepDL :: DynLogicModel s => DynLogic s -> Annotated s -> TestStep s -> [DynLogic s]
stepDL (After a k) s (Do (var := act))
  -- TOOD: make this nicer when we migrate to 9.2 where we can just bind
  -- the type variables cleanly and do `Just Refl <- eqT ...` here instead.
  | Some a == Some act = [k (unsafeCoerceVar var) (computeNextState s act (unsafeCoerceVar var))]
stepDL (AfterAny k) s (Do (var := act))
  | not (restrictedPolar act) = [k (computeNextState s act var)]
stepDL (Alt _ d d') s step = stepDL d s step ++ stepDL d' s step
stepDL (Stopping d) s step = stepDL d s step
stepDL (Weight _ d) s step = stepDL d s step
stepDL (ForAll (q :: Quantification a) f) s (Witness (a :: a') step) =
  case eqT @a @a' of
    Just Refl -> [d | isaQ q a, d <- stepDL (f a) s step]
    Nothing -> []
stepDL (Monitor _f d) s step = stepDL d s step
stepDL _ _ _ = []

stepDLW :: forall s a. (DynLogicModel s, Typeable a) => DynLogic s -> a -> [DynLogic s]
stepDLW (ForAll (q :: Quantification a') k) a =
  case eqT @a @a' of
    Just Refl -> [k a | isaQ q a]
    Nothing -> []
stepDLW (Alt _ d d') a = stepDLW d a ++ stepDLW d' a
stepDLW (Monitor _ d) a = stepDLW d a
stepDLW (Weight _ d) a = stepDLW d a
stepDLW (Stopping d) a = stepDLW d a
stepDLW _ _ = []

stepDLSeq :: DynLogicModel s => DynLogic s -> Annotated s -> TestSequence s -> DynLogic s
stepDLSeq d _ (TestSeqStopW Do{}) = d
stepDLSeq d s (TestSeqStopW (Witness a w)) = stepDLSeq (stepDLWitness d a) s (TestSeqStopW w)
stepDLSeq d s (step@(Witnesses _ (var := act)) :> ss) =
  stepDLSeq (demonicAlt $ stepDL d s step) (computeNextState s act var) ss

stepDLWitness :: forall a s. (DynLogicModel s, Typeable a) => DynLogic s -> a -> DynLogic s
stepDLWitness d a = demonicAlt $ stepDLW d a

stepDLStep :: DynLogicModel s => DynLogic s -> Annotated s -> Step s -> DynLogic s
stepDLStep d s step = stepDLSeq d s (TestSeqStep step TestSeqStop)

demonicAlt :: [DynLogic s] -> DynLogic s
demonicAlt [] = EmptySpec
demonicAlt ds = foldr1 (Alt Demonic) ds

propPruningGeneratedScriptIsNoop :: DynLogicModel s => DynLogic s -> Property
propPruningGeneratedScriptIsNoop d =
  QC.forAll (QC.sized $ \n -> QC.choose (1, max 1 n) >>= generateDLTest d) $ \test ->
    let script = case test of
          BadPrecondition s _ _ -> s
          Looping s -> s
          Stuck s _ -> s
          DLScript s -> s
     in script == pruneDLTest d script

getScript :: DynLogicTest s -> TestSequence s
getScript (BadPrecondition s _ _) = s
getScript (Looping s) = s
getScript (Stuck s _) = s
getScript (DLScript s) = s

makeTestFromPruned :: forall s. DynLogicModel s => DynLogic s -> TestSequence s -> DynLogicTest s
makeTestFromPruned dl = make dl initialAnnotatedState
  where
    make d s TestSeqStop
      | b : _ <- badActions @s d s = BadPrecondition TestSeqStop b s
      | stuck d s = Stuck TestSeqStop s
      | otherwise = DLScript TestSeqStop
    make d s (TestSeqWitness a ss) =
      onDLTestSeq (TestSeqWitness a) $
        make
          (stepDLWitness d a)
          s
          ss
    make d s (TestSeqStep step ss) =
      onDLTestSeq (TestSeqStep step) $
        make
          (stepDLStep d s step)
          (nextStateStep step s)
          ss

-- | If failed, return the prefix up to the failure. Also prunes the test in case the model has
--   changed.
unfailDLTest :: DynLogicModel s => DynLogic s -> DynLogicTest s -> DynLogicTest s
unfailDLTest d test = makeTestFromPruned d $ pruneDLTest d steps
  where
    steps = case test of
      BadPrecondition as _ _ -> as
      Stuck as _ -> as
      DLScript as -> as
      Looping as -> as

stuck :: DynLogicModel s => DynLogic s -> Annotated s -> Bool
stuck EmptySpec _ = True
stuck Stop _ = False
stuck (After _ _) _ = False
stuck (Error _ _) _ = False
stuck (AfterAny _) s =
  not $
    canGenerate
      0.01
      (computeArbitraryAction s)
      ( \case
          Some act ->
            computePrecondition s act
              && not (restrictedPolar act)
      )
stuck (Alt Angelic d d') s = stuck d s && stuck d' s
stuck (Alt Demonic d d') s = stuck d s || stuck d' s
stuck (Stopping d) s = stuck d s
stuck (Weight w d) s = w < never || stuck d s
stuck (ForAll _ _) _ = False
stuck (Monitor _ d) s = stuck d s

scriptFromDL :: DynLogicTest s -> Actions s
scriptFromDL (DLScript s) = Actions $ sequenceSteps s
scriptFromDL _ = Actions []

sequenceSteps :: TestSequence s -> [Step s]
sequenceSteps (TestSeq ss) =
  case discardWitnesses ss of
    ContStop -> []
    ContStep s ss' -> s : sequenceSteps ss'

badActionsGiven :: StateModel s => DynLogic s -> Annotated s -> Witnesses a -> [Witnesses (FailingAction s)]
badActionsGiven Stop _ _ = []
badActionsGiven EmptySpec _ _ = []
badActionsGiven AfterAny{} _ _ = []
badActionsGiven (ForAll _ k) s (Witness a step) =
  case cast a of
    Just a' -> Witness a' <$> badActionsGiven (k a') s step
    _ -> []
badActionsGiven (Alt _ d d') s w = badActionsGiven d s w ++ badActionsGiven d' s w
badActionsGiven (Stopping d) s w = badActionsGiven d s w
badActionsGiven (Weight k d) s w = if k < never then [] else badActionsGiven d s w
badActionsGiven (Monitor _ d) s w = badActionsGiven d s w
badActionsGiven d s (Do _) = Do <$> badActions d s
badActionsGiven Error{} _ _ = []
badActionsGiven After{} _ _ = []

badActions :: StateModel s => DynLogic s -> Annotated s -> [FailingAction s]
badActions EmptySpec _ = []
badActions Stop _ = []
badActions (After a _) s
  | computePrecondition s a = []
  | otherwise = [ActionFail a]
badActions (Error m _) _s = [ErrorFail m]
badActions (AfterAny _) _ = []
badActions (Alt _ d d') s = badActions d s ++ badActions d' s
badActions (Stopping d) s = badActions d s
badActions (Weight w d) s = if w < never then [] else badActions d s
badActions (ForAll _ _) _ = []
badActions (Monitor _ d) s = badActions d s

applyMonitoring :: DynLogicModel s => DynLogic s -> DynLogicTest s -> Property -> Property
applyMonitoring d (DLScript s) p =
  case findMonitoring d initialAnnotatedState s of
    Just f -> f p
    Nothing -> p
applyMonitoring _ Stuck{} p = p
applyMonitoring _ Looping{} p = p
applyMonitoring _ BadPrecondition{} p = p

findMonitoring :: DynLogicModel s => DynLogic s -> Annotated s -> TestSequence s -> Maybe (Property -> Property)
findMonitoring Stop _s TestSeqStop = Just id
findMonitoring (After a k) s (TestSeqStep (var := a') as)
  -- TODO: do nicely with eqT instead (avoids `unsafeCoerceVar`)
  | Some a == Some a' = findMonitoring (k (unsafeCoerceVar var) s') s' as
  where
    s' = computeNextState s a' (unsafeCoerceVar var)
findMonitoring (AfterAny k) s as@(TestSeqStep (_var := a) _)
  | not (restrictedPolar a) = findMonitoring (After a $ const k) s as
findMonitoring (Alt _b d d') s as =
  -- Give priority to monitoring matches to the left. Combining both
  -- results in repeated monitoring from always, which is unexpected.
  findMonitoring d s as <|> findMonitoring d' s as
findMonitoring (Stopping d) s as = findMonitoring d s as
findMonitoring (Weight _ d) s as = findMonitoring d s as
findMonitoring (ForAll (_q :: Quantification a) k) s (TestSeq (Witness (a :: a') as)) =
  case eqT @a @a' of
    Just Refl -> findMonitoring (k a) s (TestSeq as)
    Nothing -> Nothing
findMonitoring (Monitor m d) s as =
  (m .) <$> findMonitoring d s as
findMonitoring _ _ _ = Nothing