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 (\n -> rewriteFamilies =<< mkInstH n <$> reify n)
where
mkInstH name (TyConI (NewtypeD context _ vs con _))
= InstanceD context
( foldl1 AppT [ ConT ''Newtype
, bndrsToType (ConT name) vs
, head $ conTypes con
] )
[ FunD 'pack [ Clause []
(NormalB $ ConE cname) [] ]
, FunD 'unpack [ Clause [ConP cname [VarP xname]]
(NormalB $ VarE xname) [] ]
]
where
cname = conName con
xname = mkName "x"
mkInstH name _ = error $ show name ++ " is not a Newtype"
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))
justFamily (n, t, TyConI (FamilyD _ n' _ _)) = Just (n, (n', t))
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