CollectExpQ.hs

{-# LANGUAGE TemplateHaskell #-}

import Language.Haskell.TH

{--

-- collect the results of expressions as a tuple,
-- e1 >>= \a ->
-- e2 >>= \b ->
-- return (a,b)

-- for example: 
*Main> let es n = mapM (\i -> [|return i|]) [1..n]
*Main> $( es 5 >>= tupled . map return )
(1,2,3,4,5)
*Main> $( es 5 >>= listed . map return )
[1,2,3,4,5]
--}

collect :: ([ExpQ] -> ExpQ) -> [ExpQ] -> ExpQ
collect cl es = do
    let n = length es 
    as <- sequence (replicate n (newName "a"))
    let ret :: ExpQ
        ret = appE (varE $ mkName "return") (cl $ map varE as)
        bind :: (Name, ExpQ) -> ExpQ -> ExpQ
        bind (a,e) r = appE (appE (varE $ mkName ">>=") e) (lam1E (varP a) r)
    foldr bind ret $ zip as es

tupled :: [ExpQ] -> ExpQ
tupled = collect tupE

listed :: [ExpQ] -> ExpQ
listed = collect listE