Template Haskell and Language Extension Soup - Autogenerating the solver map for each day

So what if you are like me and you want to build an executable that can pick and choose which problem’s code to run? What if you are like me again and are tired of the workflow consisting of:

  1. Exporting a solver function from the problem’s module
  2. Importing the solver function in the main module
  3. Adding the solver function to a map with the correct key

My Project Euler solution’s main (not really, but actually) module consists of a preamble resembling

module Solver (solve) where

import Control.Exception (try)

import qualified Data.HashMap.Lazy as H (HashMap, fromList, lookup)

import PE001 (solve001)
import PE002 (solve002)
-- ...
import PE099 (solve099)

Every time I import a new problem, I must change two numbers. Both the module number as well as the solver number. I must, again, duplicate these numbers twice when adding an entry to the solver map.

solvers :: H.HashMap Integer (String -> Integer)
solvers = H.fromList
  [ (001, solve001), (002, solve002), (003, solve003), (004, solve004), (005, solve005)
  , (006, solve006), (007, solve007), (008, solve008), (009, solve009), (010, solve010)
  -- ...
  ]

This is too much effort spent on typing numbers. Haskell doesn’t support wildcard imports so we can give up on automating the module import. That will still have to be done. What if, however, solving a particular problem is represented as having a typeclass instance on a type that represents the problem number? e.g. what if we have

class Solved prb where
  solve :: String -> String

It turns out we need what seems like a spooky language extension at first.

<interactive>:1:24: error:
    • Could not deduce (Solved prb0)
      from the context: Solved prb
        bound by the type signature for:
                   solve :: forall prb. Solved prb => String -> String
        at <interactive>:1:24-48
      The type variable ‘prb0’ is ambiguous
    • In the ambiguity check for ‘solve’
      To defer the ambiguity check to use sites, enable AllowAmbiguousTypes
      When checking the class method:
        solve :: forall prb. Solved prb => String -> String
      In the class declaration for ‘Solved’

However, AllowAmbiguousTypes is not that scary at all. What ends up happening in the background is as follows.

  1. GHC generates the signature solve :: Solved prb => String -> String
  2. Upon compilation this becomes solve :: DictForSolved -> String -> String
    • The type checker will find a suitable dispatch table for prb and stick that table into the first argument
  3. GHC infers that at each call site for solve, there must be an unconstrained type variable prb
    • Namely, when you call solve, which instance should Solved prb resolve to?

Turning on this language extension allows GHC to defer this error to the actual call site. It is then up to the programmer to manually specialize solve and ensure that instance resolution will succeed, otherwise a type checker error will be produced. To do this, we also turn on TypeApplications. This gives us the f @t to bind the type t to the first type variable in the forall list of f. To see this in action:

Prelude> :set -XTypeApplications
Prelude> :set -fprint-explicit-foralls
Prelude> :t show
show :: forall {a}. Show a => a -> String
Prelude> :t show @Int
show @Int :: Int -> String

For whatever reason GHC will complain when you perform :t on ambiguous typeclass functions but it works in much the same way. If we declare instance Solved Int where ... and then call solve @Int, it will use the solve from that instance. We’re still left with a conundrum though. We now need a new type for each instance we want to declare! This is equivalent to the work that we did before so clearly we haven’t turned on enough language extensions. This is where DataKinds comes in.

In haskell, you have types. Each type has kind *, where kind is some sort of type of a type. The type of Show a is Constraint, if you’ve ever seen that pop up in your error messages. What DataKinds allows us to do is promote every type into a new kind and each of its data constructors into a type with that kind. That is, if data Bool = False | True, we now have kind Bool and types 'False and 'Bool of kind Bool. Of course, there is no real way to instantiate values of these types but we can keep these around at the type level to provide information and it will all get compiled out before runtime. Internally, GHC also defines a kind for natural numbers, Nat. If we turn on KindSignatures, we can now specify what kinds we want our types to be, much like how we specify what types we want our values to be. So we conjure up from the void.

