How can I use Template Haskell to build structures polymorphically? - haskell

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