{-# LANGUAGE DefaultSignatures     #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances  #-}
-- just for the type equality constraint
{-# LANGUAGE GADTs                 #-}

module PlutusCore.Quote
    ( runQuoteT
    , runQuote
    , freshUnique
    , freshName
    , freshTyName
    , freshenName
    , freshenTyName
    , QuoteT (..)
    , Quote
    , MonadQuote
    , FreshState
    , liftQuote
    , markNonFreshBelow
    , markNonFresh
    , markNonFreshMax
    ) where

import PlutusPrelude

import PlutusCore.Name

import Control.Monad.Except
import Control.Monad.Morph as MM
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Trans.Maybe
import Control.Monad.Writer
import Data.Functor.Identity
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text qualified as Text
import Hedgehog (GenT, PropertyT)

-- | The state contains the "next" 'Unique' that should be used for a name
type FreshState = Unique

emptyFreshState :: FreshState
emptyFreshState :: FreshState
emptyFreshState = Int -> FreshState
Unique Int
0

-- | The "quotation" monad transformer. Within this monad you can do safe construction of PLC terms using quasiquotation,
-- fresh-name generation, and parsing.
newtype QuoteT m a = QuoteT { QuoteT m a -> StateT FreshState m a
unQuoteT :: StateT FreshState m a }
    deriving newtype (a -> QuoteT m b -> QuoteT m a
(a -> b) -> QuoteT m a -> QuoteT m b
(forall a b. (a -> b) -> QuoteT m a -> QuoteT m b)
-> (forall a b. a -> QuoteT m b -> QuoteT m a)
-> Functor (QuoteT m)
forall a b. a -> QuoteT m b -> QuoteT m a
forall a b. (a -> b) -> QuoteT m a -> QuoteT m b
forall (m :: * -> *) a b.
Functor m =>
a -> QuoteT m b -> QuoteT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> QuoteT m a -> QuoteT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> QuoteT m b -> QuoteT m a
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> QuoteT m b -> QuoteT m a
fmap :: (a -> b) -> QuoteT m a -> QuoteT m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> QuoteT m a -> QuoteT m b
Functor, Functor (QuoteT m)
a -> QuoteT m a
Functor (QuoteT m)
-> (forall a. a -> QuoteT m a)
-> (forall a b. QuoteT m (a -> b) -> QuoteT m a -> QuoteT m b)
-> (forall a b c.
    (a -> b -> c) -> QuoteT m a -> QuoteT m b -> QuoteT m c)
-> (forall a b. QuoteT m a -> QuoteT m b -> QuoteT m b)
-> (forall a b. QuoteT m a -> QuoteT m b -> QuoteT m a)
-> Applicative (QuoteT m)
QuoteT m a -> QuoteT m b -> QuoteT m b
QuoteT m a -> QuoteT m b -> QuoteT m a
QuoteT m (a -> b) -> QuoteT m a -> QuoteT m b
(a -> b -> c) -> QuoteT m a -> QuoteT m b -> QuoteT m c
forall a. a -> QuoteT m a
forall a b. QuoteT m a -> QuoteT m b -> QuoteT m a
forall a b. QuoteT m a -> QuoteT m b -> QuoteT m b
forall a b. QuoteT m (a -> b) -> QuoteT m a -> QuoteT m b
forall a b c.
(a -> b -> c) -> QuoteT m a -> QuoteT m b -> QuoteT m c
forall (m :: * -> *). Monad m => Functor (QuoteT m)
forall (m :: * -> *) a. Monad m => a -> QuoteT m a
forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> QuoteT m b -> QuoteT m a
forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> QuoteT m b -> QuoteT m b
forall (m :: * -> *) a b.
Monad m =>
QuoteT m (a -> b) -> QuoteT m a -> QuoteT m b
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> QuoteT m a -> QuoteT m b -> QuoteT m 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
<* :: QuoteT m a -> QuoteT m b -> QuoteT m a
$c<* :: forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> QuoteT m b -> QuoteT m a
*> :: QuoteT m a -> QuoteT m b -> QuoteT m b
$c*> :: forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> QuoteT m b -> QuoteT m b
liftA2 :: (a -> b -> c) -> QuoteT m a -> QuoteT m b -> QuoteT m c
$cliftA2 :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> QuoteT m a -> QuoteT m b -> QuoteT m c
<*> :: QuoteT m (a -> b) -> QuoteT m a -> QuoteT m b
$c<*> :: forall (m :: * -> *) a b.
Monad m =>
QuoteT m (a -> b) -> QuoteT m a -> QuoteT m b
pure :: a -> QuoteT m a
$cpure :: forall (m :: * -> *) a. Monad m => a -> QuoteT m a
$cp1Applicative :: forall (m :: * -> *). Monad m => Functor (QuoteT m)
Applicative, Applicative (QuoteT m)
a -> QuoteT m a
Applicative (QuoteT m)
-> (forall a b. QuoteT m a -> (a -> QuoteT m b) -> QuoteT m b)
-> (forall a b. QuoteT m a -> QuoteT m b -> QuoteT m b)
-> (forall a. a -> QuoteT m a)
-> Monad (QuoteT m)
QuoteT m a -> (a -> QuoteT m b) -> QuoteT m b
QuoteT m a -> QuoteT m b -> QuoteT m b
forall a. a -> QuoteT m a
forall a b. QuoteT m a -> QuoteT m b -> QuoteT m b
forall a b. QuoteT m a -> (a -> QuoteT m b) -> QuoteT m b
forall (m :: * -> *). Monad m => Applicative (QuoteT m)
forall (m :: * -> *) a. Monad m => a -> QuoteT m a
forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> QuoteT m b -> QuoteT m b
forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> (a -> QuoteT m b) -> QuoteT m 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
return :: a -> QuoteT m a
$creturn :: forall (m :: * -> *) a. Monad m => a -> QuoteT m a
>> :: QuoteT m a -> QuoteT m b -> QuoteT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> QuoteT m b -> QuoteT m b
>>= :: QuoteT m a -> (a -> QuoteT m b) -> QuoteT m b
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
QuoteT m a -> (a -> QuoteT m b) -> QuoteT m b
$cp1Monad :: forall (m :: * -> *). Monad m => Applicative (QuoteT m)
Monad, m a -> QuoteT m a
(forall (m :: * -> *) a. Monad m => m a -> QuoteT m a)
-> MonadTrans QuoteT
forall (m :: * -> *) a. Monad m => m a -> QuoteT m a
forall (t :: (* -> *) -> * -> *).
(forall (m :: * -> *) a. Monad m => m a -> t m a) -> MonadTrans t
lift :: m a -> QuoteT m a
$clift :: forall (m :: * -> *) a. Monad m => m a -> QuoteT m a
MonadTrans, (forall a. m a -> n a) -> QuoteT m b -> QuoteT n b
(forall (m :: * -> *) (n :: * -> *) b.
 Monad m =>
 (forall a. m a -> n a) -> QuoteT m b -> QuoteT n b)
-> MFunctor QuoteT
forall k (t :: (* -> *) -> k -> *).
(forall (m :: * -> *) (n :: * -> *) (b :: k).
 Monad m =>
 (forall a. m a -> n a) -> t m b -> t n b)
-> MFunctor t
forall (m :: * -> *) (n :: * -> *) b.
Monad m =>
(forall a. m a -> n a) -> QuoteT m b -> QuoteT n b
hoist :: (forall a. m a -> n a) -> QuoteT m b -> QuoteT n b
$choist :: forall (m :: * -> *) (n :: * -> *) b.
Monad m =>
(forall a. m a -> n a) -> QuoteT m b -> QuoteT n b
MM.MFunctor, MonadError e, MonadReader r, Monad (QuoteT m)
Monad (QuoteT m)
-> (forall a. IO a -> QuoteT m a) -> MonadIO (QuoteT m)
IO a -> QuoteT m a
forall a. IO a -> QuoteT m a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
forall (m :: * -> *). MonadIO m => Monad (QuoteT m)
forall (m :: * -> *) a. MonadIO m => IO a -> QuoteT m a
liftIO :: IO a -> QuoteT m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> QuoteT m a
$cp1MonadIO :: forall (m :: * -> *). MonadIO m => Monad (QuoteT m)
MonadIO, MonadWriter w)

-- Need to write this by hand, deriving wants to derive the one for DefState
instance MonadState s m => MonadState s (QuoteT m) where
    get :: QuoteT m s
get = m s -> QuoteT m s
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m s
forall s (m :: * -> *). MonadState s m => m s
get
    put :: s -> QuoteT m ()
put = m () -> QuoteT m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> QuoteT m ()) -> (s -> m ()) -> s -> QuoteT m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put
    state :: (s -> (a, s)) -> QuoteT m a
state = m a -> QuoteT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> QuoteT m a)
-> ((s -> (a, s)) -> m a) -> (s -> (a, s)) -> QuoteT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (s -> (a, s)) -> m a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state

