module Control.Newtype.TH ( mkNewTypes ) where
import Control.Newtype ( Newtype(pack, unpack) )
import Control.Applicative ((<$>))
import Control.Arrow ((&&&), (***))
import Data.Function ( on )
import Data.List ( groupBy, sortBy, find, nub )
import Data.Maybe ( catMaybes )
import Data.Ord ( comparing )
import Data.Generics ( Data(gmapQ) )
import Data.Generics.Schemes ( everywhere' )
import Data.Generics.Aliases ( extT, extQ )
import Language.Haskell.TH
import Language.Haskell.Meta.Utils (conName, conTypes)
mkNewTypes :: [Name] -> Q [Dec]
mkNewTypes = mapM mkInst
where
mkInst name = rewriteFamilies =<< mkInstH name <$> reify name
mkInstH name (TyConI (NewtypeD context _ vs con _)) =
InstanceD context
(foldl1 AppT [ConT ''Newtype, bndrsToType (ConT name) vs, head $ conTypes con])
(defs (conName con))
mkInstH name _ = error $ show name ++ " is not a Newtype"
defs cname =
[ FunD 'unpack [Clause [ConP cname [VarP xname]] (NormalB $ VarE xname) []]
, FunD 'pack [Clause [] (NormalB (ConE cname)) []]
]
xname = mkName "x"
bndrsToType :: Type -> [TyVarBndr] -> Type
bndrsToType = foldl (\x y -> AppT x $ bndrToType y)
bndrToType :: TyVarBndr -> Type
bndrToType (PlainTV x) = VarT x
bndrToType (KindedTV x k) = SigT (VarT x) k
rewriteFamilies :: Dec -> Q Dec
rewriteFamilies (InstanceD preds ity ds) = do
infos <- mapM (\(n, t) -> (n, t, ) <$> reify n) $ apps ity
fams <- mapM (\(ns, t) -> (ns, t, ) . VarT <$> newName "f")
. map (nub . map snd &&& (snd . fst . head))
. groupBy ((==) `on` fst)
. sortBy (comparing ((id *** show) . fst))
. catMaybes $ map process infos
return $ InstanceD (preds' fams) (ity' fams) ds
where
process (n, t, TyConI (FamilyD _ n' _ _)) = Just ((n', t), n)
process _ = Nothing
preds' fams = map (\((n:_), t, v) -> EqualP v (AppT (ConT n) t)) fams ++ preds
ity' :: [([Name], Type, Type)] -> Type
ity' fams = everywhere' (id `extT` handleType) ity
where
handleType :: Type -> Type
handleType app@(AppT (ConT n) r)
= case find (\(ns, t, _) -> n `elem` ns && t == r) fams of
Just (_, _, v) -> v
Nothing -> app
handleType t = t
apps :: Type -> [(Name, Type)]
apps = handleType
where
handleType :: Type -> [(Name, Type)]
handleType (AppT (ConT v) r) = (v, r) : handleType r
handleType t = generic t
generic :: Data a => a -> [(Name, Type)]
generic = concat . gmapQ (const [] `extQ` handleType)
rewriteFamilies d = return d