{-# LANGUAGE BlockArguments             #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TemplateHaskell            #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.Effects.Fresh
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Polysemy effect for fresh name generation, compatible with the
-- unbound-generics library.
--
-----------------------------------------------------------------------------

module Disco.Effects.Fresh where

import           Disco.Effects.Counter
import           Disco.Names                           (QName, localName)
import           Polysemy
import           Polysemy.ConstraintAbsorber
import qualified Unbound.Generics.LocallyNameless      as U
import           Unbound.Generics.LocallyNameless.Name

-- | Fresh name generation effect, supporting raw generation of fresh
--   names, and opening binders with automatic freshening.  Simply
--   increments a global counter every time 'fresh' is called and
--   makes a variable with that numeric suffix.
data Fresh m a where
  Fresh  :: Name x -> Fresh m (Name x)

makeSem ''Fresh

-- | Dispatch the fresh name generation effect, starting at a given
--   integer.
runFresh' :: Integer -> Sem (Fresh ': r) a -> Sem r a
runFresh' :: Integer -> Sem (Fresh : r) a -> Sem r a
runFresh' Integer
i
  = Integer -> Sem (Counter : r) a -> Sem r a
forall (r :: [(* -> *) -> * -> *]) a.
Integer -> Sem (Counter : r) a -> Sem r a
runCounter' Integer
i
  (Sem (Counter : r) a -> Sem r a)
-> (Sem (Fresh : r) a -> Sem (Counter : r) a)
-> Sem (Fresh : r) a
-> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (rInitial :: [(* -> *) -> * -> *]) x.
 Fresh (Sem rInitial) x -> Sem (Counter : r) x)
-> Sem (Fresh : r) a -> Sem (Counter : r) a
forall (e1 :: (* -> *) -> * -> *) (e2 :: (* -> *) -> * -> *)
       (r :: [(* -> *) -> * -> *]) a.
FirstOrder e1 "reinterpret" =>
(forall (rInitial :: [(* -> *) -> * -> *]) x.
 e1 (Sem rInitial) x -> Sem (e2 : r) x)
-> Sem (e1 : r) a -> Sem (e2 : r) a
reinterpret \case
      Fresh x -> case Name x
x of
        Fn String
