module Control.Newtype.TH
( mkNewtype, mkNewtypes
, mkNewType, 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)
mkNewtype :: Name -> Q [Dec]
mkNewtype = mkNewTypes . (:[])
mkNewtypes :: [Name] -> Q [Dec]
mkNewtypes = mapM (\n -> rewriteFamilies =<< mkInst <$> reify n)
where
mkInst (TyConI (NewtypeD a b c d _)) = mkInstFor a b c d
mkInst (TyConI (DataD a b c [d] _)) = mkInstFor a b c d
mkInst x
= error $ show x
++ " is not a Newtype or single-field single-constructor datatype."
mkInstFor context name bnds con
= InstanceD context
( foldl1 AppT [ ConT ''Newtype
, bndrsToType (ConT name) bnds
, head $ conTypes con
] )
[ FunD 'pack [Clause [] (NormalB $ ConE cn) []]
, FunD 'unpack [Clause [ConP cn [VarP xn]] (NormalB $ VarE xn) []]
]
where
cn = conName con
xn = 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")
. mergeApps . catMaybes $ map justFamily infos
return $ InstanceD (preds' fams) (ity' fams) ds
where
justFamily :: (Name, Type, Info) -> Maybe (Name, (Name, Type))
#if __GLASGOW_HASKELL__ >= 704
justFamily (n, t, FamilyI (FamilyD _ n' _ _) _) = Just (n, (n', t))
#else
justFamily (n, t, TyConI (FamilyD _ n' _ _)) = Just (n, (n', t))
#endif
justFamily _ = Nothing
mergeApps :: [(Name, (Name, Type))] -> [([Name], Type)]
mergeApps = map (nub . map fst &&& (snd . snd . head))
. groupBy ((==) `on` snd) . sortBy (comparing snd)
preds' = (preds ++)
. map (\((n:_), t, v) -> EqualP v (AppT (ConT n) t))
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 (AppT (SigT t _) r) = handleType (AppT t r)
handleType t = generic t
generic :: Data a => a -> [(Name, Type)]
generic = concat . gmapQ (const [] `extQ` handleType)
rewriteFamilies d = return d
mkNewType :: Name -> Q [Dec]
mkNewType = mkNewtype
mkNewTypes :: [Name] -> Q [Dec]
mkNewTypes = mkNewtypes