-- | A monad that allows lifting of quoted expressions.
class Monad m => MonadQuote m where
    liftQuote :: Quote a -> m a
    -- This means we don't have to implement it when we're writing an instance for a MonadTrans monad. We can't just
    -- add an instance declaration for that, because it overlaps with the base instance.
    default liftQuote :: (MonadQuote n, MonadTrans t, t n ~ m) => Quote a -> m a
    liftQuote = n a -> t n a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (n a -> t n a) -> (Quote a -> n a) -> Quote a -> t n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Quote a -> n a
forall (m :: * -> *) a. MonadQuote m => Quote a -> m a
liftQuote

instance (Monad m) => MonadQuote (QuoteT m) where
    liftQuote :: Quote a -> QuoteT m a
liftQuote = (forall a. Identity a -> m a) -> Quote a -> QuoteT m a
forall k (t :: (* -> *) -> k -> *) (m :: * -> *) (n :: * -> *)
       (b :: k).
(MFunctor t, Monad m) =>
(forall a. m a -> n a) -> t m b -> t n b
MM.hoist (a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> (Identity a -> a) -> Identity a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity a -> a
forall a. Identity a -> a
runIdentity)

instance MonadQuote m => MonadQuote (StateT s m)
instance MonadQuote m => MonadQuote (MaybeT m)
instance MonadQuote m => MonadQuote (ExceptT e m)
instance MonadQuote m => MonadQuote (ReaderT r m)
instance MonadQuote m => MonadQuote (GenT m)
instance MonadQuote m => MonadQuote (PropertyT m)