data Part = PartA | PartB
  deriving (Eq, Show, Generic)

instance Hashable Part

class Solved (n :: Nat) (p :: Part) where
  solve :: String -> String

We also turn on DeriveGeneric because it turns out that we can automatically get Hashable if we have Generic, and my solver map uses HashMap. We also need MultiParamTypeClasses because, well, Solved is now obviously a multi parameter type class. Since Advent of Code also has two parts per day, we materialize the Part type as well. We’ve now solved the problem of having to import the solver function and change its number in our preamble. We can now do this, which will import only the instances.

-- AOC001.hs
instance Solved 1 'PartA where
  solve _ = "i don't know"
-- Solver.hs
import AOC001 ()
import AOC002 ()
-- ...
import AOC099 ()
-- ...
solvers = H.fromList
  [ ((1, PartA), solve @1 @'PartA), ((1, PartB), solve @1 @'PartB)
  , ((2, PartA), solve @2 @'PartA), ((2, PartB), solve @2 @'PartB)
  -- ...
  ]

So it appears as if we’ve shifted our work elsewhere. But now, template haskell comes to the rescue. We can now ask template haskell to obtain the typeclass definition and the in scope instances for any typeclass name we give it. Within the typeclass definition we get the type of the function we wish to build a map out of, and we can traverse the instances to produce a splice that will generate the map by applying the function to the correct types, with the correct key. Here it is in action.

generateMap :: Name -> Name -> String -> Q [Dec]
generateMap tc fn name = do
  ClassI dec i <- reify tc
  let name' = mkName name
      hmapT = ConT ''HashMap
      -- HashMap (Integer, Part) (String -> String)
      sig = hmapT `AppT` tcKey dec `AppT` tcFnSig dec fn
      -- HashMap.fromList [...]
      body = NormalB . AppE (VarE 'fromList) . ListE $ do
        InstanceD _ _ tyBindings _ <- i
        let solverFn = tyApply (VarE fn) tyBindings
        -- ((num, part), fn)
        return $ TupE [Just (TupE $ fnKey tyBindings), Just solverFn]
      rest = []
  return [SigD name' sig, ValD (VarP name') body rest]

tcKey :: Dec -> Type
tcKey (ClassD _ _ v _ _) = foldl' AppT (TupleT (length v)) (f <$> v)
  where f (PlainTV n) = ConT n
        f (KindedTV _ k) = if k == ConT ''Nat then ConT ''Integer
                           else k
tcKey _ = error "expected typeclass definition but got something else"

tcFnSig :: Dec -> Name -> Type
tcFnSig (ClassD _ _ _ _ d) fn = go d
  where go [] = error $ "could not find " <> show fn
        go (SigD n sig:ds) = if fn == n then sig else go ds
        go (_:ds) = go ds
tcFnSig _ _ = error "expected typeclass definition but got something else"

tyApply :: Exp -> Type -> Exp
tyApply f (ConT _) = f
tyApply f (AppT sub t) = AppTypeE (tyApply f sub) t
tyApply _ _ = error "wat"

fnKey :: Type -> [Maybe Exp]
fnKey (ConT _) = []
fnKey (AppT sub t) = fnKey sub <> [Just t']
  where t' = case t of
               PromotedT n -> ConE n
               LitT (NumTyLit n) -> LitE (IntegerL n)
               LitT (StrTyLit s) -> LitE (StringL s)
               _ -> error "wut"
fnKey _ = error "wat"

tcKey will look at the function declaration and convert it to the correct tuple type for our map. tcFnSig will pull out the type of the function we asked for. tyApply will apply the type variables from the instance to the function we want. Finally, fnKey takes the type variables and converts it into a tuple suitable for the map key. We can now write

generateMap ''Solved 'solve "solvers"

to generate the solvers hashmap which will apply the correct type variables to the solve function from the Solved typeclass. Pretty cool, huh.