{-# LANGUAGE TemplateHaskellQuotes #-}

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

import Control.Monad (forM)
import Language.Haskell.TH
  ( Name
  , Q
  , newName
  )
import Language.Haskell.TH.Desugar
  ( DClause (..)
  , DCon (..)
  , DDec (..)
  , DExp (..)
  , DLetDec (..)
  , DMatch (..)
  , DPat (..)
  , DType (..)
  , dCaseE
  )
import Test.Mutagen.Fragment
  ( Fragmentable
  , fragmentize
  , singleton
  )
import Test.Mutagen.TH.Util
  ( applyTyVars
  , createDPat
  , dConName
  , dTyVarBndrName
  , reifyTypeDef
  )

{-------------------------------------------------------------------------------
-- * Deriving Fragmentable instances
-------------------------------------------------------------------------------}

-- | Derive a 'Fragmentable' instance for the given data type.
deriveFragmentable :: Name -> [Name] -> Q [DDec]
deriveFragmentable :: Name -> [Name] -> Q [DDec]
deriveFragmentable 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 ignored
  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 the fragmentize clause
  insClause <- deriveFragmentize wantedCons
  -- Build the Mutable instance
  let insCxt = [Name -> DType
DConT ''Fragmentable DType -> DType -> DType
`DAppT` Name -> DType
DVarT (DTyVarBndrVis -> Name
dTyVarBndrName DTyVarBndrVis
tvb) | DTyVarBndrVis
tvb <- [DTyVarBndrVis]
dtvbs]
  let insTy = Name -> DType
DConT ''Fragmentable DType -> DType -> DType
`DAppT` DType
targetType
  let insBody = [DLetDec -> DDec
DLetDec (Name -> [DClause] -> DLetDec
DFunD 'fragmentize [DClause
insClause])]
  return [DInstanceD Nothing Nothing insCxt insTy insBody]

-- | Derive the 'fragmentize' clause for the given constructors.
--
-- This one is a bit tricky: the TH desugarer removes as (@) patterns, so the
-- only way to have a variable binding the full input is to introduce it as a
-- variable and then perform a case statement to find the actual constructor.
-- Reconstructing the input using the LHS pattern doesn't work because GHC
-- cannot guarantee its type to be the same as the one being matched against.
-- One can solve this also using TypeApplications and ScopedTypeVariables, but
-- it is unnecessary and we want to avoid needing extra extensions as much as
-- possible.
deriveFragmentize :: [DCon] -> Q DClause
deriveFragmentize :: [DCon] -> Q DClause
deriveFragmentize [DCon]
cons = do
  input_ <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"input"
  let inputFragment = Name -> DExp
DVarE 'singleton DExp -> DExp -> DExp
`DAppE` Name -> DExp
DVarE Name
input_
  let mappendExp DExp
x DExp
y = Name -> DExp
DVarE '(<>) DExp -> DExp -> DExp
`DAppE` DExp
x DExp -> DExp -> DExp
`DAppE` DExp
y
  caseCons <-
    forM cons $ \DCon
con -> do
      (vars, pat) <- DCon -> Q ([Name], DPat)
createDPat DCon
con
      let fragmentizeExprs =
            [ Name -> DExp
DVarE 'fragmentize DExp -> DExp -> DExp
`DAppE` Name -> DExp
DVarE Name
var
            | Name
var <- [Name]
vars
            ]
      let caseBody = (DExp -> DExp -> DExp) -> DExp -> [DExp] -> DExp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl DExp -> DExp -> DExp
mappendExp DExp
inputFragment [DExp]
fragmentizeExprs
      return (DMatch pat caseBody)
  let clauseBody = DExp -> [DMatch] -> DExp
dCaseE (Name -> DExp
DVarE Name
input_) [DMatch]
caseCons
  return (DClause [DVarP input_] clauseBody)