From 912f571c1f17a3577fe75c36484463f9ee67c1c2 Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Thu, 10 Oct 2019 13:42:37 -0400
Subject: [PATCH] Add a MonadBaseControl instance for PropertyT (#328)

Under an assumption I believe is reasonable (see [this issue](basvandijk/monad-control#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 | 11 ++++++++++
 hedgehog/src/Hedgehog/Internal/Tree.hs     |  6 ++++++
 3 files changed, 42 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 d6aee8c6..95ca813c 100644
--- a/hedgehog/src/Hedgehog/Internal/Property.hs
+++ b/hedgehog/src/Hedgehog/Internal/Property.hs
@@ -1,6 +1,7 @@
 {-# OPTIONS_HADDOCK not-home #-}
 {-# LANGUAGE ApplicativeDo #-}
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE DeriveFoldable #-}
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE DeriveLift #-}
@@ -203,6 +204,16 @@ newtype PropertyT m a =
 -- NOTE: Move this to the deriving list above when we drop 7.10
 deriving instance MonadResource m => MonadResource (PropertyT m)
 
+-- NOTE: Move this to the deriving list above when we drop 8.0
+#if __GLASGOW_HASKELL__ >= 802
+deriving instance MonadBaseControl b m => MonadBaseControl b (PropertyT m)
+#else
+instance MonadBaseControl b m => MonadBaseControl b (PropertyT m) where
+  type StM (PropertyT m) a = StM (TestT (GenT m)) a
+  liftBaseWith f = PropertyT $ liftBaseWith $ \rib -> f (rib . unPropertyT)
+  restoreM = PropertyT . restoreM
+#endif
+
 -- | A test monad allows the assertion of expectations.
 --
 type Test =
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 =