How can I use Template Haskell to build structures polymorphically?

Question

I can write an instance

-- In Data.Sequence.Internal
instance Lift a => Lift (Seq a) where
  ...

letting users lift fully realized sequences into splices. But suppose I want something a bit different, to build functions for creating sequences?

sequenceCode :: Quote m => Seq (Code m a) -> Code m (Seq a)
sequenceCode = ???

The idea would be that I'd be able to write something like

triple :: a -> a -> a -> Seq a
triple a b c = $$(sequenceCode (fromList [[|| a ||], [|| b ||], [|| c ||]]))

and have that function build its sequence directly with the underlying sequence constructors, rather than having to build and convert a list at run-time.

It's not very hard to write something like sequenceCode directly for sequences, using their internals (look below the jump). But, as the name suggests, sequenceCode looks a lot like sequence. Is there a way to generalize it? A moment's reflection shows that Traversable is insufficient. Would it be possible to do something with the Generic1 class in staged generics? I made a few attempts, but I don't understand that package well enough to know the right place to start. Would it be possible even just using plain old GHC generics? I'm beginning to suspect so, but I haven't tried yet and it will surely be hairy.

---

Here's the code for a Data.Sequence version:

{-# language TemplateHaskellQuotes #-}
import Data.Sequence.Internal
import qualified Language.Haskell.TH.Syntax as TH

class Functor t => SequenceCode t where
  traverseCode :: TH.Quote m => (a -> TH.Code m b) -> t a -> TH.Code m (t b)
  traverseCode f = sequenceCode . fmap f
  sequenceCode :: TH.Quote m => t (TH.Code m a) -> TH.Code m (t a)
  sequenceCode = traverseCode id

instance SequenceCode Seq where
  sequenceCode (Seq t) = [|| Seq $$(traverseCode sequenceCode t) ||]

instance SequenceCode Elem where
  sequenceCode (Elem t) = [|| Elem $$t ||]

instance SequenceCode FingerTree where
  sequenceCode (Deep s pr m sf) =
    [|| Deep s $$(sequenceCode pr) $$(traverseCode sequenceCode m) $$(sequenceCode sf) ||]
  sequenceCode (Single a) = [|| Single $$a ||]
  sequenceCode EmptyT = [|| EmptyT ||]

instance SequenceCode Digit where
  sequenceCode (One a) = [|| One $$a ||]
  sequenceCode (Two a b) = [|| Two $$a $$b ||]
  sequenceCode (Three a b c) = [|| Three $$a $$b $$c ||]
  sequenceCode (Four a b c d) = [|| Four $$a $$b $$c $$d ||]

instance SequenceCode Node where
  sequenceCode (Node2 s x y) = [|| Node2 s $$x $$y ||]
  sequenceCode (Node3 s x y z) = [|| Node3 s $$x $$y $$z ||]

Then in another module, we can define triple as above:

triple :: a -> a -> a -> Seq a
triple a b c = $$(sequenceCode (fromList [[|| a ||], [|| b ||], [|| c ||]]))

When I compile this with -ddump-splices (or -ddump-ds), I can verify that the sequence is built directly rather than using fromList.

Answer 1

I've uploaded a package that does this.

It turns out that GHC.Generics is sufficient. However, I will actually use linear-generics instead, because it has a more general version of Generic1. The idea is that by examining the generic representation of a value, we can build up all the information we need to produce a Template Haskell code for it. It's all quite low-level! First, some throat-clearing:

{-# language TemplateHaskellQuotes #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language ScopedTypeVariables #-}
{-# language DataKinds #-}
{-# language TypeOperators #-}
{-# language EmptyCase #-}
{-# language DefaultSignatures #-}

module Language.Haskell.TH.TraverseCode
  ( TraverseCode (..)
  , sequenceCode
  , genericTraverseCode
  , genericSequenceCode
  ) where
import Generics.Linear
import Language.Haskell.TH.Syntax
  (Code, Lift (..), Exp (..), Quote, Name)
import qualified Language.Haskell.TH.Syntax as TH
import Language.Haskell.TH.Lib (conE)
import Data.Kind (Type)

-- for instances
import qualified Data.Functor.Product as FProd
import qualified Data.Functor.Sum as FSum
import Data.Functor.Identity
import qualified Data.Sequence.Internal as Seq
import Data.Coerce

Now we'll get into the meat of things:

class TraverseCode t where
  traverseCode :: Quote m => (a -> Code m b) -> t a -> Code m (t b)

  default traverseCode
    :: (Quote m, GTraverseCode (Rep1 t), Generic1 t)
    => (a -> Code m b) -> t a -> Code m (t b)
  traverseCode = genericTraverseCode

sequenceCode
  :: (TraverseCode t, Quote m)
  => t (Code m a) -> Code m (t a)
sequenceCode = traverseCode id

genericSequenceCode
  :: (Quote m, GTraverseCode (Rep1 t), Generic1 t)
  => t (Code m a) -> Code m (t a)
genericSequenceCode = TH.unsafeCodeCoerce . gtraverseCode id . from1

genericTraverseCode
  :: (Quote m, GTraverseCode (Rep1 t), Generic1 t)
  => (a -> Code m b) -> t a -> Code m (t b)
genericTraverseCode f = TH.unsafeCodeCoerce . gtraverseCode f . from1

class GTraverseCode f where
  gtraverseCode :: Quote m => (a -> Code m b) -> f a -> m Exp

Why do we use untyped Template Haskell here? Simple: it's pretty easy to build the expressions we need, but working out how to make types useful for the sub-expressions would be tricky. So then, of course, we need generic instances. We'll work our way down step by step, from the outside in, gathering info along the way.

First we look at the type stuff:

instance (Datatype c, GTraverseCodeCon f)
  => GTraverseCode (D1 c f) where
  gtraverseCode f m@(M1 x) = gtraverseCodeCon pkg modl f x
    where
      pkg = packageName m
      modl = moduleName m

This gets us the names GHC uses for the package and module.

Next we look at the constructor stuff:

class GTraverseCodeCon f where
  gtraverseCodeCon :: Quote m => String -> String -> (a -> Code m b) -> f a -> m Exp

instance GTraverseCodeCon V1 where
  gtraverseCodeCon _pkg _modl _f x = case x of

instance (GTraverseCodeCon f, GTraverseCodeCon g)
  => GTraverseCodeCon (f :+: g) where
  gtraverseCodeCon pkg modl f (L1 x) = gtraverseCodeCon pkg modl f x
  gtraverseCodeCon pkg modl f (R1 y) = gtraverseCodeCon pkg modl f y

instance (Constructor c, GTraverseCodeFields f)
  => GTraverseCodeCon (C1 c f) where
  gtraverseCodeCon pkg modl f m@(M1 x) = gtraverseCodeFields (conE conN) f x
    where
      conBase = conName m
      conN :: Name
      conN = TH.mkNameG_d pkg modl conBase

The interesting case is when we reach an actual constructor (C1). Here we grab the (unqualified) name of the constructor from the Constructor instance, and combine it with the package and module names to get the Template Haskell Name of the constructor, from which we can build an expression referring to it. This expression gets passed on down to the lowest level, where we deal with fields. The rest is basically a left fold over those fields.

class GTraverseCodeFields f where
  gtraverseCodeFields :: Quote m => m Exp -> (a -> Code m b) -> f a -> m Exp

instance GTraverseCodeFields f => GTraverseCodeFields (S1 c f) where
  gtraverseCodeFields c f (M1 x) = gtraverseCodeFields c f x

instance (GTraverseCodeFields f, GTraverseCodeFields g)
  => GTraverseCodeFields (f :*: g) where
  gtraverseCodeFields c f (x :*: y) =
    gtraverseCodeFields (gtraverseCodeFields c f x) f y

instance Lift p => GTraverseCodeFields (K1 i p) where
  gtraverseCodeFields c _f (K1 x) = [| $c x |]

instance GTraverseCodeFields Par1 where
  gtraverseCodeFields cc f (Par1 ca) =
    [| $cc $(TH.unTypeCode (f ca)) |]

instance GTraverseCodeFields U1 where
  gtraverseCodeFields cc _f U1 = cc


-- Note: this instance is *different* from the one that we'd
-- write if we were using GHC.Generics, because composition works
-- differently in Generics.Linear.
instance (GTraverseCodeFields f, TraverseCode g) => GTraverseCodeFields (f :.: g) where
  gtraverseCodeFields cc f (Comp1 x) =
    gtraverseCodeFields cc (traverseCode f) x

Now we can write all sorts of instances:

instance TraverseCode Maybe
instance TraverseCode Identity
instance TraverseCode []
instance TH.Lift a => TraverseCode (Either a)
instance TH.Lift a => TraverseCode ((,) a)
instance (TraverseCode f, TraverseCode g)
  => TraverseCode (FProd.Product f g)
instance (TraverseCode f, TraverseCode g)
  => TraverseCode (FSum.Sum f g)
instance TraverseCode V1

-- The Elem instance isn't needed for the Seq instance
instance TraverseCode Seq.Elem
instance TraverseCode Seq.Digit
instance TraverseCode Seq.Node
instance TraverseCode Seq.FingerTree

For the Seq instance I was after, we need to write something by hand, because Seq isn't an instance of Generic1 (and we don't want it to be). Additionally, we don't really want the derived instance. Using a bit of coercion magic, and knowing a little something about how zipWith and replicate work on sequences, we can minimize the size of the splice and the number of types GHC has to deal with once it's compiled to Core.

instance TraverseCode Seq.Seq where
  -- Stick a single coercion on the outside, instead of having a bunch
  -- of `Elem` constructors on the inside.
  traverseCode f s = [|| coerceFT $$(traverseCode f ft') ||]
    where
      -- Use zipWith to make the tree representing the sequence
      -- nice and shallow.
      ft' = coerceSeq (Seq.zipWith (flip const) (Seq.replicate (Seq.length s) ()) s)
coerceFT :: Seq.FingerTree a -> Seq.Seq a
coerceFT = coerce
coerceSeq :: Seq.Seq a -> Seq.FingerTree a
coerceSeq = coerce