{-# LANGUAGE CPP             #-}
{-# LANGUAGE TemplateHaskell #-}
module PlutusTx.IsData.TH (unstableMakeIsData, makeIsDataIndexed) where

import Data.Foldable
import Data.Traversable

import Language.Haskell.TH qualified as TH
import Language.Haskell.TH.Datatype qualified as TH
import PlutusTx.ErrorCodes

import PlutusTx.Applicative qualified as PlutusTx

import PlutusTx.Builtins as Builtins
import PlutusTx.Builtins.Internal qualified as BI
import PlutusTx.IsData.Class
import PlutusTx.Trace (traceError)

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

toDataClause :: (TH.ConstructorInfo, Int) -> TH.Q TH.Clause
toDataClause :: (ConstructorInfo, Int) -> Q Clause
toDataClause (TH.ConstructorInfo{constructorName :: ConstructorInfo -> Name
TH.constructorName=Name
name, constructorFields :: ConstructorInfo -> [Type]
TH.constructorFields=[Type]
argTys}, Int
index) = do
    [Name]
argNames <- [Type] -> (Type -> Q Name) -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Type]
argTys ((Type -> Q Name) -> Q [Name]) -> (Type -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ \Type
_ -> String -> Q Name
TH.newName String
"arg"
    let argsList :: ExpQ
argsList = (Name -> ExpQ -> ExpQ) -> ExpQ -> [Name] -> ExpQ
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Name
v ExpQ
e -> [| BI.mkCons (toBuiltinData $(TH.varE v)) $e |]) [| BI.mkNilData BI.unitval |] [Name]
argNames
    let app :: ExpQ
app = [| BI.mkConstr index $argsList |]
    [PatQ] -> BodyQ -> [DecQ] -> Q Clause
TH.clause [Name -> [PatQ] -> PatQ
TH.conP Name
name ((Name -> PatQ) -> [Name] -> [PatQ]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> PatQ
TH.varP [Name]
argNames)] (ExpQ -> BodyQ
TH.normalB ExpQ
app) []

toDataClauses :: [(TH.ConstructorInfo, Int)] -> [TH.Q TH.Clause]
toDataClauses :: [(ConstructorInfo, Int)] -> [Q Clause]
toDataClauses [(ConstructorInfo, Int)]
indexedCons = (ConstructorInfo, Int) -> Q Clause
toDataClause ((ConstructorInfo, Int) -> Q Clause)
-> [(ConstructorInfo, Int)] -> [Q Clause]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(ConstructorInfo, Int)]
indexedCons

reconstructCase :: (TH.ConstructorInfo, Int) -> TH.Q TH.Exp -> TH.Q TH.Exp -> TH.Q TH.Exp -> TH.Q TH.Exp
reconstructCase :: (ConstructorInfo, Int) -> ExpQ -> ExpQ -> ExpQ -> ExpQ
reconstructCase (TH.ConstructorInfo{constructorName :: ConstructorInfo -> Name
TH.constructorName=Name
name, constructorFields :: ConstructorInfo -> [Type]
TH.constructorFields=[Type]
argTys}, Int
index) ExpQ
ixExpr ExpQ
argsExpr ExpQ
kont = do
    [Name]
argNames <- [Type] -> (Type -> Q Name) -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Type]
argTys ((Type -> Q Name) -> Q [Name]) -> (Type -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ \Type
_ -> String -> Q Name
TH.newName String
"arg"

    -- Applicatively build the constructor application, assuming that all the arguments are in scope
    let app :: ExpQ
app = (ExpQ -> Name -> ExpQ) -> ExpQ -> [Name] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\ExpQ
h Name
v -> [| $h PlutusTx.<*> fromBuiltinData $(TH.varE v) |]) [| PlutusTx.pure $(TH.conE name) |] [Name]
argNames

    -- Takes a list of argument names, and safely takes one element off the list for each, binding it to the name.
    -- Finally, invokes 'app'.
    let handleList :: [TH.Name] -> TH.Q TH.Exp -> TH.Q TH.Exp
        handleList :: [Name] -> ExpQ -> ExpQ
handleList [] ExpQ
lExp = [| matchList $lExp $app (\_ _ -> Nothing) |]
        handleList (Name
argName:[Name]
rest) ExpQ
lExp = do
            Name
tailName <- String -> Q Name
TH.newName String
"t"
            [|
             let consCase = \ $(TH.varP argName) $(TH.varP tailName) -> $(handleList rest (TH.varE tailName))
             in matchList $lExp Nothing consCase
             |]
    -- Check that the index matches the expected one, otherwise fallthrough to 'kont'
    let body :: ExpQ