s Integer
_  -> String -> Integer -> Name x
forall a. String -> Integer -> Name a
Fn String
s (Integer -> Name x)
-> Sem (Counter : r) Integer -> Sem (Counter : r) (Name x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem (Counter : r) Integer
forall (r :: [(* -> *) -> * -> *]).
Member Counter r =>
Sem r Integer
next
        nm :: Name x
nm@Bn{} -> Name x -> Sem (Counter : r) (Name x)
forall (m :: * -> *) a. Monad m => a -> m a
return Name x
nm

    -- Above code copied from
    -- https://hackage.haskell.org/package/unbound-generics-0.4.1/docs/src/Unbound.Generics.LocallyNameless.Fresh.html ;
    -- see instance Monad m => Fresh (FreshMT m) .

    -- It turns out to make things much simpler to reimplement the
    -- Fresh effect ourselves in terms of a state effect, since then
    -- we can immediately dispatch it.  The alternative would be to
    -- implement it in terms of (Embed U.FreshM), but then we are
    -- stuck with that constraint.  Given the constraint-absorbing
    -- machinery below, just impementing the 'fresh' effect itself
    -- means we can then reuse other things from unbound-generics that
    -- depend on a Fresh constraint, such as the 'unbind' function
    -- below.

-- | Run a computation requiring fresh name generation, beginning with
--   0 for the initial freshly generated name.
runFresh :: Sem (Fresh ': r) a -> Sem r a
runFresh :: Sem (Fresh : r) a -> Sem r a
runFresh = Integer -> Sem (Fresh : r) a -> Sem r a
forall (r :: [(* -> *) -> * -> *]) a.
Integer -> Sem (Fresh : r) a -> Sem r a
runFresh' Integer
0

-- | Run a computation requiring fresh name generation, beginning with
--   1 instead of 0 for the initial freshly generated name.
runFresh1 :: Sem (Fresh ': r) a -> Sem r a
runFresh1 :: Sem (Fresh : r) a -> Sem r a
runFresh1 = Integer -> Sem (Fresh : r) a -> Sem r a
forall (r :: [(* -> *) -> * -> *]) a.
Integer -> Sem (Fresh : r) a -> Sem r a
runFresh' Integer
1

------------------------------------------------------------
-- Other functions

-- | Open a binder, automatically creating fresh names for the bound
--   variables.
unbind :: (Member Fresh r, U.Alpha p, U.Alpha t) => U.Bind p t -> Sem r (p, t)
unbind :: Bind p t -> Sem r (p, t)
unbind Bind p t
b = (Fresh (Sem r) => Sem r (p, t)) -> Sem r (p, t)
forall (r :: [(* -> *) -> * -> *]) a.
Member Fresh r =>
(Fresh (Sem r) => Sem r a) -> Sem r a
absorbFresh (Bind p t -> Sem r (p, t)
forall p t (m :: * -> *).
(Alpha p, Alpha t, Fresh m) =>
Bind p t -> m (p, t)
U.unbind Bind p t
b)

-- | Generate a fresh (local, free) qualified name based on a given
--   string.
freshQ :: (Member Fresh r) => String -> Sem r (QName a)
freshQ :: String -> Sem r (QName a)
freshQ String
s = Name a -> QName a
forall a. Name a -> QName a
localName (Name a -> QName a) -> Sem r (Name a) -> Sem r (QName a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name a -> Sem r (Name a)
forall (r :: [(* -> *) -> * -> *]) x.
Member Fresh r =>
Name x -> Sem r (Name x)
fresh (String -> Name a
forall a. String -> Name a
string2Name String
s)

------------------------------------------------------------
-- Machinery for absorbing MTL-style constraint.
-- See https://hackage.haskell.org/package/polysemy-zoo-0.7.0.1/docs/Polysemy-ConstraintAbsorber.html
-- Used https://hackage.haskell.org/package/polysemy-zoo-0.7.0.1/docs/src/Polysemy.ConstraintAbsorber.MonadState.html#absorbState as a template.

-- | Run a 'Sem' computation requiring a 'U.Fresh' constraint (from
--   the @unbound-generics@ library) in terms of an available 'Fresh'
--   effect.
absorbFresh :: Member Fresh r => (U.Fresh (Sem r) => Sem r a) -> Sem r a
absorbFresh :: (Fresh (Sem r) => Sem r a) -> Sem r a
absorbFresh = FreshDict (Sem r)
-> (forall s.
    Reifies s (FreshDict (Sem r)) :- Fresh (Action (Sem r) s))
-> (Fresh (Sem r) => Sem r a)
-> Sem r a
forall (p :: (* -> *) -> Constraint) (x :: (* -> *) -> * -> * -> *)
       d (r :: [(* -> *) -> * -> *]) a.
d
-> (forall s. Reifies s d :- p (x (Sem r) s))
-> (p (Sem r) => Sem r a)
-> Sem r a
absorbWithSem @U.Fresh @Action ((forall x. Name x -> Sem r (Name x)) -> FreshDict (Sem r)
forall (m :: * -> *).
(forall x. Name x -> m (Name x)) -> FreshDict m
FreshDict forall (r :: [(* -> *) -> * -> *]) x.
Member Fresh r =>
Name x -> Sem r (Name x)
forall x. Name x -> Sem r (Name x)
fresh) ((Reifies s (FreshDict (Sem r)) => Dict (Fresh (Action (Sem r) s)))
-> Reifies s (FreshDict (Sem r)) :- Fresh (Action (Sem r) s)
forall (a :: Constraint) (b :: Constraint). (a => Dict b) -> a :- b
Sub Reifies s (FreshDict (Sem r)) => Dict (Fresh (Action (Sem r) s))
forall (a :: Constraint). a => Dict a
Dict)
{-# INLINEABLE absorbFresh #-}

newtype FreshDict m = FreshDict { FreshDict m -> forall x. Name x -> m (Name x)
fresh_ :: forall x. Name x -> m (Name x) }

-- | Wrapper for a monadic action with phantom type parameter for reflection.
--   Locally defined so that the instance we are going to build with reflection
--   must be coherent, that is there cannot be orphans.
newtype Action m s' a = Action (m a)
  deriving (a -> Action m s' b -> Action m s' a
(a -> b) -> Action m s' a -> Action m s' b
(forall a b. (a -> b) -> Action m s' a -> Action m s' b)
-> (forall a b. a -> Action m s' b -> Action m s' a)
-> Functor (Action m s')
forall a b. a -> Action m s' b -> Action m s' a
forall a b. (a -> b) -> Action m s' a -> Action m s' b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
a -> Action m s' b -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
(a -> b) -> Action m s' a -> Action m s' b
<$ :: a -> Action m s' b -> Action m s' a
$c<$ :: forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
a -> Action m s' b -> Action m s' a
fmap :: (a -> b) -> Action m s' a -> Action m s' b
$cfmap :: forall (m :: * -> *) k (s' :: k) a b.
Functor m =>
(a -> b) -> Action m s' a -> Action m s' b
Functor, Functor (Action m s')
a -> Action m s' a
Functor (Action m s')
-> (forall a. a -> Action m s' a)
-> (forall a b.
    Action m s' (a -> b) -> Action m s' a -> Action m s' b)
-> (forall a b c.
    (a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c)
-> (forall a b. Action m s' a -> Action m s' b -> Action m s' b)
-> (forall a b. Action m s' a -> Action m s' b -> Action m s' a)
-> Applicative (Action m s')
Action m s' a -> Action m s' b -> Action m s' b
Action m s' a -> Action m s' b -> Action m s' a
Action m s' (a -> b) -> Action m s' a -> Action m s' b
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
forall a. a -> Action m s' a
forall a b. Action m s' a -> Action m s' b -> Action m s' a
forall a b. Action m s' a -> Action m s' b -> Action m s' b
forall a b. Action m s' (a -> b) -> Action m s' a -> Action m s' b
forall a b c.
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
forall (f :: * -> *).
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
forall (m :: * -> *) k (s' :: k).
Applicative m =>
Functor (Action m s')
forall (m :: * -> *) k (s' :: k) a.
Applicative m =>
a -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' b
forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' (a -> b) -> Action m s' a -> Action m s' b
forall (m :: * -> *) k (s' :: k) a b c.
Applicative m =>
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
<* :: Action m s' a -> Action m s' b -> Action m s' a
$c<* :: forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' a
*> :: Action m s' a -> Action m s' b -> Action m s' b
$c*> :: forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' a -> Action m s' b -> Action m s' b
liftA2 :: (a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
$cliftA2 :: forall (m :: * -> *) k (s' :: k) a b c.
Applicative m =>
(a -> b -> c) -> Action m s' a -> Action m s' b -> Action m s' c
<*> :: Action m s' (a -> b) -> Action m s' a -> Action m s' b
$c<*> :: forall (m :: * -> *) k (s' :: k) a b.
Applicative m =>
Action m s' (a -> b) -> Action m s' a -> Action m s' b
pure :: a -> Action m s' a
$cpure :: forall (m :: * -> *) k (s' :: k) a.
Applicative m =>
a -> Action m s' a
$cp1Applicative :: forall (m :: * -> *) k (s' :: k).
Applicative m =>
Functor (Action m s')
Applicative, Applicative (Action m s')
a -> Action m s' a
Applicative (Action m s')
-> (forall a b.
    Action m s' a -> (a -> Action m s' b) -> Action m s' b)
-> (forall a b. Action m s' a -> Action m s' b -> Action m s' b)
-> (forall a. a -> Action m s' a)
-> Monad (Action m s')
Action m s' a -> (a -> Action m s' b) -> Action m s' b
Action m s' a -> Action m s' b -> Action m s' b
forall a. a -> Action m s' a
forall a b. Action m s' a -> Action m s' b -> Action m s' b
forall a b. Action m s' a -> (a -> Action m s' b) -> Action m s' b
forall (m :: * -> *).
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
forall (m :: * -> *) k (s' :: k).
Monad m =>
Applicative (Action m s')
forall (m :: * -> *) k (s' :: k) a. Monad m => a -> Action m s' a
forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> Action m s' b -> Action m s' b
forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> (a -> Action m s' b) -> Action m s' b
return :: a -> Action m s' a
$creturn :: forall (m :: * -> *) k (s' :: k) a. Monad m => a -> Action m s' a
>> :: Action m s' a -> Action m s' b -> Action m s' b
$c>> :: forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> Action m s' b -> Action m s' b
>>= :: Action m s' a -> (a -> Action m s' b) -> Action m s' b
$c>>= :: forall (m :: * -> *) k (s' :: k) a b.
Monad m =>
Action m s' a -> (a -> Action m s' b) -> Action m s' b
$cp1Monad :: forall (m :: * -> *) k (s' :: k).
Monad m =>
Applicative (Action m s')
Monad)

instance ( Monad m
         , Reifies s' (FreshDict m)
         ) => U.Fresh (Action m s') where
  fresh :: Name a -> Action m s' (Name a)
fresh Name a
x = m (Name a) -> Action m s' (Name a)
forall k k (m :: k -> *) (s' :: k) (a :: k). m a -> Action m s' a
Action (m (Name a) -> Action m s' (Name a))
-> m (Name a) -> Action m s' (Name a)
forall a b. (a -> b) -> a -> b
$ FreshDict m -> Name a -> m (Name a)
forall (m :: * -> *). FreshDict m -> forall x. Name x -> m (Name x)
fresh_ (Proxy s' -> FreshDict m
forall k (s :: k) a (proxy :: k -> *). Reifies s a => proxy s -> a
reflect (Proxy s' -> FreshDict m) -> Proxy s' -> FreshDict m
forall a b. (a -> b) -> a -> b
$ Proxy s'
forall k (t :: k). Proxy t
Proxy @s') Name a
x
  {-# INLINEABLE fresh #-}