Composing Declarations in Template Haskell

I have recently tried to use Template Haskell to generate both a datatype and lenses for accessing the fields of this datatype, and it was harder than it should have been. In this post, I will demonstrate the problem, I will pinpoint its cause, and I will propose a solution.

The Problem

Consider the following code. I'm using a simple, contrived example instead of a more realistic one because it will be easier to write Template Haskell code for generating this silly code than it would be to write Template Haskell code generating lenses and such.

class Default a where
  def :: a


data X = X
data Y = Y X

data DoubledX = DoubledX
data DoubledY = DoubledY X X

instance Default X        where def = X
instance Default Y        where def = Y def
instance Default DoubledX where def = DoubledX
instance Default DoubledY where def = DoubledY def def

Most of that code is boilerplate, and I would like to generate that boilerplate using Template Haskell. I hasten to note that Template Haskell is a tool of last resort, to be used only when none of Haskell's many other abstraction facilities would have sufficed. In this example, I would probably use some generics library to define a default implementation of def for any algebraic datatype:

{-# LANGUAGE DefaultSignatures, DeriveGeneric, FlexibleContexts #-}
import Generics.Eot

class Default a where
  def :: a
  
  default def :: (HasEot a, Default (Eot a)) => a
  def = fromEot def

instance Default () where
  def = ()

instance (Default a, Default b) => Default (a, b) where
  def = (def, def)

instance Default a => Default (Either a b) where
  def = Left def


instance Default X
instance Default Y
instance Default DoubledX
instance Default DoubledY

This works fine, but today I want to talk about one of Template Haskell's limitations, so let's write a Template Haskell implementation instead.

{-# LANGUAGE TemplateHaskell #-}
import Data.List
import Language.Haskell.TH

-- > data Foo = Bar Int String
-- > generateDefault ''Foo
--
-- generates
--
-- > instance Default Foo where def = Bar def def
generateDefault :: Name -> Q [Dec]
generateDefault name = do
  --      data Foo =             Bar   Int...
  TyConI (DataD _ _ _ _ (NormalC cname cargs:_) _) <- reify name
  
  -- Bar def...
  let expr = foldl' (\c _ -> [|$c def|]) (conE cname) cargs
  
  [d|instance Default $(conT name) where def = $expr|]


data X = X
data Y = Y X

data DoubledX = DoubledX
data DoubledY = DoubledY X X

generateDefault ''X
generateDefault ''Y
generateDefault ''DoubledX
generateDefault ''DoubledY

In addition to the Default instances, we can also generate the Doubled datatypes, they are a version of the original datatype which has two copies of each field:

-- > data Foo = Bar Int String
-- > generateDoubled ''Foo
--
-- generates
--
-- > data DoubledFoo = DoubledBar Int String Int String
generateDoubled :: Name -> Q [Dec]
generateDoubled name = do
  --    data Foo =             Bar   Int...
  TyConI (DataD _ _ _ _ (NormalC cname cargs:_) _) <- reify name
  
  let cons = [NormalC (doubledName cname) (cargs ++ cargs)]
  pure [DataD [] (doubledName name) [] Nothing cons []]

doubledName :: Name -> Name
doubledName = mkName . ("Doubled" ++) . nameBase


data X = X
data Y = Y X

generateDoubled ''X
generateDoubled ''Y

generateDefault ''X
generateDefault ''Y
generateDefault ''DoubledX
generateDefault ''DoubledY

So, we can write a Template Haskell function which generates a datatype, and we can write one which generates an instance for that datatype. But can we write one which generates both the datatype and its instance? Both of our functions are Q actions which produce a [Dec], so it looks like a no brainer: we can simply run both Q actions one after the other and concatenate the resulting lists.

generateBoth :: Name -> Q [Dec]
generateBoth name = (++) <$> generateDoubled name
                         <*> generateDefault (doubledName name)


data X = X

-- error: ‘DoubledX’ is not in scope at a reify
generateBoth ''X

Sadness, it does not work :(

The Cause

The reason DoubledX is not in scope when generateDefault calls reify ''DoubledX is that the DoubledX datatype is not created as a side-effect of the generateDoubled ''X action, but as a side-effect of splicing the resulting [Dec] into the surrounding code. When concatenating the two lists, this doesn't happen until after both lists have been generated, and so DoubledX cannot be "reified" while generating the second list.

I didn't thought I'd ever say something like this, but: this pure API was a mistake, an API based on side-effects would be better! I'll qualify that statement though: since reify obtains information about datatypes (and other named entities) via a side-effect, namely reading from some global symbol table, I think there should be a corresponding action for adding new names to this table. As we have seen, with the current API in which names are added by returning a pure [Dec] value, declaration templates don't compose, so I think that API was a mistake.

I should note that there is, in fact, an action for adding new declarations as a side-effect:

addTopDecls :: [Dec] -> Q ()

Unfortunately, as of this writing, addTopDecls is unable to add datatypes:

import Language.Haskell.TH.Syntax

generateBoth :: Name -> Q [Dec]
generateBoth name = do
  decs <- generateDoubled name
  addTopDecls decs
  generateDefault (doubledName name)


data X = X

-- error: Only function, value, annotation, and foreign import 
--        declarations may be added with addTopDecl
-- error: ‘DoubledX’ is not in scope at a reify
generateBoth ''X

So the real solution would be to fix the implementation so it also supports datatypes, but until then, I have a simpler solution.

Simple Solution

We can't change declaration-splicing's Q [Dec] API, but if we create our own API, we can write an adapter which turns it into a Q [Dec].

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

newtype LocalQ a = LocalQ { unLocalQ :: StateT [Dec] Q a }
  deriving (Functor, Applicative, Monad, MonadFail)

runLocalQ :: LocalQ () -> Q [Dec]
runLocalQ = flip execStateT [] . unLocalQ

The idea is that in addition to Q's global symbol table, LocalQ also has access to a local symbol table holding the declarations which have been added within the LocalQ computation.

addLocalDecls :: [Dec] -> LocalQ ()
addLocalDecls decls = LocalQ $ modify (++ decls)

reifyLocallyFirst :: Name -> LocalQ Info
reifyLocallyFirst name = do
  decls <- LocalQ get
  case find (hasName name) decls of
    Just dec -> pure $ TyConI dec
    Nothing  -> liftQ $ reify name

-- for this simple example, I'm only interested
-- in datatype declarations
hasName :: Name -> Dec -> Bool
hasName expected (DataD _ actual _ _ _ _) = actual == expected
hasName _        _                        = False

liftQ :: Q a -> LocalQ a
liftQ = LocalQ . lift

If we reimplement our declaration templates to use this new, better API...

locallyGenerateDefault :: Name -> LocalQ ()
locallyGenerateDefault name = do
  info <- reifyLocallyFirst name
  let TyConI (DataD _ _ _ _ (NormalC cname cargs:_) _) = info
  let expr = foldl' (\c _ -> [|$c def|]) (conE cname) cargs
  decls <- liftQ [d|instance Default $(conT name) where def = $expr|]
  addLocalDecls decls

locallyGenerateDoubled :: Name -> LocalQ ()
locallyGenerateDoubled name = do
  info <- reifyLocallyFirst name
  let TyConI (DataD _ _ _ _ (NormalC cname cargs:_) _) = info
  let cons = [NormalC (doubledName cname) (cargs ++ cargs)]
  addLocalDecls [DataD [] (doubledName name) [] Nothing cons []]

...then this time we can compose them just fine:

locallyGenerateBoth :: Name -> LocalQ ()
locallyGenerateBoth name = do
  locallyGenerateDoubled name
  locallyGenerateDefault (doubledName name)


data X = X
data Y = Y X

runLocalQ $ locallyGenerateDefault ''X
runLocalQ $ locallyGenerateDefault ''Y
runLocalQ $ locallyGenerateBoth ''X
runLocalQ $ locallyGenerateBoth ''Y

Happiness, it works! Now all that's left is to convince everybody to rewrite their declaration templates using LocalQ instead of Q, and we'll finally be able to reuse each other's code.

Final Solution

Okay, so that last part is probably not going to happen. If only there was a way to monkey-patch existing Q code so it would use reifyLocallyFirst instead of reify...

Well, here's a little-known fact about Q:

newtype Q a = Q { unQ :: forall m. Quasi m => m a }

That's right, Q isn't some magic Monad which only the compiler can execute! It's a concrete type, which we can examine and manipulate however we want. The finally-tagless encoding might be a bit intimidating, but in practice, a Q a value is basically an AST listing which Q actions need to be performed in order to produce an a. So we should be able to dive in and replace all the reify calls with reifyLocallyFirst calls, no problem.

The finally-tagless way to do that is to write a Quasi instance which instantiates reify with reifyLocallyFirst, and delegates all the other operations to some other Quasi instance:

instance Quasi LocalQ where
  qReify              = reifyLocallyFirst
  qAddTopDecls        = addLocalDecls
  
  -- Most of those aren't being exercised by my simple example,
  -- so I can afford to use 'undefined' for the trickier methods.
  qGetQ                   = undefined
  qPutQ                   = undefined
  qRecover                = undefined
  qAddDependentFile   x   = liftQ $ qAddDependentFile   x
  qAddModFinalizer    x   = liftQ $ qAddModFinalizer    x
  qExtsEnabled            = liftQ $ qExtsEnabled
  qIsExtEnabled       x   = liftQ $ qIsExtEnabled       x
  qLocation               = liftQ $ qLocation
  qLookupName         x y = liftQ $ qLookupName         x y
  qNewName            x   = liftQ $ qNewName            x
  qReifyAnnotations   x   = liftQ $ qReifyAnnotations   x
  qReifyConStrictness x   = liftQ $ qReifyConStrictness x
  qReifyFixity        x   = liftQ $ qReifyFixity        x
  qReifyInstances     x y = liftQ $ qReifyInstances     x y
  qReifyModule        x   = liftQ $ qReifyModule        x
  qReifyRoles         x   = liftQ $ qReifyRoles         x
  qReport             x y = liftQ $ qReport             x y
  qRunIO              x   = liftQ $ qRunIO              x

With this instance, I can now transform Q [Dec] templates into LocalQ () templates, and the transformed version will use reifyLocallyFirst instead of reify.

localize :: Q [Dec] -> LocalQ ()
localize declarationTemplate = do
  decls <- unQ declarationTemplate
  addLocalDecls decls

generateBoth :: Name -> Q [Dec]
generateBoth name = runLocalQ $ do
  localize $ generateDoubled name
  localize $ generateDefault (doubledName name)


data X = X
data Y = Y X

generateDefault ''X
generateDefault ''Y
generateBoth ''X
generateBoth ''Y

Notice that I am reusing the original generateDefault and generateDoubled implementations, those which caused the reify error when I first tried to implement generateBoth. I am not using the locallyGenerateDefault and locallyGenerateDoubled reimplementations from the previous section. This means that (with a more fleshed out implementation of LocalQ), I should be able to reuse any existing declaration template out there, including Control.Lens.TH.makeLenses! :D