Use state to update a Command's input

hedgehog/hedgehog.cabal

@@ -138,6 +138,7 @@ test-suite test
     Test.Hedgehog.Maybe
     Test.Hedgehog.Seed
     Test.Hedgehog.Skip
+    Test.Hedgehog.State
     Test.Hedgehog.Text
     Test.Hedgehog.Zip
 

hedgehog/src/Hedgehog/Internal/State.hs

@@ -43,7 +43,7 @@ module Hedgehog.Internal.State (
   , Parallel(..)
   , takeVariables
   , variablesOK
-  , dropInvalid
+  , rethreadState
   , action
   , sequential
   , parallel
@@ -380,6 +380,7 @@ callbackEnsure callbacks s0 s i o =
 -- your 'Command' list to 'sequential' or 'parallel'.
 --
 data Command gen m (state :: (Type -> Type) -> Type) =
+  -- | A "simple" command.
   forall input output.
   (TraversableB input, Show (input Symbolic), Show output, Typeable output) =>
   Command {
@@ -402,11 +403,47 @@ data Command gen m (state :: (Type -> Type) -> Type) =
         [Callback input output state]
     }
 
+  |
+  -- | An "advanced" command.
+  forall input0 input output.
+  (TraversableB input, Show (input Symbolic), Show output, Typeable output) =>
+  CommandA {
+    -- | A generator which provides random arguments for a command. If the
+    --   command cannot be executed in the current state, it should return
+    --   'Nothing'.
+    --
+      commandAGen ::
+        state Symbolic -> Maybe (gen input0)
+
+    -- | Turns the randomly generated argument into the command's input by
+    --   examining the state. This allows the input to depend on previous steps,
+    --   in a way that gets preserved during shrinking. If this returns
+    --   'Nothing', then the generated argument is invalid on the current state,
+    --   and the action will be dropped as with 'Require'.
+    --
+    , commandAMkInput ::
+        state Symbolic -> input0 -> Maybe (input Symbolic)
+
+    -- | Executes a command using the arguments generated by 'commandAGen' and
+    --   'commandAMkInput'.
+    --
+    , commandAExecute ::
+        input Concrete -> m output
+
+    -- | A set of callbacks which provide optional command configuration such
+    --   as pre-condtions, post-conditions and state updates.
+    --
+    , commandACallbacks ::
+        [Callback input output state]
+    }
+
 -- | Checks that input for a command can be executed in the given state.
 --
 commandGenOK :: Command gen m state -> state Symbolic -> Bool
 commandGenOK (Command inputGen _ _) state =
   Maybe.isJust (inputGen state)
+commandGenOK (CommandA inputGen _ _ _) state =
+  Maybe.isJust (inputGen state)
 
 -- | An instantiation of a 'Command' which can be executed, and its effect
 --   evaluated.
@@ -418,6 +455,9 @@ data Action m (state :: (Type -> Type) -> Type) =
       actionInput ::
         input Symbolic
 
+    , actionRefreshInput ::
+        state Symbolic -> Maybe (input Symbolic)
+
     , actionOutput ::
         Symbolic output
 
@@ -435,7 +475,7 @@ data Action m (state :: (Type -> Type) -> Type) =
     }
 
 instance Show (Action m state) where
-  showsPrec p (Action input (Symbolic (Name output)) _ _ _ _) =
+  showsPrec p (Action input _ (Symbolic (Name output)) _ _ _ _) =
     showParen (p > 10) $
       showString "Var " .
       showsPrec 11 output .
@@ -512,26 +552,28 @@ contextNewVar = do
   put $ Context state (insertSymbolic var vars)
   pure var
 
--- | Drops invalid actions from the sequence.
+-- | Pass the state through the actions, updating inputs and dropping invalid
+--   ones.
 --
-dropInvalid :: [Action m state] -> State (Context state) [Action m state]
-dropInvalid =
+rethreadState :: [Action m state] -> State (Context state) [Action m state]
+rethreadState =
   let
-    loop step@(Action input output _execute require update _ensure) = do
+    loop (Action _ refreshInput output exec require update ensure) = do
       Context state0 vars0 <- get
 
-      if require state0 input && variablesOK input vars0 then do
-        let
-          state =
-            update state0 input (Var output)
+      case refreshInput state0 of
+        Just input | require state0 input && variablesOK input vars0 -> do
+          let
+            state =
+              update state0 input (Var output)
 
-          vars =
-            insertSymbolic output vars0
+            vars =
+              insertSymbolic output vars0
 
