[Haskell-cafe] Compile-time evaluation

[Nicholas Messenger]

{-# OPTIONS_GHC -fglasgow-exts -fno-monomorphism-restriction #-}

-- Many people ask if GHC will evaluate toplevel constants at compile
-- time, you know, since Haskell is pure it'd be great if those
-- computations could be done once and not use up cycles during
-- runtime.  Not an entirely bad idea, I think.
-- 
-- So I set about allowing just that: for arbitrary expressions to be
-- evaluated, and the expanded expression spliced into client code.
-- 
-- If you had some data in a file just out of convenience, you could say:
-- > yourData = $(compileTimeIO $ parseFile $ readFile "data.txt")
-- 
-- Or if you had an expensive computation that you want done at compile:
-- > result = $(compileTimeEval $ expensiveComputation)
-- 
-- I would appreciate comments.  I wrote this completely blind with just
-- the TH and Generics haddocks, so if I'm doing something tremendously
-- stupid that can be improved, let me know. :)  Especially if you can
-- think of a less awkward way to go from Generics' data to TH
-- expressions than using 'showConstr'...
-- 
-- I wrote this with 6.6.1, in case there's any incompatibilities.  Copy/
-- paste this post into CompileTime.hs, load into ghci, :set -fth, and
-- futz around with the splices.
-- 
-- -- Nicholas Messenger (nmessenger at gmail.com, omnId on #haskell)

module CompileTime(compileTimeEval, compileTimeIO) where

import Data.Generics
import Language.Haskell.TH
import Control.Monad
import Data.Tree
import Data.Ratio

-- Expands a datum into an expression tree to be spliced into
-- client code.
compileTimeEval :: Data a => a -> ExpQ
compileTimeEval = return . toExp

-- Runs the IO action and splices in the evaluated result datum.
compileTimeIO :: Data a => IO a -> ExpQ
compileTimeIO = liftM toExp . runIO

-- Does the work. :)  toTree gets us a tree of constructors, so
-- we mostly just have to fold the tree with AppE, except for
-- TH's bizarre TupE.
toExp :: Data d => d -> Exp
toExp = applyAll . toTree
 where
  applyAll (Node k args)
    | isTuple k = TupE (map applyAll args)
    | otherwise = foldl AppE k (map applyAll args)

  isTuple (ConE n) = all (==',') (nameBase n)
  isTuple _        = False

-- Synonym to shorten the definition of exp below
type Ex a = a -> Exp

-- Turns some datum into a tree of TH expressions representing
-- that datum.  The Exp at each node represents the constructor,
-- the subtrees are its arguments.
toTree :: Data d => d -> Tree Exp
toTree x = Node (exp x) (gmapQ toTree x)
 where
  -- The various ways to turn a (Data d => d) into an
  -- Exp representing its constructor.
  any  = ConE . mkName . deparen . showConstr . toConstr
  char = LitE . CharL
  int  = sigged $ LitE . IntegerL . toInteger
  rat  = sigged $ LitE . RationalL . toRational
  sigged f x = SigE (f x) (ConT . mkName . show $ typeOf x)

  -- The above functions combined together for different types.
  -- This is what gives the constructor Exp at each Node.  There
  -- are definitely more types to cover that 'any' gets wrong...
  exp = any `extQ` (int::Ex Int)    `extQ` (int::Ex Integer)
            `extQ` char             `extQ` (rat::Ex Float)
            `extQ` (rat::Ex Double) `extQ` (rat::Ex Rational)

  -- Generics' showConstr puts parens around infix
  -- constructors.  TH's ConE doesn't like 'em.
  deparen s = (if last s == ')' then init else id) .
              (if head s == '(' then tail else id) $ s