module Control.Monad.Operational.TH (makeSingletons) where
import Prelude hiding (mapM)
import Language.Haskell.TH
import Control.Monad.Operational.Class
import Data.Char
import qualified Data.Map as Map
import Data.List (nub)
import Data.Traversable
renameType :: Map.Map Name Type -> Type -> Type
renameType m (VarT n) = case n `Map.lookup` m of
Just t -> t
Nothing -> VarT n
renameType m (SigT t k) = SigT (renameType m t) k
renameType m (AppT l r) = AppT (renameType m l) (renameType m r)
renameType _ t = t
tyVars :: Type -> [Name]
tyVars (VarT n) = [n]
tyVars (AppT l r) = tyVars l ++ tyVars r
tyVars _ = []
makeSingletons :: Name -> Q [Dec]
makeSingletons name = do
TyConI dec <- reify name
case dec of
DataD _ _ vs cs _ -> fmap concat $ mapM (fromCon (map fromTyVarBndr vs)) cs
_ -> fail "Expecting a type construcor"
where
gen vs_ eqs_ conName argTypes_ resultType_ = do
let bodyName = let (b:bs) = nameBase conName in mkName (toLower b : bs)
let refresh1 m i = case Map.lookup i m of
Just (VarT v) -> v
_ -> i
let ref = Map.fromList [(v, VarT $ mkName $ "v" ++ show i) | (i, v) <- zip [(0 :: Int)..] vs_]
let vs = map (refresh1 ref) vs_
let resultType = renameType ref resultType_
let argTypes = map (renameType ref) argTypes_
let eqs = [EqualP (renameType ref s) (renameType ref t) | EqualP s t <- eqs_]
let eqm = Map.fromList [(v, t) | EqualP (VarT v) t <- eqs]
let vs' = map (refresh1 eqm) vs
let resultType' = renameType eqm resultType
let argTypes' = map (renameType eqm) argTypes
let instr = renameType eqm $ foldl AppT (ConT name) $ map VarT (init vs')
let m = mkName "m"
let vars = map PlainTV $ (m :) $ nub
$ tyVars instr ++ tyVars resultType' ++ concatMap tyVars argTypes'
let sig = SigD bodyName $ ForallT vars [ClassP ''(:!) [instr, VarT m]]
$ foldr (\x y -> AppT ArrowT x `AppT` y) (AppT (VarT m) resultType') argTypes'
ps <- mapM (newName . ("p"++) . show) [0..length argTypes 1]
let body = AppE (VarE 'singleton) $ foldl AppE (ConE conName) (map VarE ps)
return [sig, FunD bodyName [Clause (map VarP ps) (NormalB body) []]]
fromCon vs (ForallC _ eqs (NormalC conName ts)) = gen vs eqs conName (map snd ts) (VarT $ last vs)
fromCon vs (NormalC conName ts) = gen vs [] conName (map snd ts) (VarT $ last vs)
fromCon _ _ = fail "Unsupported data constructor"
fromTyVarBndr :: TyVarBndr -> Name
fromTyVarBndr (PlainTV n) = n
fromTyVarBndr (KindedTV n _) = n