-        put $ Context state vars
-        pure $ Just step
-      else
-        pure Nothing
+          put $ Context state vars
+          pure $ Just $ Action input refreshInput output exec require update ensure
+        _ ->
+          pure Nothing
   in
     fmap Maybe.catMaybes . traverse loop
 
@@ -545,34 +587,56 @@ action commands =
   Gen.justT $ do
     Context state0 _ <- get
 
-    Command mgenInput exec callbacks <-
+    cmd <-
       Gen.element_ $ filter (\c -> commandGenOK c state0) commands
 
     -- If we shrink the input, we still want to use the same output. Otherwise
     -- any actions using this output as part of their input will be dropped. But
     -- the existing output is still in the context, so `contextNewVar` will
     -- create a new one. To avoid that, we generate the output before the input.
-    output <- contextNewVar
-
-    input <-
-      case mgenInput state0 of
-        Nothing ->
-          error "genCommand: internal error, tried to use generator with invalid state."
-        Just gen ->
-          hoist lift $ Gen.toGenT gen
-
-    if not $ callbackRequire callbacks state0 input then
-      pure Nothing
-
-    else do
-      contextUpdate $
-        callbackUpdate callbacks state0 input (Var output)
 
-      pure . Just $
-        Action input output exec
-          (callbackRequire callbacks)
-          (callbackUpdate callbacks)
-          (callbackEnsure callbacks)
+    case cmd of
+      Command mgenInput exec callbacks -> do
+        output <- contextNewVar
+        input <-
+          case mgenInput state0 of
+            Nothing ->
+              error "genCommand: internal error, tried to use generator with invalid state."
+            Just gen ->
+              hoist lift $ Gen.toGenT gen
+
+        if not $ callbackRequire callbacks state0 input then
+          pure Nothing
+        else do
+          contextUpdate $
+            callbackUpdate callbacks state0 input (Var output)
+
+          pure . Just $
+            Action input (const $ Just input) output exec
+              (callbackRequire callbacks)
+              (callbackUpdate callbacks)
+              (callbackEnsure callbacks)
+      CommandA mgenInput mkInput exec callbacks -> do
+        output <- contextNewVar
+        input0 <-
+          case mgenInput state0 of
+            Nothing ->
+              error "genCommand: internal error, tried to use generator with invalid state."
+            Just gen ->
+              hoist lift $ Gen.toGenT gen
+
+        case mkInput state0 input0 of
+          Just input | callbackRequire callbacks state0 input -> do
+            contextUpdate $
+              callbackUpdate callbacks state0 input (Var output)
+
+            pure . Just $
+              Action input (flip mkInput input0) output exec
+                (callbackRequire callbacks)
+                (callbackUpdate callbacks)
+                (callbackEnsure callbacks)
+          _ ->
+            pure Nothing
 
 genActions ::
      (MonadGen gen, MonadTest m)
@@ -583,7 +647,7 @@ genActions ::
 genActions range commands ctx = do
   xs <- Gen.fromGenT . (`evalStateT` ctx) . distributeT $ Gen.list range (action commands)
   pure $
-    dropInvalid xs `runState` ctx
+    rethreadState xs `runState` ctx
 
 -- | A sequence of actions to execute.
 --
@@ -594,7 +658,7 @@ newtype Sequential m state =
     }
 
 renderAction :: Action m state -> [String]
-renderAction (Action input (Symbolic (Name output)) _ _ _ _) =
+renderAction (Action input _ (Symbolic (Name output)) _ _ _ _) =
   let
     prefix0 =
       "Var " ++ show output ++ " = "
@@ -610,7 +674,7 @@ renderAction (Action input (Symbolic (Name output)) _ _ _ _) =
         fmap (prefix ++) xs
 
 renderActionResult :: Environment -> Action m state -> [String]
-renderActionResult env (Action _ output@(Symbolic (Name name)) _ _ _ _) =
+renderActionResult env (Action _ _ output@(Symbolic (Name name)) _ _ _ _) =
   let
     prefix0 =
       "Var " ++ show name ++ " = "
@@ -709,7 +773,7 @@ data ActionCheck state =
     }
 
 execute :: (MonadTest m, HasCallStack) => Action m state -> StateT Environment m (ActionCheck state)
-execute (Action sinput soutput exec _require update ensure) =
+execute (Action sinput _ soutput exec _require update ensure) =
   withFrozenCallStack $ do
     env0 <- get
     input <- evalEither $ reify env0 sinput
@@ -736,7 +800,7 @@ executeUpdateEnsure ::
   => (state Concrete, Environment)
   -> Action m state
   -> m (state Concrete, Environment)
