Skip to content

Instantly share code, notes, and snippets.

@glguy
Created February 24, 2014 18:48
Show Gist options
  • Save glguy/9194384 to your computer and use it in GitHub Desktop.
Save glguy/9194384 to your computer and use it in GitHub Desktop.
Stand-alone implementation of the automatic Each instance generation
{-# 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"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment