From 8f4623d2d9b552f1f4938bdf3f565970741ab412 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/monadtranscontrol-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      | 15 +++++++++++++++
 hedgehog/src/Hedgehog/Internal/Property.hs |  1 +
 hedgehog/src/Hedgehog/Internal/Tree.hs     |  6 ++++++
 3 files changed, 22 insertions(+)

diff --git a/hedgehog/src/Hedgehog/Internal/Gen.hs b/hedgehog/src/Hedgehog/Internal/Gen.hs
index 9a8ca552..6e3f117b 100644
--- a/hedgehog/src/Hedgehog/Internal/Gen.hs
+++ b/hedgehog/src/Hedgehog/Internal/Gen.hs
@@ -165,6 +165,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 +191,7 @@ import           Data.Bifunctor (first)
 import           Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
 import qualified Data.Char as Char
+import           Data.Coerce (coerce)
 import           Data.Foldable (for_, toList)
 import           Data.Functor.Identity (Identity(..))
 import           Data.Int (Int8, Int16, Int32, Int64)
@@ -587,6 +589,19 @@ instance MonadBase b m => MonadBase b (GenT m) where
   liftBase =
     lift . liftBase
 
+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
+
 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..b67aa970 100644
--- a/hedgehog/src/Hedgehog/Internal/Property.hs
+++ b/hedgehog/src/Hedgehog/Internal/Property.hs
@@ -178,6 +178,7 @@ newtype PropertyT m a =
     , Monad
     , MonadIO
     , MonadBase b
+    , MonadBaseControl b
     , MonadThrow
     , MonadCatch
     , MonadReader r
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 =