-- | Run a quote from an empty identifier state. Note that the resulting term cannot necessarily
-- be safely combined with other terms - that should happen inside 'QuoteT'.
runQuoteT ::  Monad m => QuoteT m a -> m a
runQuoteT :: QuoteT m a -> m a
runQuoteT QuoteT m a
q = StateT FreshState m a -> FreshState -> m a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (QuoteT m a -> StateT FreshState m a
forall (m :: * -> *) a. QuoteT m a -> StateT FreshState m a
unQuoteT QuoteT m a
q) FreshState
emptyFreshState

-- | A non-transformer version of 'QuoteT'.
type Quote = QuoteT Identity

-- | See 'runQuoteT'.
runQuote :: Quote a -> a
runQuote :: Quote a -> a
runQuote = Identity a -> a
forall a. Identity a -> a
runIdentity (Identity a -> a) -> (Quote a -> Identity a) -> Quote a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Quote a -> Identity a
forall (m :: * -> *) a. Monad m => QuoteT m a -> m a
runQuoteT

-- | Get a fresh 'Unique'.
freshUnique :: MonadQuote m => m Unique
freshUnique :: m FreshState
freshUnique = Quote FreshState -> m FreshState
forall (m :: * -> *) a. MonadQuote m => Quote a -> m a
liftQuote (Quote FreshState -> m FreshState)
-> Quote FreshState -> m FreshState
forall a b. (a -> b) -> a -> b
$ do
    FreshState
nextU <- StateT FreshState Identity FreshState -> Quote FreshState
forall (m :: * -> *) a. StateT FreshState m a -> QuoteT m a
QuoteT StateT FreshState Identity FreshState
forall s (m :: * -> *). MonadState s m => m s
get
    StateT FreshState Identity () -> QuoteT Identity ()
