{-# LANGUAGE TemplateHaskell #-} module Data.Geniplate(universeBi, universeBiT, transformBi, transformBiT) where import Control.Exception(assert) import Control.Monad.State.Strict import Data.Maybe import Language.Haskell.TH import Language.Haskell.TH.Syntax hiding (lift) -- | Generate TH code for a function that extracts all subparts of a certain type. -- The argument to 'universeBi' is a name with the type @S -> [T]@, for some types -- @S@ and @T@. The function will extract all subparts of type @T@ from @S@. universeBi :: Name -> Q Exp universeBi = universeBiT [] -- | Same as 'universeBi', but does not look inside any types mention in the -- list of types. universeBiT :: [TypeQ] -> Name -> Q Exp universeBiT stops name = do (_tvs, from, tos) <- getNameType name let to = unList tos -- qRunIO $ print (from, to) (ds, f) <- uniBiQ stops from to x <- newName "_x" let e = LamE [VarP x] $ LetE ds $ AppE (AppE f (VarE x)) (ListE []) -- qRunIO $ putStrLn $ pprint e return e type U = StateT (Map Type Dec, Map Type Bool) Q uniBiQ :: [TypeQ] -> Type -> Type -> Q ([Dec], Exp) uniBiQ stops from ato = do ss <- sequence stops to <- expandSyn ato (f, (m, _)) <- runStateT (uniBi from to) (mEmpty, mFromList $ zip ss (repeat False)) return (mElems m, f) uniBi :: Type -> Type -> U Exp uniBi afrom to = do (m, c) <- get from <- lift $ expandSyn afrom case mLookup from m of Just (FunD n _) -> return $ VarE n _ -> do f <- lift $ newName "_f" cs <- if from == to then lift $ fmap unFunD [d| f _x _r = _x : _r |] else do b <- contains to from if b then do put (mInsert from (FunD f [Clause [] (NormalB $ TupE []) []]) m, c) -- insert something to break recursion, will be replaced below. uniBiCase from to else -- No occurrences of to inside from, so add nothing. lift $ fmap unFunD [d| f _ _r = _r |] let d = FunD f cs modify $ \ (m', c') -> (mInsert from d m', c') return $ VarE f -- Check if the second type is contained anywhere in the first type. contains :: Type -> Type -> U Bool contains to afrom = do -- lift $ qRunIO $ print ("contains", to, from) (m, c) <- get from <- lift $ expandSyn afrom case mLookup from c of Just b -> return b Nothing -> do if from == to then return True -- Don't bother caching; we should reach this case where caching matters else do let (con, ts) = splitTypeApp from put (m, mInsert from False c) -- To make the fixpoint of the recursion false. b <- case con of ConT n -> containsCon n to ts TupleT _ -> fmap or $ mapM (contains to) ts ArrowT -> return False ListT -> contains to (head ts) t -> genError $ "contains: unexpected type: " ++ pprint from ++ " (" ++ show t ++ ")" modify $ \ (m', c') -> (m', mInsert from b c') return b containsCon :: Name -> Type -> [Type] -> U Bool containsCon con to ts = do (tvs, cons) <- lift $ getTyConInfo con let conCon (NormalC _ xs) = fmap or $ mapM (field . snd) xs conCon (InfixC x1 _ x2) = fmap or $ mapM field [snd x1, snd x2] conCon (RecC _ xs) = fmap or $ mapM field [ t | (_,_,t) <- xs ] conCon c = genError $ "containsCon: " ++ show c s = mkSubst tvs ts field t = contains to (subst s t) fmap or $ mapM conCon cons unFunD :: [Dec] -> [Clause] unFunD [FunD _ cs] = cs unFunD _ = genError $ "unFunD" uniBiCase :: Type -> Type -> U [Clause] uniBiCase from to = do let (con, ts) = splitTypeApp from case con of ConT n -> uniBiCon n ts to TupleT _ -> uniBiTuple ts to -- ArrowT -> lift $ fmap unFunD [d| f _ _r = _r |] -- Stop at functions ListT -> uniBiList (head ts) to t -> genError $ "uniBiCase: unexpected type: " ++ pprint from ++ " (" ++ show t ++ ")" uniBiList :: Type -> Type -> U [Clause] uniBiList t to = do uni <- uniBi t to rec <- uniBi (AppT ListT t) to lift $ fmap unFunD [d| f [] _r = _r; f (_x:_xs) _r = $(return uni) _x ($(return rec) _xs _r) |] uniBiTuple :: [Type] -> Type -> U [Clause] uniBiTuple ts to = fmap (:[]) $ mkArm to [] TupP ts uniBiCon :: Name -> [Type] -> Type -> U [Clause] uniBiCon con ts to = do (tvs, cons) <- lift $ getTyConInfo con let genArm (NormalC c xs) = arm (ConP c) xs genArm (InfixC x1 c x2) = arm (\ [p1, p2] -> InfixP p1 c p2) [x1, x2] genArm (RecC c xs) = arm (ConP c) [ (b,t) | (_,b,t) <- xs ] genArm c = genError $ "uniBiCon: " ++ show c s = mkSubst tvs ts arm c xs = mkArm to s c $ map snd xs if null cons then -- No constructurs, return nothing lift $ fmap unFunD [d| f _ _r = _r |] else mapM genArm cons mkArm :: Type -> Subst -> ([Pat] -> Pat) -> [Type] -> U Clause mkArm to s c ts = do r <- lift $ newName "_r" vs <- mapM (const $ lift $ newName "_x") ts let sub v t = do let t' = subst s t uni <- uniBi t' to return $ AppE (AppE uni (VarE v)) es <- zipWithM sub vs ts let body = foldr ($) (VarE r) es return $ Clause [c (map VarP vs), VarP r] (NormalB body) [] type Subst = [(Name, Type)] mkSubst :: [TyVarBndr] -> [Type] -> Subst mkSubst vs ts = let vs' = map un vs un (PlainTV v) = v un (KindedTV v _) = v in assert (length vs' == length ts) $ zip vs' ts subst :: Subst -> Type -> Type subst s (ForallT v c t) = ForallT v c $ subst s t subst s t@(VarT n) = fromMaybe t $ lookup n s subst s (AppT t1 t2) = AppT (subst s t1) (subst s t2) subst s (SigT t k) = SigT (subst s t) k subst _ t = t getTyConInfo :: Name -> Q ([TyVarBndr], [Con]) getTyConInfo con = do info <- qReify con case info of TyConI (DataD _ _ tvs cs _) -> return (tvs, cs) PrimTyConI{} -> return ([], []) i -> genError $ "unexpected TyCon: " ++ show i getNameType :: Name -> Q ([TyVarBndr], Type, Type) getNameType name = do info <- qReify name let split (ForallT tvs _ t) = (tvs ++ tvs', from, to) where (tvs', from, to) = split t split (AppT (AppT ArrowT from) to) = ([], from, to) split t = genError $ "Type is not an arrow: " ++ pprint t case info of VarI _ t _ _ -> return $ split t _ -> genError $ "Name is not variable: " ++ pprint name unList :: Type -> Type unList (AppT (ConT n) t) | n == ''[] = t unList (AppT ListT t) = t unList t = genError $ "universeBi: Type is not a list: " ++ pprint t -- ++ " (" ++ show t ++ ")" splitTypeApp :: Type -> (Type, [Type]) splitTypeApp (AppT a r) = (c, rs ++ [r]) where (c, rs) = splitTypeApp a splitTypeApp t = (t, []) expandSyn :: Type -> Q Type expandSyn (ForallT tvs ctx t) = liftM (ForallT tvs ctx) $ expandSyn t expandSyn t@AppT{} = expandSynApp t [] expandSyn t@ConT{} = expandSynApp t [] expandSyn (SigT t k) = liftM (flip SigT k) $ expandSyn t expandSyn t = return t expandSynApp :: Type -> [Type] -> Q Type expandSynApp (AppT t1 t2) ts = do t2' <- expandSyn t2; expandSynApp t1 (t2':ts) expandSynApp t@(ConT n) ts = do info <- qReify n case info of TyConI (TySynD _ tvs rhs) -> let (ts', ts'') = splitAt (length tvs) ts s = mkSubst tvs ts' rhs' = subst s rhs in expandSynApp rhs' ts'' _ -> return $ foldl AppT t ts expandSynApp t ts = do t' <- expandSyn t; return $ foldl AppT t' ts genError :: String -> a genError msg = error $ "Data.Geniplate: " ++ msg ---------------------------------------------------- -- Exp has type (S -> S) -> T -> T, for some S and T -- | Generate TH code for a function that transforms all subparts of a certain type. -- The argument to 'transformBi' is a name with the type @(S->S) -> T -> T@, for some types -- @S@ and @T@. The function will transform all subparts of type @S@ inside @T@ using the given function. transformBi :: Name -> Q Exp transformBi = transformBiT [] -- | Same as 'transformBi', but does not look inside any types mention in the -- list of types. transformBiT :: [TypeQ] -> Name -> Q Exp transformBiT stops name = do (_tvs, fcn, res) <- getNameType name f <- newName "_f" (ds, tr) <- case (fcn, res) of (AppT (AppT ArrowT s) s', AppT (AppT ArrowT t) t') | s == s' && t == t' -> trBiQ stops f s t _ -> genError $ "transformBi: malformed type: " ++ pprint (AppT (AppT ArrowT fcn) res) ++ ", should have form (S->S) -> (T->T)" x <- newName "_x" let e = LamE [VarP f, VarP x] $ LetE ds $ AppE tr (VarE x) -- qRunIO $ putStrLn $ pprint e return e trBiQ :: [TypeQ] -> Name -> Type -> Type -> Q ([Dec], Exp) trBiQ stops f aft st = do ss <- sequence stops ft <- expandSyn aft (tr, (m, _)) <- runStateT (trBi (VarE f) ft st) (mEmpty, mFromList $ zip ss (repeat False)) return (mElems m, tr) trBi :: Exp -> Type -> Type -> U Exp trBi f ft ast = do (m, c) <- get st <- lift $ expandSyn ast -- lift $ qRunIO $ print (ft, st) case mLookup st m of Just (FunD n _) -> return $ VarE n _ -> do tr <- lift $ newName "_tr" cs <- if ft == st then lift $ fmap unFunD [d| _f _x = $(return f) _x |] else do b <- contains ft st -- lift $ qRunIO $ print (b, ft, st) if b then do put (mInsert st (FunD tr [Clause [] (NormalB $ TupE []) []]) m, c) -- insert something to break recursion, will be replaced below. trBiCase f ft st else lift $ fmap unFunD [d| f _x = _x |] let d = FunD tr cs modify $ \ (m', c') -> (mInsert st d m', c') return $ VarE tr trBiCase :: Exp -> Type -> Type -> U [Clause] trBiCase f ft st = do let (con, ts) = splitTypeApp st case con of ConT n -> trBiCon f n ft ts TupleT _ -> trBiTuple f ft ts -- ArrowT -> lift $ fmap unFunD [d| f _ _r = _r |] -- Stop at functions ListT -> trBiList f ft (head ts) _ -> genError $ "trBiCase: unexpected type: " ++ pprint st ++ " (" ++ show st ++ ")" trBiList :: Exp -> Type -> Type -> U [Clause] trBiList f ft st = do tr <- trBi f ft st rec <- trBi f ft (AppT ListT st) lift $ fmap unFunD [d| _f [] = []; _f (_x:_xs) = ($(return tr) _x) : ($(return rec) _xs) |] trBiTuple :: Exp -> Type -> [Type] -> U [Clause] trBiTuple f ft ts = fmap (:[]) $ trMkArm f ft [] TupP TupE ts trBiCon :: Exp -> Name -> Type -> [Type] -> U [Clause] trBiCon f con ft ts = do (tvs, cons) <- lift $ getTyConInfo con let genArm (NormalC c xs) = arm (ConP c) (foldl AppE $ ConE c) xs genArm (InfixC x1 c x2) = arm (\ [p1, p2] -> InfixP p1 c p2) (\ [e1, e2] -> InfixE (Just e1) (ConE c) (Just e2)) [x1, x2] genArm (RecC c xs) = arm (ConP c) (foldl AppE $ ConE c) [ (b,t) | (_,b,t) <- xs ] genArm c = genError $ "trBiCon: " ++ show c s = mkSubst tvs ts arm c ec xs = trMkArm f ft s c ec $ map snd xs mapM genArm cons trMkArm :: Exp -> Type -> Subst -> ([Pat] -> Pat) -> ([Exp] -> Exp) -> [Type] -> U Clause trMkArm f ft s c ec ts = do vs <- mapM (const $ lift $ newName "_x") ts let sub v t = do let t' = subst s t tr <- trBi f ft t' return $ AppE tr (VarE v) es <- zipWithM sub vs ts let body = ec es return $ Clause [c (map VarP vs)] (NormalB body) [] ---------------------------------------------------- -- Can't use Data.Map since TH stuff is not in Ord newtype Map a b = Map [(a, b)] mEmpty :: Map a b mEmpty = Map [] mLookup :: (Eq a) => a -> Map a b -> Maybe b mLookup a (Map xys) = lookup a xys mInsert :: (Eq a) => a -> b -> Map a b -> Map a b mInsert a b (Map xys) = Map $ (a, b) : filter ((/= a) . fst) xys mElems :: Map a b -> [b] mElems (Map xys) = map snd xys mFromList :: [(a, b)] -> Map a b mFromList xys = Map xys