{-# LANGUAGE TemplateHaskellQuotes #-}

-- | Derive a 'Lazy' instance for a given data type.
module Test.Mutagen.TH.Lazy
  ( deriveLazy
  )
where

import Language.Haskell.TH
  ( Lit (..)
  , Name
  , Q
  , newName
  )
import Language.Haskell.TH.Desugar
  ( DClause (..)
  , DCon (..)
  , DDec (..)
  , DExp (..)
  , DLetDec (..)
  , DPat (..)
  , DType (..)
  )
import Test.Mutagen
  ( Lazy
  , lazyNode
  , __evaluated__
  )
import Test.Mutagen.TH.Util
  ( applyTyVars
  , createDPat
  , dConName
  , dTyVarBndrName
  , mkConDExp
  , reifyTypeDef
  )

{-------------------------------------------------------------------------------
-- * Deriving Lazy instances
-------------------------------------------------------------------------------}

-- | Derive a 'Lazy' instance for the given data type.
deriveLazy :: Name -> [Name] -> Q [DDec]
deriveLazy :: Name -> [Name] -> Q [DDec]
deriveLazy Name
typeName [Name]
ignoredCons = do
  -- Reify the type definition
  (dtvbs, dcons) <- Name -> Q ([DTyVarBndrVis], [DCon])
reifyTypeDef Name
typeName
  -- Apply the context type variables to the type name to get 'Type'-kinded
  -- target type to derive the instance for
  let targetType = Name -> [DTyVarBndrVis] -> DType
applyTyVars Name
typeName [DTyVarBndrVis]
dtvbs
  -- Keep only the constructors that are not explicitly ignored by the user
  let wantedCons = (DCon -> Bool) -> [DCon] -> [DCon]
forall a. (a -> Bool) -> [a] -> [a]
filter (\DCon
con -> DCon -> Name
dConName DCon
con Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
ignoredCons) [DCon]
dcons
  -- Derive 'lazyNode' for each constructor separately
  insClauses <- mapM deriveLazyNode wantedCons
  -- Build the Lazy instance
  let insCxt = [Name -> DType
DConT ''Lazy DType -> DType -> DType
`DAppT` Name -> DType
DVarT (DTyVarBndrVis -> Name
dTyVarBndrName DTyVarBndrVis
tvb) | DTyVarBndrVis
tvb <- [DTyVarBndrVis]
dtvbs]
  let insTy = Name -> DType
DConT ''Lazy DType -> DType -> DType
`DAppT` DType
targetType
  let insBody = [DLetDec -> DDec
DLetDec (Name -> [DClause] -> DLetDec
DFunD 'lazyNode [DClause]
insClauses)]
  return [DInstanceD Nothing Nothing insCxt insTy insBody]

-- | Derive a single clause of the 'lazyNode' method for a given constructor.
deriveLazyNode :: DCon -> Q DClause
deriveLazyNode :: DCon -> Q DClause
deriveLazyNode DCon
con = do
  acc_ <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"acc"
  (vars, pat) <- createDPat con

  let lazyNodeExprs =
        [ Name -> DExp
DVarE 'lazyNode
            DExp -> DExp -> DExp
`DAppE` (Name -> DExp
DConE '(:) DExp -> DExp -> DExp
`DAppE` Lit -> DExp
DLitE (Integer -> Lit
IntegerL Integer
idx) DExp -> DExp -> DExp
`DAppE` Name -> DExp
DVarE Name
acc_)
            DExp -> DExp -> DExp
`DAppE` Name -> DExp
DVarE Name
var
        | (Integer
idx, Name
var) <- [Integer] -> [Name] -> [(Integer, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0 ..] [Name]
vars
        ]
  let clauseBody =
        Name -> DExp
DVarE '__evaluated__
          DExp -> DExp -> DExp
`DAppE` Name -> DExp
DVarE Name
acc_
          DExp -> DExp -> DExp
`DAppE` Name -> [DExp] -> DExp
mkConDExp (DCon -> Name
dConName DCon
con) [DExp]
lazyNodeExprs
  return (DClause [DVarP acc_, pat] clauseBody)