forall (m :: * -> *) a. StateT FreshState m a -> QuoteT m a
QuoteT (StateT FreshState Identity () -> QuoteT Identity ())
-> StateT FreshState Identity () -> QuoteT Identity ()
forall a b. (a -> b) -> a -> b
$ FreshState -> StateT FreshState Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (FreshState -> StateT FreshState Identity ())
-> FreshState -> StateT FreshState Identity ()
forall a b. (a -> b) -> a -> b
$ Int -> FreshState
Unique (FreshState -> Int
unUnique FreshState
nextU Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    FreshState -> Quote FreshState
forall (f :: * -> *) a. Applicative f => a -> f a
pure FreshState
nextU

-- | Get a fresh 'Name', given the annotation and the 'Text.Text' name.
freshName :: MonadQuote m => Text.Text -> m Name
freshName :: Text -> m Name
freshName Text
str = Text -> FreshState -> Name
Name Text
str (FreshState -> Name) -> m FreshState -> m Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m FreshState
forall (m :: * -> *). MonadQuote m => m FreshState
freshUnique

-- | Make a copy of the given 'Name' that is distinct from the old one.
freshenName :: MonadQuote m => Name -> m Name
freshenName :: Name -> m Name
freshenName (Name Text
str FreshState
_) = Text -> FreshState -> Name
Name Text
str (FreshState -> Name) -> m FreshState -> m Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m FreshState
forall (m :: * -> *). MonadQuote m => m FreshState
freshUnique

-- | Get a fresh 'TyName', given the annotation and the 'Text.Text' name.
freshTyName :: MonadQuote m => Text.Text -> m TyName
freshTyName :: Text -> m TyName
freshTyName = (Name -> TyName) -> m Name -> m TyName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> TyName
TyName (m Name -> m TyName) -> (Text -> m Name) -> Text -> m TyName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> m Name
forall (m :: * -> *). MonadQuote m => Text -> m Name
freshName

-- | Make a copy of the given 'TyName' that is distinct from the old one.
freshenTyName :: MonadQuote m => TyName -> m TyName
freshenTyName :: TyName -> m TyName
freshenTyName (TyName Name
name) = Name -> TyName
TyName (Name -> TyName) -> m Name -> m TyName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> m Name
forall (m :: * -> *). MonadQuote m => Name -> m Name
freshenName Name
name

-- | Mark a all 'Unique's less than the given 'Unique' as used, so they will not be generated in future.
markNonFreshBelow :: MonadQuote m => Unique -> m ()
markNonFreshBelow :: FreshState -> m ()
markNonFreshBelow = QuoteT Identity () -> m ()
forall (m :: * -> *) a. MonadQuote m => Quote a -> m a
liftQuote (QuoteT Identity () -> m ())
-> (FreshState -> QuoteT Identity ()) -> FreshState -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StateT FreshState Identity () -> QuoteT Identity ()
forall (m :: * -> *) a. StateT FreshState m a -> QuoteT m a
QuoteT (StateT FreshState Identity () -> QuoteT Identity ())
-> (FreshState -> StateT FreshState Identity ())
-> FreshState
-> QuoteT Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FreshState -> FreshState) -> StateT FreshState Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FreshState -> FreshState) -> StateT FreshState Identity ())
-> (FreshState -> FreshState -> FreshState)
-> FreshState
-> StateT FreshState Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FreshState -> FreshState -> FreshState
forall a. Ord a => a -> a -> a
max

-- | Mark a given 'Unique' (and implicitly all 'Unique's less than it) as used, so they will not be generated in future.
markNonFresh :: MonadQuote m => Unique -> m ()
markNonFresh :: FreshState -> m ()
markNonFresh = FreshState -> m ()
forall (m :: * -> *). MonadQuote m => FreshState -> m ()
markNonFreshBelow (FreshState -> m ())
-> (FreshState -> FreshState) -> FreshState -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FreshState -> FreshState
forall a. Enum a => a -> a
succ

-- | Mark the maximal 'Unique' from a set of 'Unique's (and implicitly all 'Unique's less than it)
-- as used, so they will not be generated in future.
markNonFreshMax :: MonadQuote m => Set Unique -> m ()
markNonFreshMax :: Set FreshState -> m ()
markNonFreshMax = FreshState -> m ()
forall (m :: * -> *). MonadQuote m => FreshState -> m ()
markNonFresh (FreshState -> m ())
-> (Set FreshState -> FreshState) -> Set FreshState -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FreshState -> Maybe FreshState -> FreshState
forall a. a -> Maybe a -> a
fromMaybe (Int -> FreshState
Unique Int
0) (Maybe FreshState -> FreshState)
-> (Set FreshState -> Maybe FreshState)
-> Set FreshState
-> FreshState
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set FreshState -> Maybe FreshState
forall a. Set a -> Maybe a
Set.lookupMax