{-# LANGUAGE FunctionalDependencies #-}

-- | Module: Control.Monad.HashCons
--
-- Provides a transformer, and a capability type class in the style of @mtl@,
-- for hash consing. See the Covenant wiki for how this works.
module Control.Monad.HashCons
  ( -- * Transformer
    HashConsT,
    runHashConsT,
    hashCons,
    lookupRef_,

    -- * Capability type class
    MonadHashCons (..),
  )
where

import Control.Monad.State.Strict
  ( StateT,
    get,
    modify,
    runStateT,
  )
import Control.Monad.Trans (MonadTrans (lift))
import Control.Monad.Trans.Except (ExceptT)
import Control.Monad.Trans.Maybe (MaybeT)
import Control.Monad.Trans.RWS.CPS (RWST)
import Control.Monad.Trans.Reader (ReaderT)
import Control.Monad.Trans.Writer.CPS (WriterT)
import Data.Bimap (Bimap)
import Data.Bimap qualified as Bimap
import Data.Kind (Type)

-- | A transformer implementing hash consing capabilities, with references of
-- type @r@ and referents of type @e@. It is assumed that values of type @e@
-- contain values of type @r@ in their capacity as references, though this is
-- not a requirement of this transformer.
--
-- = Important note
--
-- This implementation is not suitable for any @m@ that throws exceptions. This
-- includes @IO@, @ST@ and anything stacked atop them. For the reasons why, see
-- [here](https://github.com/haskell-effectful/effectful/blob/master/transformers.md#statet).
--
-- @since 1.0.0
newtype HashConsT (r :: Type) (e :: Type) (m :: Type -> Type) (a :: Type)
  = HashConsT (StateT (Bimap r e) m a)
  deriving
    ( -- | @since 1.0.0
      (forall a b. (a -> b) -> HashConsT r e m a -> HashConsT r e m b)
-> (forall a b. a -> HashConsT r e m b -> HashConsT r e m a)
-> Functor (HashConsT r e m)
forall a b. a -> HashConsT r e m b -> HashConsT r e m a
forall a b. (a -> b) -> HashConsT r e m a -> HashConsT r e m b
forall r e (m :: Type -> Type) a b.
Functor m =>
a -> HashConsT r e m b -> HashConsT r e m a
forall r e (m :: Type -> Type) a b.
Functor m =>
(a -> b) -> HashConsT r e m a -> HashConsT r e m b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall r e (m :: Type -> Type) a b.
Functor m =>
(a -> b) -> HashConsT r e m a -> HashConsT r e m b
fmap :: forall a b. (a -> b) -> HashConsT r e m a -> HashConsT r e m b
$c<$ :: forall r e (m :: Type -> Type) a b.
Functor m =>
a -> HashConsT r e m b -> HashConsT r e m a
<$ :: forall a b. a -> HashConsT r e m b -> HashConsT r e m a
Functor,
      -- | @since 1.0.0
      Functor (HashConsT r e m)
Functor (HashConsT r e m) =>
(forall a. a -> HashConsT r e m a)
-> (forall a b.
    HashConsT r e m (a -> b) -> HashConsT r e m a -> HashConsT r e m b)
-> (forall a b c.
    (a -> b -> c)
    -> HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m c)
-> (forall a b.
    HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b)
-> (forall a b.
    HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m a)
-> Applicative (HashConsT r e m)
forall a. a -> HashConsT r e m a
forall a b.
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m a
forall a b.
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
forall a b.
HashConsT r e m (a -> b) -> HashConsT r e m a -> HashConsT r e m b
forall a b c.
(a -> b -> c)
-> HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m c
forall r e (m :: Type -> Type).
Monad m =>
Functor (HashConsT r e m)
forall r e (m :: Type -> Type) a. Monad m => a -> HashConsT r e m a
forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m a
forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m (a -> b) -> HashConsT r e m a -> HashConsT r e m b
forall r e (m :: Type -> Type) a b c.
Monad m =>
(a -> b -> c)
-> HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m c
forall (f :: Type -> Type).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall r e (m :: Type -> Type) a. Monad m => a -> HashConsT r e m a
pure :: forall a. a -> HashConsT r e m a
$c<*> :: forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m (a -> b) -> HashConsT r e m a -> HashConsT r e m b
<*> :: forall a b.
HashConsT r e m (a -> b) -> HashConsT r e m a -> HashConsT r e m b
$cliftA2 :: forall r e (m :: Type -> Type) a b c.
Monad m =>
(a -> b -> c)
-> HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m c
liftA2 :: forall a b c.
(a -> b -> c)
-> HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m c
$c*> :: forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
*> :: forall a b.
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
$c<* :: forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m a
<* :: forall a b.
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m a
Applicative,
      -- | @since 1.0.0
      Applicative (HashConsT r e m)
Applicative (HashConsT r e m) =>
(forall a b.
 HashConsT r e m a -> (a -> HashConsT r e m b) -> HashConsT r e m b)
-> (forall a b.
    HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b)
-> (forall a. a -> HashConsT r e m a)
-> Monad (HashConsT r e m)
forall a. a -> HashConsT r e m a
forall a b.
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
forall a b.
HashConsT r e m a -> (a -> HashConsT r e m b) -> HashConsT r e m b
forall r e (m :: Type -> Type).
Monad m =>
Applicative (HashConsT r e m)
forall r e (m :: Type -> Type) a. Monad m => a -> HashConsT r e m a
forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> (a -> HashConsT r e m b) -> HashConsT r e m b
forall (m :: Type -> Type).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> (a -> HashConsT r e m b) -> HashConsT r e m b
>>= :: forall a b.
HashConsT r e m a -> (a -> HashConsT r e m b) -> HashConsT r e m b
$c>> :: forall r e (m :: Type -> Type) a b.
Monad m =>
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
>> :: forall a b.
HashConsT r e m a -> HashConsT r e m b -> HashConsT r e m b
$creturn :: forall r e (m :: Type -> Type) a. Monad m => a -> HashConsT r e m a
return :: forall a. a -> HashConsT r e m a
Monad
    )
    via (StateT (Bimap r e) m)
  deriving
    ( -- | @since 1.0.0
      (forall (m :: Type -> Type). Monad m => Monad (HashConsT r e m)) =>
(forall (m :: Type -> Type) a. Monad m => m a -> HashConsT r e m a)
-> MonadTrans (HashConsT r e)
forall r e (m :: Type -> Type). Monad m => Monad (HashConsT r e m)
forall r e (m :: Type -> Type) a.
Monad m =>
m a -> HashConsT r e m a
forall (m :: Type -> Type). Monad m => Monad (HashConsT r e m)
forall (m :: Type -> Type) a. Monad m => m a -> HashConsT r e m a
forall (t :: (Type -> Type) -> Type -> Type).
(forall (m :: Type -> Type). Monad m => Monad (t m)) =>
(forall (m :: Type -> Type) a. Monad m => m a -> t m a)
-> MonadTrans t
$clift :: forall r e (m :: Type -> Type) a.
Monad m =>
m a -> HashConsT r e m a
lift :: forall (m :: Type -> Type) a. Monad m => m a -> HashConsT r e m a
MonadTrans
    )
    via StateT (Bimap r e)

-- | Execute the computation described, returning both the result and the unique
-- pairings of @r@ and @e@ produced as part of it.
--
-- @since 1.0.0
runHashConsT ::
  forall (r :: Type) (e :: Type) (m :: Type -> Type) (a :: Type).
  HashConsT r e m a ->
  m (a, Bimap r e)
runHashConsT :: forall r e (m :: Type -> Type) a.
HashConsT r e m a -> m (a, Bimap r e)
runHashConsT (HashConsT StateT (Bimap r e) m a
comp) = StateT (Bimap r e) m a -> Bimap r e -> m (a, Bimap r e)
forall s (m :: Type -> Type) a. StateT s m a -> s -> m (a, s)
runStateT StateT (Bimap r e) m a
comp Bimap r e
forall a b. Bimap a b
Bimap.empty

-- | Given a value of type @e@, produce the unique value of type @r@ acting as a
-- reference to it. This @r@ will be cached, ensuring any future requests for
-- the reference for this value of type @e@ will be the same.
--
-- @since 1.0.0
hashCons ::
  forall (r :: Type) (e :: Type) (m :: Type -> Type).
  (Ord r, Ord e, Bounded r, Enum r, Monad m) =>
  e ->
  HashConsT r e m r
hashCons :: forall r e (m :: Type -> Type).
(Ord r, Ord e, Bounded r, Enum r, Monad m) =>
e -> HashConsT r e m r
hashCons e
x = StateT (Bimap r e) m r -> HashConsT r e m r
forall r e (m :: Type -> Type) a.
StateT (Bimap r e) m a -> HashConsT r e m a
HashConsT (StateT (Bimap r e) m r -> HashConsT r e m r)
-> StateT (Bimap r e) m r -> HashConsT r e m r
forall a b. (a -> b) -> a -> b
$ do
  Bimap r e
binds <- StateT (Bimap r e) m (Bimap r e)
forall s (m :: Type -> Type). MonadState s m => m s
get
  case e -> Bimap r e -> Maybe r
forall a b (m :: Type -> Type).
(Ord a, Ord b, MonadThrow m) =>
b -> Bimap a b -> m a
Bimap.lookupR e
x Bimap r e
binds of
    Maybe r
Nothing ->
      if Bimap r e -> Bool
forall a b. Bimap a b -> Bool
Bimap.null Bimap r e
binds
        then r
forall a. Bounded a => a
minBound r -> StateT (Bimap r e) m () -> StateT (Bimap r e) m r
forall a b. a -> StateT (Bimap r e) m b -> StateT (Bimap r e) m a
forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ (Bimap r e -> Bimap r e) -> StateT (Bimap r e) m ()
forall s (m :: Type -> Type). MonadState s m => (s -> s) -> m ()
modify (r -> e -> Bimap r e -> Bimap r e
forall a b. (Ord a, Ord b) => a -> b -> Bimap a b -> Bimap a b
Bimap.insert r
forall a. Bounded a => a
minBound e
x)
        else do
          let largestOldRef :: r
largestOldRef = (r, e) -> r
forall a b. (a, b) -> a
fst ((r, e) -> r) -> (Bimap r e -> (r, e)) -> Bimap r e -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bimap r e -> (r, e)
forall a b. Bimap a b -> (a, b)
Bimap.findMax (Bimap r e -> r) -> Bimap r e -> r
forall a b. (a -> b) -> a -> b
$ Bimap r e
binds
          let newRef :: r
newRef = r -> r
forall a. Enum a => a -> a
succ r
largestOldRef
          r
newRef r -> StateT (Bimap r e) m () -> StateT (Bimap r e) m r
forall a b. a -> StateT (Bimap r e) m b -> StateT (Bimap r e) m a
forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ (Bimap r e -> Bimap r e) -> StateT (Bimap r e) m ()
forall s (m :: Type -> Type). MonadState s m => (s -> s) -> m ()
modify (r -> e -> Bimap r e -> Bimap r e
forall a b. (Ord a, Ord b) => a -> b -> Bimap a b -> Bimap a b
Bimap.insert r
newRef e
x)
    Just r
ref -> r -> StateT (Bimap r e) m r
forall a. a -> StateT (Bimap r e) m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure r
ref

-- | Given a value of type @r@, fetch the cached @e@ value, if it exists.
--
-- @since 1.0.0
lookupRef_ ::
  forall (r :: Type) (e :: Type) (m :: Type -> Type).
  (Monad m, Ord e, Ord r) =>
  r ->
  HashConsT r e m (Maybe e)
lookupRef_ :: forall r e (m :: Type -> Type).
(Monad m, Ord e, Ord r) =>
r -> HashConsT r e m (Maybe e)
lookupRef_ r
r = StateT (Bimap r e) m (Maybe e) -> HashConsT r e m (Maybe e)
forall r e (m :: Type -> Type) a.
StateT (Bimap r e) m a -> HashConsT r e m a
HashConsT (r -> Bimap r e -> Maybe e
forall a b (m :: Type -> Type).
(Ord a, Ord b, MonadThrow m) =>
a -> Bimap a b -> m b
Bimap.lookup r
r (Bimap r e -> Maybe e)
-> StateT (Bimap r e) m (Bimap r e)
-> StateT (Bimap r e) m (Maybe e)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT (Bimap r e) m (Bimap r e)
forall s (m :: Type -> Type). MonadState s m => m s
get)

-- | An @mtl@-style capability type class for hash consing capability, using
-- references of type @r@ and values of type @e@.
--
-- = Laws
--
-- 1. @'refTo' x '>>' 'refTo' x@ @=@ @'refTo' x@
-- 2. @'liftA2' ('/=') ('refTo' x) ('refTo' y)@ @=@ @'refTo' x '*>' 'refTo' y '*>' 'pure' (x '/=' y)@
-- 3. @'refTo' x '>>=' (\\r -> 'lookupRef' r '>>=' (\\y -> 'pure' (y, r)))@ @=@ @('Just' x, ) '<$>' 'refTo' x@
--
-- @since 1.0.0
class
  (Eq e, Eq r, Monad m) =>
  MonadHashCons (r :: Type) (e :: Type) (m :: Type -> Type)
    | m -> e r
  where
  -- | Produce the unique value of type @r@ that acts as a reference for the
  -- given value of type @e@.
  --
  -- @since 1.0.0
  refTo :: e -> m r

  -- | Given a value of type @r@, fetch the cached value of type @e@.
  --
  -- @since 1.0.0
  lookupRef :: r -> m (Maybe e)

-- | @since 1.0.0
instance (Ord r, Ord e, Bounded r, Enum r, Monad m) => MonadHashCons r e (HashConsT r e m) where
  {-# INLINEABLE refTo #-}
  refTo :: e -> HashConsT r e m r
refTo = e -> HashConsT r e m r
forall r e (m :: Type -> Type).
(Ord r, Ord e, Bounded r, Enum r, Monad m) =>
e -> HashConsT r e m r
hashCons
  {-# INLINEABLE lookupRef #-}
  lookupRef :: r -> HashConsT r e m (Maybe e)
lookupRef = r -> HashConsT r e m (Maybe e)
forall r e (m :: Type -> Type).
(Monad m, Ord e, Ord r) =>
r -> HashConsT r e m (Maybe e)
lookupRef_

-- | @since 1.0.0
instance (Ord r, Ord e, MonadHashCons r e m) => MonadHashCons r e (MaybeT m) where
  {-# INLINEABLE refTo #-}
  refTo :: e -> MaybeT m r
refTo e
e = m r -> MaybeT m r
forall (m :: Type -> Type) a. Monad m => m a -> MaybeT m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (e -> m r
forall r e (m :: Type -> Type). MonadHashCons r e m => e -> m r
refTo e
e)
  {-# INLINEABLE lookupRef #-}
  lookupRef :: r -> MaybeT m (Maybe e)
lookupRef r
r = m (Maybe e) -> MaybeT m (Maybe e)
forall (m :: Type -> Type) a. Monad m => m a -> MaybeT m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (r -> m (Maybe e)
forall r e (m :: Type -> Type).
MonadHashCons r e m =>
r -> m (Maybe e)
lookupRef r
r)

-- | @since 1.0.0
instance (MonadHashCons r e m) => MonadHashCons r e (ReaderT r' m) where
  {-# INLINEABLE refTo #-}
  refTo :: e -> ReaderT r' m r
refTo e
e = m r -> ReaderT r' m r
forall (m :: Type -> Type) a. Monad m => m a -> ReaderT r' m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (e -> m r
forall r e (m :: Type -> Type). MonadHashCons r e m => e -> m r
refTo e
e)
  {-# INLINEABLE lookupRef #-}
  lookupRef :: r -> ReaderT r' m (Maybe e)
lookupRef r
r = m (Maybe e) -> ReaderT r' m (Maybe e)
forall (m :: Type -> Type) a. Monad m => m a -> ReaderT r' m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (r -> m (Maybe e)
forall r e (m :: Type -> Type).
MonadHashCons r e m =>
r -> m (Maybe e)
lookupRef r
r)

-- | @since 1.0.0
instance (MonadHashCons r e m) => MonadHashCons r e (StateT s m) where
  {-# INLINEABLE refTo #-}
  refTo :: e -> StateT s m r
refTo e
e = m r -> StateT s m r
forall (m :: Type -> Type) a. Monad m => m a -> StateT s m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (e -> m r
forall r e (m :: Type -> Type). MonadHashCons r e m => e -> m r
refTo e
e)
  {-# INLINEABLE lookupRef #-}
  lookupRef :: r -> StateT s m (Maybe e)
lookupRef r
r = m (Maybe e) -> StateT s m (Maybe e)
forall (m :: Type -> Type) a. Monad m => m a -> StateT s m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (r -> m (Maybe e)
forall r e (m :: Type -> Type).
MonadHashCons r e m =>
r -> m (Maybe e)
lookupRef r
r)

-- | @since 1.0.0
instance (MonadHashCons r e m) => MonadHashCons r e (WriterT w m) where
  {-# INLINEABLE refTo #-}
  refTo :: e -> WriterT w m r
refTo e
e = m r -> WriterT w m r
forall (m :: Type -> Type) a. Monad m => m a -> WriterT w m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (e -> m r
forall r e (m :: Type -> Type). MonadHashCons r e m => e -> m r
refTo e
e)
  {-# INLINEABLE lookupRef #-}
  lookupRef :: r -> WriterT w m (Maybe e)
lookupRef r
r = m (Maybe e) -> WriterT w m (Maybe e)
forall (m :: Type -> Type) a. Monad m => m a -> WriterT w m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (r -> m (Maybe e)
forall r e (m :: Type -> Type).
MonadHashCons r e m =>
r -> m (Maybe e)
lookupRef r
r)

-- | @since 1.0.0
instance (MonadHashCons r e m) => MonadHashCons r e (RWST r' w s m) where
  {-# INLINEABLE refTo #-}
  refTo :: e -> RWST r' w s m r
refTo e
e = m r -> RWST r' w s m r
forall (m :: Type -> Type) a. Monad m => m a -> RWST r' w s m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (e -> m r
forall r e (m :: Type -> Type). MonadHashCons r e m => e -> m r
refTo e
e)
  {-# INLINEABLE lookupRef #-}
  lookupRef :: r -> RWST r' w s m (Maybe e)
lookupRef r
r = m (Maybe e) -> RWST r' w s m (Maybe e)
forall (m :: Type -> Type) a. Monad m => m a -> RWST r' w s m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (r -> m (Maybe e)
forall r e (m :: Type -> Type).
MonadHashCons r e m =>
r -> m (Maybe e)
lookupRef r
r)

-- | @since 1.0.0
instance (MonadHashCons r e m) => MonadHashCons r e (ExceptT e' m) where
  {-# INLINEABLE refTo #-}
  refTo :: e -> ExceptT e' m r
refTo e
e = m r -> ExceptT e' m r
forall (m :: Type -> Type) a. Monad m => m a -> ExceptT e' m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (e -> m r
forall r e (m :: Type -> Type). MonadHashCons r e m => e -> m r
refTo e
e)
  {-# INLINEABLE lookupRef #-}
  lookupRef :: r -> ExceptT e' m (Maybe e)
lookupRef r
r = m (Maybe e) -> ExceptT e' m (Maybe e)
forall (m :: Type -> Type) a. Monad m => m a -> ExceptT e' m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (r -> m (Maybe e)
forall r e (m :: Type -> Type).
MonadHashCons r e m =>
r -> m (Maybe e)
lookupRef r
r)