Stand-alone implementation of the automatic Each instance generation

Each.hs

{-# LANGUAGE TemplateHaskell #-}

module Each (makeEachInstance) where

import Control.Applicative (Applicative, (<*>), (<$>), pure)
import Control.Lens        (Each(each), _2, folded, lengthOf, to, toListOf, view)
import Control.Monad       (replicateM)
import Data.Map            (Map, fromList)
import Data.Traversable    (sequenceA)
import Language.Haskell.TH.Lens (HasTypeVars, conFields, name, substTypeVars, typeVars)
import Language.Haskell.TH

-- | Automatically generate 'Each' instances for types whose constructors
-- contain only fields of a single type. (The field type defaults to '()' for
-- types without fields).
makeEachInstance :: Name -> DecsQ
makeEachInstance typeName =
  do info <- reify typeName
     case info of
       TyConI (DataD    [] n bs cs _) -> buildInstance (makeType n bs) cs
       TyConI (NewtypeD [] n bs c  _) -> buildInstance (makeType n bs) [c]
       _                              -> fail badNameMsg
  where
  makeType :: Name -> [TyVarBndr] -> Type
  makeType typeConName binders =
     foldl AppT (ConT typeConName) (toListOf (typeVars . to VarT) binders)


-- | Generate an 'Each' instance for a given type and the constructors for that type.
buildInstance :: Type -> [Con] -> DecsQ
buildInstance s cs =

  do a     <- singleFieldType cs
     (t,b) <- computeTypeChange s a

     let eachClauses
           | null cs = [emptyClause]
           | otherwise = buildClause <$> cs

     sequence
       [ instanceD
           (cxt [])
           [t|Each $(pure s) $(pure t) $(pure a) $(pure b)|]
           [funD 'each eachClauses]
       ]

-- | Given the outer type and inner type compute a new outer,inner
-- pair by replacing the free variables in the inner type with
-- fresh variables.
computeTypeChange :: Type -> Type -> Q (Type, Type)
computeTypeChange s a =
  do m <- freshMap a
     let t = substTypeVars m s
         b = substTypeVars m a
     return (t, b)

-- | Create a clause that passes the given bottom along as the result.
-- This is used for datatypes without constructors.
-- each _ x = pure x
emptyClause :: ClauseQ
emptyClause =
  do x <- newName "x"
     clause [wildP, varP x] (normalB [|pure $(varE x)|]) []

-- | Build a single clause for the 'each' implementation given a constructor.
buildClause :: Con -> ClauseQ
buildClause con = case lengthOf conFields con of
                    0 -> zeroClause
                    n -> nonzeroClause n
  where
  conName = view name con
  conExpr = conE conName
  conPat  = conP conName

  -- each _ Con             = pure Con
  zeroClause = clause [wildP, conPat []] (normalB [| pure $conExpr |]) []

  -- each f (Con x y z ...) = Con <$> f x <*> f y <*> f z ...
  nonzeroClause n =
    do f    <-                newName "f"
       x:xs <- replicateM n $ newName "x"

       let fExpr  = varE f
           xExpr  = varE x
           xExprs = varE <$> xs

           base     = [| $conExpr <$> $fExpr $xExpr |]
           step l r = [| $l       <*> $fExpr $r     |]
           body     = foldl step base xExprs

       clause [varP f, conP conName (varP <$> x:xs)] (normalB body) []


-- | Returns the type which is equal across all fields of
-- all constructors if one exists. Types without fields
-- return a default of '()'
singleFieldType :: [Con] -> TypeQ
singleFieldType cs
  = case toListOf (folded . conFields . _2) cs of
      []                  -> tupleT 0
      t:ts | all (==t) ts -> return t
           | otherwise    -> fail badFieldTypesMsg


-- | Given a value 't' with type variables, compute a map from
-- every free variable in 't' to a fresh variable.
freshMap :: HasTypeVars t => t -> Q (Map Name Name)
freshMap t = sequenceA (fromList pairMs)
  where
  names  = toListOf typeVars t
  pairMs = [ (n, newName (nameBase n)) | n <- names ]

badFieldTypesMsg :: String
badFieldTypesMsg = "Automatic Each instances requires fields have single type"

badNameMsg :: String
badNameMsg = "Automatic Each instance generation needs the name of a datatype or newtype"