-executeUpdateEnsure (state0, env0) (Action sinput soutput exec _require update ensure) =
+executeUpdateEnsure (state0, env0) (Action sinput _ soutput exec _require update ensure) =
   withFrozenCallStack $ do
     input <- evalEither $ reify env0 sinput
     output <- exec input

hedgehog/test/Test/Hedgehog/State.hs

@@ -0,0 +1,95 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TupleSections #-}
+
+module Test.Hedgehog.State where
+
+import           Control.Applicative (Const(..))
+import           Control.Monad (void)
+import           Control.Monad.IO.Class (liftIO)
+import qualified Data.IORef as IORef
+import           Hedgehog
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Internal.Config as Config
+import qualified Hedgehog.Internal.Property as Property
+import qualified Hedgehog.Internal.Runner as Runner
+import qualified Hedgehog.Range as Range
+
+-- | Test that 'commandAMkInput' works as expected when shrinking.
+--
+-- We create a state machine that always generates two actions. Initially, one
+-- will have the number 5 as input and put it in state. The other will have
+-- (True, 5) as input. It checks the number is less than 5. Since it's not, we
+-- start shrinking.
+--
+-- We shrink the first action initially, through 4,3,2,1,0. Each of these
+-- changes the input to the second action, even though we're not shrinking that,
+-- because `commandMkInput` looks at the state. The second action passes with
+-- each of these.
+--
+-- So then we shrink the second action, to (False, 5). That fails again, so we
+-- go back to shrinking the first one. All of those shrinks pass again.
+--
+-- We log the list of inputs to the second action, and after running this state
+-- machine (and ignoring its result) we check that this list is correct.
+--
+-- This depends on the order shrinks are performed in state machines. Hopefully
+-- it won't be too fragile.
+prop_mkInput :: Property
+prop_mkInput =
+  withTests 1 . property $ do
+    actionListsRef <- liftIO $ IORef.newIORef []
+    let
+      prop = property $ do
+        actions <- forAll $ Gen.sequential
+          (Range.linear 2 2)
+          (Const Nothing)
+          [ let
+              commandGen = \case
+                Const Nothing ->
+                  Just $ Const <$> Gen.shrink (\n -> reverse [0..n-1])
+                                              (pure (5 :: Int))
+                Const (Just _) -> Nothing
+              commandExecute _ = pure ()
+              commandCallbacks =
+                [Update $ \_ (Const input) _ -> Const $ Just input]
+            in
+              Command { .. }
+          , let
+              commandAGen = \case
+                Const Nothing ->
+                  Nothing
+                Const (Just _) ->
+                  Just $ Gen.shrink (\b -> if b then [False] else [])
+                                    (pure True)
+              commandAMkInput (Const st) inputB = case st of
+                Nothing ->
+                  Nothing
+                Just stateN ->
+                  Just $ Const (stateN, inputB)
+              commandAExecute (Const (stateN, inputB)) = liftIO $ do
+                IORef.modifyIORef' actionListsRef ((stateN, inputB) :)
+              commandACallbacks =
+                [Ensure $ \_ _ (Const (stateN, _)) _ -> diff stateN (<) 5]
+            in
+              CommandA { .. }
+          ]
+        executeSequential (Const Nothing) actions
+
+    -- We could simply use `check` here, but that prints its output to the test
+    -- logs.
+    seed <- Config.resolveSeed Nothing
+    void $ liftIO $ Runner.checkReport (Property.propertyConfig prop)
+                                       0
+                                       seed
+                                       (Property.propertyTest prop)
+                                       (const $ pure ())
+
+    actionLists <- liftIO $ reverse <$> IORef.readIORef actionListsRef
+    actionLists === ((, True) <$> [5,4..0]) ++ ((, False) <$> [5,4..0])
+
+tests :: IO Bool
+tests =
+  checkParallel $$(discover)

hedgehog/test/test.hs

@@ -6,6 +6,7 @@ import qualified Test.Hedgehog.Filter
 import qualified Test.Hedgehog.Maybe
 import qualified Test.Hedgehog.Seed
 import qualified Test.Hedgehog.Skip
+import qualified Test.Hedgehog.State
 import qualified Test.Hedgehog.Text
 import qualified Test.Hedgehog.Zip
 
@@ -19,6 +20,7 @@ main =
     , Test.Hedgehog.Maybe.tests
     , Test.Hedgehog.Seed.tests
     , Test.Hedgehog.Skip.tests
+    , Test.Hedgehog.State.tests
     , Test.Hedgehog.Text.tests
     , Test.Hedgehog.Zip.tests
     ]