body =
            [|
                let indexMatchCase = $(handleList argNames argsExpr)
                    fallthrough = $kont
                in BI.ifThenElse ($ixExpr `BI.equalsInteger` (index :: Integer)) (const indexMatchCase) (const fallthrough) BI.unitval
            |]
    ExpQ
body

fromDataClause :: [(TH.ConstructorInfo, Int)] -> TH.Q TH.Clause
fromDataClause :: [(ConstructorInfo, Int)] -> Q Clause
fromDataClause [(ConstructorInfo, Int)]
indexedCons = do
    Name
dName <- String -> Q Name
TH.newName String
"d"
    Name
indexName <- String -> Q Name
TH.newName String
"index"
    Name
argsName <- String -> Q Name
TH.newName String
"args0"
    -- Call the clause for each constructor, falling through to the next one, until we get to the end in which case we return 'Nothing'
    let cases :: ExpQ
cases =
            (ExpQ -> (ConstructorInfo, Int) -> ExpQ)
-> ExpQ -> [(ConstructorInfo, Int)] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
            (\ExpQ
kont (ConstructorInfo, Int)
ixCon -> (ConstructorInfo, Int) -> ExpQ -> ExpQ -> ExpQ -> ExpQ
reconstructCase (ConstructorInfo, Int)
ixCon (Name -> ExpQ
TH.varE Name
indexName) (Name -> ExpQ
TH.varE Name
argsName) ExpQ
kont)
            [| Nothing |]
            [(ConstructorInfo, Int)]
indexedCons
    let body :: ExpQ
body =
          [|
            let constrMatchCase = \ $(TH.varP indexName) $(TH.varP argsName) -> $cases
            in matchData' $(TH.varE dName) constrMatchCase (const Nothing) (const Nothing) (const Nothing) (const Nothing)
          |]
    [PatQ] -> BodyQ -> [DecQ] -> Q Clause
TH.clause [Name -> PatQ
TH.varP Name
dName] (ExpQ -> BodyQ
TH.normalB ExpQ
body) []

unsafeReconstructCase :: (TH.ConstructorInfo, Int) -> TH.Q TH.Exp -> TH.Q TH.Exp -> TH.Q TH.Exp -> TH.Q TH.Exp
unsafeReconstructCase :: (ConstructorInfo, Int) -> ExpQ -> ExpQ -> ExpQ -> ExpQ
unsafeReconstructCase (TH.ConstructorInfo{constructorName :: ConstructorInfo -> Name
TH.constructorName=Name
name, constructorFields :: ConstructorInfo -> [Type]
TH.constructorFields=[Type]
argTys}, Int
index) ExpQ
ixExpr ExpQ
argsExpr ExpQ
kont = do
    [Name]
argNames <- [Type] -> (Type -> Q Name) -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Type]
argTys ((Type -> Q Name) -> Q [Name]) -> (Type -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ \Type
_ -> String -> Q Name
TH.newName String
"arg"

    -- Build the constructor application, assuming that all the arguments are in scope
    let app :: ExpQ
app = (ExpQ -> Name -> ExpQ) -> ExpQ -> [Name] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\ExpQ
h Name
v -> [| $h (unsafeFromBuiltinData $(TH.varE v)) |]) (Name -> ExpQ
TH.conE Name
name) [Name]
argNames

    -- Takes a list of argument names, and takes one element off the list for each, binding it to the name.
    -- Finally, invokes 'app'.
    let handleList :: [TH.Name] -> TH.Q TH.Exp -> TH.Q TH.Exp
        handleList :: [Name] -> ExpQ -> ExpQ
handleList [] ExpQ
_ = [| $app |]
        handleList (Name
argName:[Name]
rest) ExpQ
lExp = do
            [|
             let
                 t = $lExp
                 $(TH.varP argName) = BI.head t
             in $(handleList rest [| BI.tail t |])
             |]
    -- Check that the index matches the expected one, otherwise fallthrough to 'kont'
    let body :: ExpQ
body =
            [|
                let indexMatchCase = $(handleList argNames argsExpr)
                    fallthrough = $kont
                in BI.ifThenElse ($ixExpr `BI.equalsInteger` (index :: Integer)) (const indexMatchCase) (const fallthrough) BI.unitval
            |]
    ExpQ
body

unsafeFromDataClause :: [(TH.ConstructorInfo, Int)] -> TH.Q TH.Clause
unsafeFromDataClause :: [(ConstructorInfo, Int)] -> Q Clause
unsafeFromDataClause [(ConstructorInfo, Int)]
indexedCons = do
    Name
dName <- String -> Q Name
TH.newName String
"d"
    Name
indexName <- String -> Q Name
TH.newName String
"index"
    Name
