Code Generation with Template Haskell

Question

Background:

Recently, I've been working through Coroutining Folds with Hyperfunctions. Don't bail out now if the whole thing makes little sense - it's just the motivation for the existence of my code, not what I actually want reviewed.

I followed a rather typical approach of starting with the data structures and types in the paper, then alternating between re-implementing them from scratch and comparing my implementations with those in the paper.

The basic hyperfunction operations end up as such:

-- Used to make documenting types in some nested expressions work
{-# language ScopedTypeVariables #-}

-- This pragma and these imports are used in later sections
{-# language TemplateHaskell #-}
import Control.Monad
import qualified Language.Haskell.TH as TH

newtype H a b = H { apply :: H b a -> b }

(!) :: H a b -> H b a -> b
(!) = apply
infixr 1 !

self :: H a a
self = lift id

lift :: (a -> b) -> H a b
lift f = let x = f << x in x

(#) :: H b c -> H a b -> H a c
f # g = H $ \k -> f ! g # k
infixr 8 #

(<<) :: (a -> b) -> H a b -> H a b
f << g = H $ \k -> f (k ! g)

base :: b -> H a b
base = H . const

run :: H a a -> a
run f = f ! self

Everything's good. This is all basically a direct implementation of the primitive operations from the paper. (The pragmas and imports are there for upcoming code.)

Next up is the first part of the paper's core - a fold operation that folds a list into a hyperfunction.

hfold :: (a -> b -> c) -> c -> [a] -> H b c
hfold f z = go
  where
    go []     = base z
    go (x:xs) = f x << go xs

This is also directly from the paper. The first major observation the paper makes is that folds over multiple lists can be combined into a single coroutined fold by making use of the hyperfunctions that result from hfold.

I spent a lot of time working on re-creating this observation. Eventually I managed to put together a foldr variant that traversed two lists simultaneously.

foldr2 :: forall a b r. (a -> b -> r -> r) -> r -> [a] -> [b] -> r
foldr2 f r xs ys = run (hfy # hfx) r
  where
    hfx :: H (r -> r) (b -> r -> r)
    hfx = hfold (\x g y -> f x y . g) (const id) xs

    hfy :: H (b -> r -> r) (r -> r)
    hfy = hfold (flip id) id ys

Note that this is where ScopedTypeVariables comes into play, to let me keep track of the types of the hyperfunctions involved.

I'm not going to pretend this was easy to recreate, even having read the paper through a couple times. It's not the exact formulation used by the paper anyway, as it focuses on zip functions, and I chose to go with the slightly more primitive foldr functions.

I assumed that because of the flexibility of hyperfunctions, I could easily extend the same idea to three input lists.

foldr3 :: forall a b c r. (a -> b -> c -> r -> r) -> r -> [a] -> [b] -> [c] -> r
foldr3 f r xs ys zs = run (hfz # hfy # hfx) r
  where
    hfx :: H (r -> r) (b -> c -> r -> r)
    hfx = hfold (\x g y z -> f x y z . g) (const $ const id) xs

    hfy :: H (b -> c -> r -> r) (c -> r -> r)
    hfy = hfold (flip id) (const id) ys

    hfz = hfold (flip id) id zs

And.. It wasn't very hard at all, in comparison. In fact, it mostly followed mechanically from the 2-list version.

And this is where I got curious about generalizing this pattern. I figured I could write some template haskell to give me the foldr for any number of input lists. So I did, but I don't feel like it's very good.

Code to actually review starts here:

foldrN :: Int -> TH.ExpQ
foldrN n = do
    names <- replicateM n $ TH.newName "xs"
    r <- TH.newName "r"
    f <- TH.newName "f"
    folds <- forM (zip [1..] names) $ \(i, name) ->
        if i == n
        then [| hfold $(foldedFunc f n) $(lastArg n) $(TH.varE name) |]
        else [| hfold (flip id) $(lastArg i) $(TH.varE name) |]

    let chain = foldr1 (\x y -> TH.InfixE (Just x) ((TH.VarE '(#))) (Just y)) folds
        pats = map TH.VarP $ f : r : reverse names

    body <- [| run $(return chain) $(TH.varE r) |]
    return $ TH.LamE pats body

lastArg :: Int -> TH.ExpQ
lastArg 1 = [| id |]
lastArg n = [| const $(lastArg (n - 1)) |]

foldedFunc :: TH.Name -> Int -> TH.ExpQ
foldedFunc fName n = do
    names <- replicateM n $ TH.newName "x"
    fPat <- TH.newName "g"
    let pats = map TH.VarP $ head names : fPat : tail names
        funApp = foldl1 TH.AppE (map TH.VarE $ fName : names)
    body <- [| $(return funApp) . $(TH.varE fPat) |]
    return $ TH.LamE pats (body)

This all works correctly:

>>> :t $(foldrN 1)
$(foldrN 1) :: (a -> b -> b) -> b -> [a] -> b
>>> :t $(foldrN 2)
$(foldrN 2) :: (a -> b -> b1 -> b1) -> b1 -> [a] -> [b] -> b1
>>> :t $(foldrN 3)
$(foldrN 3) :: (a -> b1 -> b -> b2 -> b2) -> b2 -> [a] -> [b1] -> [b] -> b2

However, the only thing I'm happy with the definition of is lastArg. The definition of chain is especially painful, for me.

Is there any way to make this kind of template haskell code better?