{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE TemplateHaskell #-}
module PlutusTx.Lift.THUtils where

import PlutusIR
import PlutusIR.Compiler.Names

import PlutusCore.Quote

import Control.Monad

import Data.Text qualified as T

import Language.Haskell.TH qualified as TH
import Language.Haskell.TH.Datatype qualified as TH
import Language.Haskell.TH.Syntax qualified as TH

-- We do not use qualified import because the whole module contains off-chain code
import Prelude as Haskell

-- | Very nearly the same as 'TH.showName', but doesn't print uniques, since we don't need to
-- incorporate them into our names.
showName :: TH.Name -> T.Text
showName :: Name -> Text
showName Name
n = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ case Name
n of
    TH.Name OccName
occ NameFlavour
TH.NameS         -> OccName -> String
TH.occString OccName
occ
    TH.Name OccName
occ (TH.NameQ ModName
m)     -> ModName -> String
TH.modString ModName
m String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ OccName -> String
TH.occString OccName
occ
    TH.Name OccName
occ (TH.NameG NameSpace
_ PkgName
_ ModName
m) -> ModName -> String
TH.modString ModName
m String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ OccName -> String
TH.occString OccName
occ
    TH.Name OccName
occ (TH.NameU Uniq
_)     -> OccName -> String
TH.occString OccName
occ
    TH.Name OccName
occ (TH.NameL Uniq
_)     -> OccName -> String
TH.occString OccName
occ

-- | Normalize a type, in particular getting rid of things like 'TH.ListT' in favour of applications of the actual name.
normalizeType :: TH.Type -> TH.Type
normalizeType :: Type -> Type
normalizeType = \case
    TH.ForallT [TyVarBndr]
b Cxt
c Type
t       -> [TyVarBndr] -> Cxt -> Type -> Type
TH.ForallT [TyVarBndr]
b Cxt
c (Type -> Type
normalizeType Type
t)
    TH.AppT Type
t1 Type
t2          -> Type -> Type -> Type
TH.AppT (Type -> Type
normalizeType Type
t1) (Type -> Type
normalizeType Type
t2)
    TH.SigT Type
t Type
_            -> Type -> Type
normalizeType Type
t
    TH.InfixT Type
t1 Name
n Type
t2      -> Name -> Type
TH.ConT Name
n Type -> Type -> Type
`TH.AppT` Type -> Type
normalizeType Type
t1 Type -> Type -> Type
`TH.AppT` Type -> Type
normalizeType Type
t2
    TH.UInfixT Type
t1 Name
n Type
t2     -> Name -> Type
TH.ConT Name
n Type -> Type -> Type
`TH.AppT` Type -> Type
normalizeType Type
t1 Type -> Type -> Type
`TH.AppT` Type -> Type
normalizeType Type
t2
    TH.ParensT Type
t           -> Type -> Type
normalizeType Type
t
    Type
TH.ListT               -> Name -> Type
TH.ConT ''[]
    TH.TupleT Int
arity        -> Name -> Type
TH.ConT (Int -> Name
TH.tupleTypeName Int
arity)
    TH.UnboxedTupleT Int
arity -> Name -> Type
TH.ConT (Int -> Name
TH.unboxedTupleTypeName Int
arity)
    TH.UnboxedSumT Int
arity   -> Name -> Type
TH.ConT (Int -> Name
TH.unboxedSumTypeName Int
arity)
    -- some of this stuff probably should be normalized (like tuples) but I don't know quite what to do
    Type
t                      -> Type
t

requireExtension :: TH.Extension -> TH.Q ()
requireExtension :: Extension -> Q ()
requireExtension Extension
ext = do
    Bool
enabled <- Extension -> Q Bool
TH.isExtEnabled Extension
ext
    Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
enabled (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q ()) -> String -> Q ()
forall a b. (a -> b) -> a -> b
$ String
"Extension must be enabled: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Extension -> String
forall a. Show a => a -> String
show Extension
ext

mkTyVarDecl :: (MonadQuote m) => TH.Name -> Kind () -> m (TH.Name, TyVarDecl TyName ())
mkTyVarDecl :: Name -> Kind () -> m (Name, TyVarDecl TyName ())
mkTyVarDecl Name
name Kind ()
kind = do
    TyName
tyName <- Text -> m TyName
forall (m :: * -> *). MonadQuote m => Text -> m TyName
safeFreshTyName (Text -> m TyName) -> Text -> m TyName
forall a b. (a -> b) -> a -> b
$ Name -> Text
showName Name
name
    (Name, TyVarDecl TyName ()) -> m (Name, TyVarDecl TyName ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name
name, () -> TyName -> Kind () -> TyVarDecl TyName ()
forall tyname ann.
ann -> tyname -> Kind ann -> TyVarDecl tyname ann
TyVarDecl () TyName
tyName Kind ()
kind)

isNewtype :: TH.DatatypeInfo -> Bool
isNewtype :: DatatypeInfo -> Bool
isNewtype TH.DatatypeInfo{datatypeVariant :: DatatypeInfo -> DatatypeVariant
TH.datatypeVariant=DatatypeVariant
variant} = case DatatypeVariant
variant of
    DatatypeVariant
TH.Newtype -> Bool
True
    DatatypeVariant
_          -> Bool
False

-- | "Safe" wrapper around 'TH.listE' for typed TH.
tyListE :: [TH.TExpQ a] -> TH.TExpQ [a]
tyListE :: [TExpQ a] -> TExpQ [a]
tyListE [TExpQ a]
texps = Q Exp -> TExpQ [a]
forall a. Q Exp -> Q (TExp a)
TH.unsafeTExpCoerce [| $(TH.listE (fmap TH.unTypeQ texps)) |]