-- editorconfig-checker-disable-file
{-# LANGUAGE BangPatterns    #-}
{-# 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 <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Type]
argTys forall a b. (a -> b) -> a -> b
$ \Type
_ -> forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"arg"
    let argsList :: Q Exp
argsList = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Name
v Q Exp
e -> [| BI.mkCons (toBuiltinData $(TH.varE v)) $e |]) [| BI.mkNilData BI.unitval |] [Name]
argNames
    let app :: Q Exp
app = [| BI.mkConstr index $argsList |]
    forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
TH.clause [forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
TH.conP Name
name (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: * -> *). Quote m => Name -> m Pat
TH.varP [Name]
argNames)] (forall (m :: * -> *). Quote m => m Exp -> m Body
TH.normalB Q Exp
app) []

toDataClauses :: [(TH.ConstructorInfo, Int)] -> [TH.Q TH.Clause]
toDataClauses :: [(ConstructorInfo, Int)] -> [Q Clause]
toDataClauses [(ConstructorInfo, Int)]
indexedCons = (ConstructorInfo, Int) -> Q Clause
toDataClause 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) -> Q Exp -> Q Exp -> Q Exp -> Q Exp
reconstructCase (TH.ConstructorInfo{constructorName :: ConstructorInfo -> Name
TH.constructorName=Name
name, constructorFields :: ConstructorInfo -> [Type]
TH.constructorFields=[Type]
argTys}, Int
index) Q Exp
ixExpr Q Exp
argsExpr Q Exp
kont = do
    [Name]
argNames <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Type]
argTys forall a b. (a -> b) -> a -> b
$ \Type
_ -> forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"arg"

    -- Applicatively build the constructor application, assuming that all the arguments are in scope
    let app :: Q Exp
app = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Q Exp
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] -> Q Exp -> Q Exp
handleList [] Q Exp
lExp = [| matchList $lExp $app (\_ _ -> Nothing) |]
        handleList (Name
argName:[Name]
rest) Q Exp
lExp = do
            Name
tailName <- forall (m :: * -> *). Quote m => String -> m 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 :: Q Exp
body =
            [|
                let !indexMatchCase = $(handleList argNames argsExpr)
                    !fallthrough = $kont
                in BI.ifThenElse ($ixExpr `BI.equalsInteger` (index :: Integer)) (const indexMatchCase) (const fallthrough) BI.unitval
            |]
    Q Exp
body

fromDataClause :: [(TH.ConstructorInfo, Int)] -> TH.Q TH.Clause
fromDataClause :: [(ConstructorInfo, Int)] -> Q Clause
fromDataClause [(ConstructorInfo, Int)]
indexedCons = do
    Name
dName <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"d"
    Name
indexName <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"index"
    Name
argsName <- forall (m :: * -> *). Quote m => String -> m 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 :: Q Exp
cases =
            forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
            (\Q Exp
kont (ConstructorInfo, Int)
ixCon -> (ConstructorInfo, Int) -> Q Exp -> Q Exp -> Q Exp -> Q Exp
reconstructCase (ConstructorInfo, Int)
ixCon (forall (m :: * -> *). Quote m => Name -> m Exp
TH.varE Name
indexName) (forall (m :: * -> *). Quote m => Name -> m Exp
TH.varE Name
argsName) Q Exp
kont)
            [| Nothing |]
            [(ConstructorInfo, Int)]
indexedCons
    let body :: Q Exp
body =
          [|
            let !constrMatchCase = \ $(TH.varP indexName) $(TH.varP argsName) -> $cases
            in matchData' $(TH.varE dName) constrMatchCase (const Nothing) (const Nothing) (const Nothing) (const Nothing)
          |]
    forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
TH.clause [forall (m :: * -> *). Quote m => Name -> m Pat
TH.varP Name
dName] (forall (m :: * -> *). Quote m => m Exp -> m Body
TH.normalB Q Exp
body) []

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

    -- Build the constructor application, assuming that all the arguments are in scope
    let app :: Q Exp
app = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Q Exp
h Name
v -> [| $h (unsafeFromBuiltinData $(TH.varE v)) |]) (forall (m :: * -> *). Quote m => Name -> m Exp
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] -> Q Exp -> Q Exp
handleList [] Q Exp
_ = [| $app |]
        handleList (Name
argName:[Name]
rest) Q Exp
lExp = do
            [|
             let
                 !t = $lExp
                 $(TH.bangP $ 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 :: Q Exp
body =
            [|
                let !indexMatchCase = $(handleList argNames argsExpr)
                    !fallthrough = $kont
                in BI.ifThenElse ($ixExpr `BI.equalsInteger` (index :: Integer)) (const indexMatchCase) (const fallthrough) BI.unitval
            |]
    Q Exp
body

unsafeFromDataClause :: [(TH.ConstructorInfo, Int)] -> TH.Q TH.Clause
unsafeFromDataClause :: [(ConstructorInfo, Int)] -> Q Clause
unsafeFromDataClause [(ConstructorInfo, Int)]
indexedCons = do
    Name
dName <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"d"
    Name
indexName <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"index"
    Name
tupName <- forall (m :: * -> *). Quote m => String -> m 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 :: Q Exp
cases =
            forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
            (\Q Exp
kont (ConstructorInfo, Int)
ixCon -> (ConstructorInfo, Int) -> Q Exp -> Q Exp -> Q Exp -> Q Exp
unsafeReconstructCase (ConstructorInfo, Int)
ixCon (forall (m :: * -> *). Quote m => Name -> m Exp
TH.varE Name
indexName) [| BI.snd $(TH.varE tupName) |] Q Exp
kont)
            [| traceError reconstructCaseError |]
            [(ConstructorInfo, Int)]
indexedCons
    let body :: Q Exp
body =
          [|
            let $(TH.bangP $ TH.varP tupName) = BI.unsafeDataAsConstr $(TH.varE dName)
                $(TH.bangP $ TH.varP indexName) = BI.fst $(TH.varE tupName)
            in $cases
          |]
    forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
TH.clause [forall (m :: * -> *). Quote m => Name -> m Pat
TH.varP Name
dName] (forall (m :: * -> *). Quote m => m Exp -> m Body
TH.normalB Q Exp
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
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (ConstructorInfo -> Name
TH.constructorName 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 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 <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for (DatatypeInfo -> [ConstructorInfo]
TH.datatypeCons DatatypeInfo
info) forall a b. (a -> b) -> a -> b
$ \ConstructorInfo
c -> case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (ConstructorInfo -> Name
TH.constructorName ConstructorInfo
c) [(Name, Int)]
indices of
            Just Int
i  -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (ConstructorInfo
c, Int
i)
            Maybe Int
Nothing -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"No index given for constructor" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (ConstructorInfo -> Name
TH.constructorName ConstructorInfo
c)

    Dec
toDataInst <- do
        let constraints :: [Type]
constraints = 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 (forall {flag}. TyVarBndr flag -> Name
tyvarbndrName TyVarBndr ()
t)]) (DatatypeInfo -> [TyVarBndr ()]
TH.datatypeVars DatatypeInfo
info)
        Dec
toDataDecl <- forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
TH.funD 'toBuiltinData ([(ConstructorInfo, Int)] -> [Q Clause]
toDataClauses [(ConstructorInfo, Int)]
indexedCons)
        Dec
toDataPrag <- forall (m :: * -> *).
Quote m =>
Name -> Inline -> RuleMatch -> Phases -> m Dec
TH.pragInlD 'toBuiltinData Inline
TH.Inlinable RuleMatch
TH.FunLike Phases
TH.AllPhases
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
TH.InstanceD 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 = 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 (forall {flag}. TyVarBndr flag -> Name
tyvarbndrName TyVarBndr ()
t)]) (DatatypeInfo -> [TyVarBndr ()]
TH.datatypeVars DatatypeInfo
info)
        Dec
fromDataDecl <- forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
TH.funD 'fromBuiltinData [[(ConstructorInfo, Int)] -> Q Clause
fromDataClause [(ConstructorInfo, Int)]
indexedCons]
        Dec
fromDataPrag <- forall (m :: * -> *).
Quote m =>
Name -> Inline -> RuleMatch -> Phases -> m Dec
TH.pragInlD 'fromBuiltinData Inline
TH.Inlinable RuleMatch
TH.FunLike Phases
TH.AllPhases
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
TH.InstanceD 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 = 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 (forall {flag}. TyVarBndr flag -> Name
tyvarbndrName TyVarBndr ()
t)]) (DatatypeInfo -> [TyVarBndr ()]
TH.datatypeVars DatatypeInfo
info)
        Dec
unsafeFromDataDecl <- forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
TH.funD 'unsafeFromBuiltinData [[(ConstructorInfo, Int)] -> Q Clause
unsafeFromDataClause [(ConstructorInfo, Int)]
indexedCons]
        Dec
unsafeFromDataPrag <- forall (m :: * -> *).
Quote m =>
Name -> Inline -> RuleMatch -> Phases -> m Dec
TH.pragInlD 'unsafeFromBuiltinData Inline
TH.Inlinable RuleMatch
TH.FunLike Phases
TH.AllPhases
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
TH.InstanceD forall a. Maybe a
Nothing [Type]
constraints (Name -> [Type] -> Type
TH.classPred ''UnsafeFromData [Type
appliedType]) [Dec
unsafeFromDataPrag, Dec
unsafeFromDataDecl]

    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 :: TyVarBndr flag -> Name
tyvarbndrName (TH.PlainTV Name
n flag
_)    = Name
n
        tyvarbndrName (TH.KindedTV Name
n flag
_ Type
_) = Name
n
#else
        tyvarbndrName (TH.PlainTV n)      = n
        tyvarbndrName (TH.KindedTV n _)   = n
#endif