TH.hs

module Data.Record.Label.TH (mkLabels) where

import Control.Monad (liftM)
import Data.Char (toLower, toUpper)
import Language.Haskell.TH ( Body (NormalB)
                           , Clause (Clause)
                           , Con (RecC)
                           , Dec (DataD, FunD)
                           , Exp (AppE, ConE, LamE, RecUpdE, VarE)
                           , Info (TyConI)
                           , Name
                           , Pat (VarP)
                           , Q
                           , mkName
                           , nameBase
                           , reify)
import Language.Haskell.TH.Syntax (VarStrictType)

mkLabels :: [Name] -> Q [Dec]
mkLabels = liftM concat . mapM mkLabels1

mkLabels1 :: Name -> Q [Dec]
mkLabels1 n = do
    i <- reify n
    let cs = case i of
                 TyConI (DataD _ _ _ cs _) -> cs -- only process data declarations
                 _ -> []
        ls = [ l | (RecC _ ls) <- cs, l <- ls ] -- we're only interested in labels of record constructors
    return $ map mkLabel ls

mkLabel :: VarStrictType -> Dec
mkLabel (name, _, ty) =
    -- Generate a name for the label:
    -- * If the original selector starts with an _, remove it and make
    --   the next character lowercase.
    -- * Otherwise, add 'l', and make the next character uppercase.
    let n = mkName $ case nameBase name of
                ('_' : c : rest) -> toLower c : rest
                (f : rest)   -> 'l' : toUpper f : rest
    in FunD n [Clause [] (NormalB (
           AppE (AppE (ConE (mkName "Label"))
                      (VarE name)) -- getter
                (LamE [VarP (mkName "b"), VarP (mkName "a")] -- setter
                      (RecUpdE (VarE (mkName "a")) [(name, VarE (mkName "b"))]))
                                   )) []]

isRec :: Con -> Bool
isRec (RecC _ _) = True
isRec _          = False