tripping with a monadic action

hedgehog/src/Hedgehog.hs

@@ -106,6 +106,7 @@ module Hedgehog (
   , (===)
   , (/==)
   , tripping
+  , trippingM
 
   , eval
   , evalNF
@@ -196,7 +197,7 @@ import           Hedgehog.Internal.State (Action, Sequential(..), Parallel(..))
 import           Hedgehog.Internal.State (executeSequential, executeParallel)
 import           Hedgehog.Internal.State (Var(..), Symbolic, Concrete(..), concrete, opaque)
 import           Hedgehog.Internal.TH (discover, discoverPrefix)
-import           Hedgehog.Internal.Tripping (tripping)
+import           Hedgehog.Internal.Tripping (tripping, trippingM)
 
 
 -- $functors

hedgehog/src/Hedgehog/Internal/State.hs

@@ -548,6 +548,12 @@ action commands =
     Command mgenInput exec callbacks <-
       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 ->
@@ -559,8 +565,6 @@ action commands =
       pure Nothing
 
     else do
-      output <- contextNewVar
-
       contextUpdate $
         callbackUpdate callbacks state0 input (Var output)
 

hedgehog/src/Hedgehog/Internal/Tripping.hs

@@ -1,6 +1,7 @@
 {-# OPTIONS_HADDOCK not-home #-}
 module Hedgehog.Internal.Tripping (
     tripping
+  , trippingM
   ) where
 
 import           Hedgehog.Internal.Property (MonadTest, Diff(..), success, failWith)
@@ -28,18 +29,25 @@ tripping ::
   -> (b -> f a)
   -> m ()
 tripping x encode decode =
+  trippingM x (pure . encode) (pure . decode)
+
+
+-- | Similar to tripping, but with a monadic action.
+trippingM ::
+     (MonadTest m, Applicative f, Show b, Show (f a), Eq (f a), HasCallStack)
+  => a
+  -> (a -> m b)
+  -> (b -> m (f a))
+  -> m ()
+trippingM x encode decode = do
   let
     mx =
       pure x
 
-    i =
-      encode x
-
-    my =
-      decode i
-  in
-    if mx == my then
-      success
+  i <- encode x
+  my <- decode i
+  if mx == my
+    then success
     else
       case valueDiff <$> mkValue mx <*> mkValue my of
         Nothing ->