tupName <- String -> Q Name
TH.newName String
"tup"
    -- Call the clause for each constructor, falling through to the next one, until we get to the end in which case we call 'error'
    let cases :: ExpQ
cases =
            (ExpQ -> (ConstructorInfo, Int) -> ExpQ)
-> ExpQ -> [(ConstructorInfo, Int)] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
            (\ExpQ
kont (ConstructorInfo, Int)
ixCon -> (ConstructorInfo, Int) -> ExpQ -> ExpQ -> ExpQ -> ExpQ
unsafeReconstructCase (ConstructorInfo, Int)
ixCon (Name -> ExpQ
TH.varE Name
indexName) [| BI.snd $(TH.varE tupName) |] ExpQ
kont)
            [| traceError reconstructCaseError |]
            [(ConstructorInfo, Int)]
indexedCons
    let body :: ExpQ
body =
          [|
            let $(TH.varP tupName) = BI.unsafeDataAsConstr $(TH.varE dName)
                $(TH.varP indexName) = BI.fst $(TH.varE tupName)
            in $cases
          |]
    [PatQ] -> BodyQ -> [DecQ] -> Q Clause
TH.clause [Name -> PatQ
TH.varP Name
dName] (ExpQ -> BodyQ
TH.normalB ExpQ
body) []

defaultIndex :: TH.Name -> TH.Q [(TH.Name, Int)]
defaultIndex :: Name -> Q [(Name, Int)]
defaultIndex Name
name = do
    DatatypeInfo
info <- Name -> Q DatatypeInfo
TH.reifyDatatype Name
name
    [(Name, Int)] -> Q [(Name, Int)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(Name, Int)] -> Q [(Name, Int)])
-> [(Name, Int)] -> Q [(Name, Int)]
forall a b. (a -> b) -> a -> b
$ [Name] -> [Int] -> [(Name, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (ConstructorInfo -> Name
TH.constructorName (ConstructorInfo -> Name) -> [ConstructorInfo] -> [Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DatatypeInfo -> [ConstructorInfo]
TH.datatypeCons DatatypeInfo
info) [Int
0..]

-- | Generate a 'FromData' and a 'ToData' instance for a type. This may not be stable in the face of constructor additions,
-- renamings, etc. Use 'makeIsDataIndexed' if you need stability.
unstableMakeIsData :: TH.Name -> TH.Q [TH.Dec]
unstableMakeIsData :: Name -> Q [Dec]
unstableMakeIsData Name
name = Name -> [(Name, Int)] -> Q [Dec]
makeIsDataIndexed Name
name ([(Name, Int)] -> Q [Dec]) -> Q [(Name, Int)] -> Q [Dec]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> Q [(Name, Int)]
defaultIndex Name
name

-- | Generate a 'FromData' and a 'ToData' instance for a type, using an explicit mapping of constructor names to indices. Use
-- this for types where you need to keep the representation stable.
makeIsDataIndexed :: TH.Name -> [(TH.Name, Int)] -> TH.Q [TH.Dec]
makeIsDataIndexed :: Name -> [(Name, Int)] -> Q [Dec]
makeIsDataIndexed Name
name [(Name, Int)]
indices = do

    DatatypeInfo
info <- Name -> Q DatatypeInfo
TH.reifyDatatype Name
name
    let appliedType :: Type
appliedType = DatatypeInfo -> Type
TH.datatypeType DatatypeInfo
info

    [(ConstructorInfo, Int)]
indexedCons <- [ConstructorInfo]
-> (ConstructorInfo -> Q (ConstructorInfo, Int))
-> Q [(ConstructorInfo, Int)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for (DatatypeInfo -> [ConstructorInfo]
TH.datatypeCons DatatypeInfo
info) ((ConstructorInfo -> Q (ConstructorInfo, Int))
 -> Q [(ConstructorInfo, Int)])
-> (ConstructorInfo -> Q (ConstructorInfo, Int))
-> Q [(ConstructorInfo, Int)]
forall a b. (a -> b) -> a -> b
$ \ConstructorInfo
c -> case Name -> [(Name, Int)] -> Maybe Int
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (ConstructorInfo -> Name
TH.constructorName ConstructorInfo
c) [(Name, Int)]
indices of
            Just Int
i  -> (ConstructorInfo, Int) -> Q (ConstructorInfo, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ConstructorInfo
c, Int
i)
            Maybe Int
Nothing -> String -> Q (ConstructorInfo, Int)
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q (ConstructorInfo, Int))
-> String -> Q (ConstructorInfo, Int)
forall a b. (a -> b) -> a -> b
$ String
"No index given for constructor" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show (ConstructorInfo -> Name
TH.constructorName ConstructorInfo
c)

    Dec
toDataInst <- do
        let constraints :: [Type]
constraints = (TyVarBndr -> Type) -> [TyVarBndr] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\TyVarBndr
t -> Name -> [Type] -> Type
TH.classPred ''ToData [Name -> Type
TH.VarT (TyVarBndr -> Name
tyvarbndrName TyVarBndr
t)]) (DatatypeInfo -> [TyVarBndr]
TH.datatypeVars DatatypeInfo
info)
        Dec
toDataDecl <- Name -> [Q Clause] -> DecQ
TH.funD 'toBuiltinData ([(ConstructorInfo, Int)] -> [Q Clause]
toDataClauses [(ConstructorInfo, Int)]
indexedCons)
        Dec
toDataPrag <- Name -> Inline -> RuleMatch -> Phases -> DecQ
TH.pragInlD 'toBuiltinData Inline
TH.Inlinable RuleMatch
TH.FunLike Phases
TH.AllPhases
        Dec -> DecQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec -> DecQ) -> Dec -> DecQ
