From a439b09e5e5531bff08466721787f6a604c1216a Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Wed, 25 Sep 2019 16:32:30 -0400
Subject: [PATCH] Add a MonadBaseControl instance for PropertyT

Under an assumption I believe is reasonable (see
[this issue](https://github.com/basvandijk/monad-control/issues/48) and
[this SO question](https://stackoverflow.com/questions/58105759/monadbasecontrol-laws)
I believe it is possible to write a lawful `MonadBaseControl` instance
for `CofreeT`, and therefore also for the `TreeT`, `GenT`, and
`PropertyT` types in this package. I can't say for sure whether it
"makes sense", but hey, maybe it does.
---
 hedgehog/src/Hedgehog/Internal/Gen.hs      | 25 ++++++++++++++++++++++
 hedgehog/src/Hedgehog/Internal/Property.hs |  1 +
 hedgehog/src/Hedgehog/Internal/Tree.hs     |  6 ++++++
 3 files changed, 32 insertions(+)

diff --git a/hedgehog/src/Hedgehog/Internal/Gen.hs b/hedgehog/src/Hedgehog/Internal/Gen.hs
index 9a8ca552..b93c3ed5 100644
--- a/hedgehog/src/Hedgehog/Internal/Gen.hs
+++ b/hedgehog/src/Hedgehog/Internal/Gen.hs
@@ -18,6 +18,9 @@
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE UndecidableInstances #-} -- MonadBase
+#if __GLASGOW_HASKELL__ >= 806
+{-# LANGUAGE DerivingVia #-}
+#endif
 module Hedgehog.Internal.Gen (
   -- * Transformer
     Gen
@@ -165,6 +168,7 @@ module Hedgehog.Internal.Gen (
 import           Control.Applicative (Alternative(..),liftA2)
 import           Control.Monad (MonadPlus(..), filterM, replicateM, join)
 import           Control.Monad.Base (MonadBase(..))
+import           Control.Monad.Trans.Control (MonadBaseControl(..))
 import           Control.Monad.Catch (MonadThrow(..), MonadCatch(..))
 import           Control.Monad.Error.Class (MonadError(..))
 import           Control.Monad.Fail (MonadFail (..))
@@ -190,6 +194,9 @@ import           Data.Bifunctor (first)
 import           Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
 import qualified Data.Char as Char
+#if __GLASGOW_HASKELL__ < 806
+import           Data.Coerce (coerce)
+#endif
 import           Data.Foldable (for_, toList)
 import           Data.Functor.Identity (Identity(..))
 import           Data.Int (Int8, Int16, Int32, Int64)
@@ -587,6 +594,24 @@ instance MonadBase b m => MonadBase b (GenT m) where
   liftBase =
     lift . liftBase
 
+#if __GLASGOW_HASKELL__ >= 806
+deriving via (ReaderT Size (ReaderT Seed (TreeT (MaybeT m))))
+  instance MonadBaseControl b m => MonadBaseControl b (GenT m)
+#else
+instance MonadBaseControl b m => MonadBaseControl b (GenT m) where
+  type StM (GenT m) a = StM (GloopT m) a
+  liftBaseWith g = gloopToGen $ liftBaseWith $ \q -> g (\gen -> q (genToGloop gen))
+  restoreM = gloopToGen . restoreM
+
+type GloopT m = ReaderT Size (ReaderT Seed (TreeT (MaybeT m)))
+
+gloopToGen :: GloopT m a -> GenT m a
+gloopToGen = coerce
+
+genToGloop :: GenT m a -> GloopT m a
+genToGloop = coerce
+#endif
+
 instance MonadThrow m => MonadThrow (GenT m) where
   throwM =
     lift . throwM
diff --git a/hedgehog/src/Hedgehog/Internal/Property.hs b/hedgehog/src/Hedgehog/Internal/Property.hs
index 3fbf44d7..efd27d03 100644
--- a/hedgehog/src/Hedgehog/Internal/Property.hs
+++ b/hedgehog/src/Hedgehog/Internal/Property.hs
@@ -186,6 +186,7 @@ newtype PropertyT m a =
     )
 -- NOTE: Move this to the deriving list above when we drop 7.10
 deriving instance MonadResource m => MonadResource (PropertyT m)
+deriving instance MonadBaseControl b m => MonadBaseControl b (PropertyT m)
 
 -- | A test monad allows the assertion of expectations.
 --
diff --git a/hedgehog/src/Hedgehog/Internal/Tree.hs b/hedgehog/src/Hedgehog/Internal/Tree.hs
index 6d749640..1649ecf4 100644
--- a/hedgehog/src/Hedgehog/Internal/Tree.hs
+++ b/hedgehog/src/Hedgehog/Internal/Tree.hs
@@ -66,6 +66,7 @@ import           Data.Functor.Classes (showsUnaryWith, showsBinaryWith)
 import qualified Data.Maybe as Maybe
 
 import           Hedgehog.Internal.Distributive
+import           Control.Monad.Trans.Control (MonadBaseControl (..))
 
 import           Prelude hiding (filter)
 
@@ -91,6 +92,11 @@ newtype TreeT m a =
       runTreeT :: m (NodeT m a)
     }
 
+instance MonadBaseControl b m => MonadBaseControl b (TreeT m) where
+  type StM (TreeT m) a = StM m (NodeT m a)
+  liftBaseWith f = TreeT $ liftBaseWith (\g -> pure <$> f (g . runTreeT))
+  restoreM = TreeT . restoreM
+
 -- | A node in a rose tree.
 --
 type Node =