From 7e27b1575c6044fdc4abc4c8514d3d617d779d53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Baldur=20Bl=C3=B6ndal?= <baldur.blondal@iohk.io>
Date: Mon, 26 Feb 2024 15:25:09 +0000
Subject: [PATCH] Add oneShot to Codensity functions.

---
 src/Control/Monad/Codensity.hs | 36 +++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/Control/Monad/Codensity.hs b/src/Control/Monad/Codensity.hs
index b524e5f..f934df3 100644
--- a/src/Control/Monad/Codensity.hs
+++ b/src/Control/Monad/Codensity.hs
@@ -50,7 +50,7 @@ import Data.Functor.Kan.Ran
 import Data.Functor.Plus
 import Data.Functor.Rep
 import Data.Type.Equality (type (~~))
-import GHC.Exts (TYPE)
+import GHC.Exts (TYPE, oneShot)
 
 -- |
 -- @'Codensity' f@ is the Monad generated by taking the right Kan extension
@@ -71,7 +71,7 @@ newtype Codensity (m :: k -> TYPE rep) a = Codensity
   }
 
 instance Functor (Codensity (k :: j -> TYPE rep)) where
-  fmap f (Codensity m) = Codensity (\k -> m (\x -> k (f x)))
+  fmap f (Codensity m) = Codensity $ oneShot (\k -> m (\x -> k (f x)))
   {-# INLINE fmap #-}
 
 instance Apply (Codensity (f :: k -> TYPE rep)) where
@@ -81,13 +81,13 @@ instance Apply (Codensity (f :: k -> TYPE rep)) where
 instance Applicative (Codensity (f :: k -> TYPE rep)) where
   pure x = Codensity (\k -> k x)
   {-# INLINE pure #-}
-  Codensity f <*> Codensity g = Codensity (\bfr -> f (\ab -> g (\x -> bfr (ab x))))
+  Codensity f <*> Codensity g = Codensity $ oneShot (\bfr -> f $ oneShot (\ab -> g $ oneShot (\x -> bfr (ab x))))
   {-# INLINE (<*>) #-}
 
 instance Monad (Codensity (f :: k -> TYPE rep)) where
   return = pure
   {-# INLINE return #-}
-  m >>= k = Codensity (\c -> runCodensity m (\a -> runCodensity (k a) c))
+  Codensity m >>= k = Codensity $ oneShot (\c -> m (oneShot $ \a -> runCodensity (k a) c))
   {-# INLINE (>>=) #-}
 
 -- Writing instances like
@@ -101,7 +101,7 @@ instance Monad (Codensity (f :: k -> TYPE rep)) where
 
 instance (f ~~ f', Fail.MonadFail f')
   => Fail.MonadFail (Codensity (f :: k -> TYPE rep)) where
-  fail msg = Codensity $ \ _ -> Fail.fail msg
+  fail msg = Codensity $ oneShot $ \ _ -> Fail.fail msg
   {-# INLINE fail #-}
 
 instance (m ~~ m', MonadIO m')
@@ -110,16 +110,16 @@ instance (m ~~ m', MonadIO m')
   {-# INLINE liftIO #-}
 
 instance MonadTrans Codensity where
-  lift m = Codensity (m >>=)
+  lift m = Codensity $ oneShot (m >>=)
   {-# INLINE lift #-}
 
 instance (v ~~ v', Alt v')
   => Alt (Codensity (v :: k -> TYPE rep)) where
-  Codensity m <!> Codensity n = Codensity (\k -> m k <!> n k)
+  Codensity m <!> Codensity n = Codensity $ oneShot (\k -> m k <!> n k)
   {-# INLINE (<!>) #-}
 
 instance (v ~~ v', Plus v') => Plus (Codensity (v :: k -> TYPE rep)) where
-  zero = Codensity (const zero)
+  zero = Codensity $ oneShot (const zero)
   {-# INLINE zero #-}
 
 {-
@@ -134,9 +134,9 @@ instance Plus v => MonadPlus (Codensity v) where
 
 instance (v ~~ v', Alternative v')
   => Alternative (Codensity (v :: k -> TYPE rep)) where
-  empty = Codensity (\_ -> empty)
+  empty = Codensity $ oneShot (\_ -> empty)
   {-# INLINE empty #-}
-  Codensity m <|> Codensity n = Codensity (\k -> m k <|> n k)
+  Codensity m <|> Codensity n = Codensity $ oneShot (\k -> m k <|> n k)
   {-# INLINE (<|>) #-}
 
 instance (v ~~ v', Alternative v')
@@ -172,7 +172,7 @@ codensityToAdjunction r = runCodensity r unit
 {-# INLINE codensityToAdjunction #-}
 
 adjunctionToCodensity :: Adjunction f g => g (f a) -> Codensity g a
-adjunctionToCodensity f = Codensity (\a -> fmap (rightAdjunct a) f)
+adjunctionToCodensity f = Codensity $ oneShot (\a -> fmap (rightAdjunct a) f)
 {-# INLINE adjunctionToCodensity #-}
 
 -- | The 'Codensity' monad of a representable 'Functor' is isomorphic to the
@@ -198,7 +198,7 @@ codensityToComposedRep (Codensity f) = f (\a -> tabulate $ \e -> (e, a))
 -- 'composedRepToCodensity' = 'ranToCodensity' . 'composedRepToRan'
 -- @
 composedRepToCodensity :: Representable u => u (Rep u, a) -> Codensity u a
-composedRepToCodensity hfa = Codensity $ \k -> fmap (\(e, a) -> index (k a) e) hfa
+composedRepToCodensity hfa = Codensity $ oneShot $ \k -> fmap (\(e, a) -> index (k a) e) hfa
 {-# INLINE composedRepToCodensity #-}
 
 -- | The 'Codensity' 'Monad' of a 'Functor' @g@ is the right Kan extension ('Ran')
@@ -218,21 +218,21 @@ ranToCodensity (Ran m) = Codensity m
 
 instance (m ~~ m', Functor f, MonadFree f m')
   => MonadFree f (Codensity (m :: k -> TYPE rep)) where
-  wrap t = Codensity (\h -> wrap (fmap (\p -> runCodensity p h) t))
+  wrap t = Codensity $ oneShot (\h -> wrap (fmap (\p -> runCodensity p h) t))
   {-# INLINE wrap #-}
 
 instance (m ~~ m', MonadReader r m')
   => MonadState r (Codensity (m :: k -> TYPE rep)) where
-  get = Codensity (ask >>=)
+  get = Codensity $ oneShot (ask >>=)
   {-# INLINE get #-}
-  put s = Codensity (\k -> local (const s) (k ()))
+  put s = Codensity $ oneShot (\k -> local (const s) (k ()))
   {-# INLINE put #-}
 
 instance (m ~~ m', MonadReader r m')
   => MonadReader r (Codensity (m :: k -> TYPE rep)) where
-  ask = Codensity (ask >>=)
+  ask = Codensity $ oneShot (ask >>=)
   {-# INLINE ask #-}
-  local f m = Codensity $ \c -> ask >>= \r -> local f . runCodensity m $ local (const r) . c
+  local f m = Codensity $ oneShot $ \c -> ask >>= \r -> local f . runCodensity m $ local (const r) . c
   {-# INLINE local #-}
 
 -- | Right associate all binds in a computation that generates a free monad
@@ -257,7 +257,7 @@ improve m = lowerCodensity m
 --
 -- > wrapCodensity (`finally` putStrLn "Done.")
 wrapCodensity :: (forall a. m a -> m a) -> Codensity m ()
-wrapCodensity f = Codensity (\k -> f (k ()))
+wrapCodensity f = Codensity $ oneShot (\k -> f (k ()))
 
 -- | @'reset' m@ delimits the continuation of any 'shift' inside @m@.
 --