forall a b. (a -> b) -> a -> b
$ Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
TH.InstanceD Maybe Overlap
forall a. Maybe a
Nothing [Type]
constraints (Name -> [Type] -> Type
TH.classPred ''ToData [Type
appliedType]) [Dec
toDataPrag, Dec
toDataDecl]

    Dec
fromDataInst <- do
        let constraints :: [Type]
constraints = (TyVarBndr -> Type) -> [TyVarBndr] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\TyVarBndr
t -> Name -> [Type] -> Type
TH.classPred ''FromData [Name -> Type
TH.VarT (TyVarBndr -> Name
tyvarbndrName TyVarBndr
t)]) (DatatypeInfo -> [TyVarBndr]
TH.datatypeVars DatatypeInfo
info)
        Dec
fromDataDecl <- Name -> [Q Clause] -> DecQ
TH.funD 'fromBuiltinData [[(ConstructorInfo, Int)] -> Q Clause
fromDataClause [(ConstructorInfo, Int)]
indexedCons]
        Dec
fromDataPrag <- Name -> Inline -> RuleMatch -> Phases -> DecQ
TH.pragInlD 'fromBuiltinData Inline
TH.Inlinable RuleMatch
TH.FunLike Phases
TH.AllPhases
        Dec -> DecQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec -> DecQ) -> Dec -> DecQ
forall a b. (a -> b) -> a -> b
$ Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
TH.InstanceD Maybe Overlap
forall a. Maybe a
Nothing [Type]
constraints (Name -> [Type] -> Type
TH.classPred ''FromData [Type
appliedType]) [Dec
fromDataPrag, Dec
fromDataDecl]

    Dec
unsafeFromDataInst <- do
        let constraints :: [Type]
constraints = (TyVarBndr -> Type) -> [TyVarBndr] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\TyVarBndr
t -> Name -> [Type] -> Type
TH.classPred ''UnsafeFromData [Name -> Type
TH.VarT (TyVarBndr -> Name
tyvarbndrName TyVarBndr
t)]) (DatatypeInfo -> [TyVarBndr]
TH.datatypeVars DatatypeInfo
info)
        Dec
unsafeFromDataDecl <- Name -> [Q Clause] -> DecQ
TH.funD 'unsafeFromBuiltinData [[(ConstructorInfo, Int)] -> Q Clause
unsafeFromDataClause [(ConstructorInfo, Int)]
indexedCons]
        Dec
unsafeFromDataPrag <- Name -> Inline -> RuleMatch -> Phases -> DecQ
TH.pragInlD 'unsafeFromBuiltinData Inline
TH.Inlinable RuleMatch
TH.FunLike Phases
TH.AllPhases
        Dec -> DecQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec -> DecQ) -> Dec -> DecQ
forall a b. (a -> b) -> a -> b
$ Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
TH.InstanceD Maybe Overlap
forall a. Maybe a
Nothing [Type]
constraints (Name -> [Type] -> Type
TH.classPred ''UnsafeFromData [Type
appliedType]) [Dec
unsafeFromDataPrag, Dec
unsafeFromDataDecl]

    [Dec] -> Q [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Dec
toDataInst, Dec
fromDataInst, Dec
unsafeFromDataInst]
    where
#if MIN_VERSION_template_haskell(2,17,0)
        tyvarbndrName (TH.PlainTV n _)    = n
        tyvarbndrName (TH.KindedTV n _ _) = n
#else
        tyvarbndrName :: TyVarBndr -> Name
tyvarbndrName (TH.PlainTV Name
n)      = Name
n
        tyvarbndrName (TH.KindedTV Name
n Type
_)   = Name
n
#endif