datatype/MkCons.hs

   1 {-# LANGUAGE TemplateHaskell #-}
   2 module MkCons where
   3
   4 import Language.Haskell.TH.Syntax
   5 import Language.Haskell.TH
   6 import Data.Char
   7 import Control.Monad
   8
   9 className :: Name -> Name
  10 className = mkName . (++"'") . stringName
  11 constructorName :: Name -> Name
  12 constructorName = mkName . map toLower . stringName
  13 selectorName :: Name -> Name
  14 selectorName = mkName . map toLower . (++"'") . stringName
  15 stringName :: Name -> String
  16 stringName (Name occ _) = occString occ
  17
  18 numberedArgs :: [a] -> [String]
  19 numberedArgs = zipWith (\i _->"f" ++ show i) [0 :: Int ..]
  20
  21 mkConsClass :: Name -> DecsQ
  22 mkConsClass typename = reify typename >>= \info->case info of
  23     TyConI dec
  24         -> case dec of
  25         DataD _ _ tyvars _ constructors _
  26             -> sequence
  27                 [ mkDerivation tyvars
  28                 , mkConstructorClasses tyvars constructors
  29                 , mkPrinterInstances tyvars constructors
  30                 , mkCompilerInstances tyvars constructors
  31                 ]
  32         _
  33             -> fail "mkConsClass only supports data types"
  34     _
  35         -> fail "mkConsClass only supports types"
  36   where
  37     mkDerivation :: [TyVarBndr] -> DecQ
  38     mkDerivation tyvars = let names = (map (VarT . getNameTyVarBndr) tyvars) in pure $
  39         InstanceD Nothing
  40             [ConT (mkName "Serialise") `AppT` t | t <- names]
  41             (ConT (mkName "Serialise") `AppT` foldl AppT (ConT typename) names)
  42             []
  43
  44     mkConstructorClasses :: [TyVarBndr] -> [Con] -> DecQ
  45     mkConstructorClasses tyvars constructors = do
  46         cclasses <- mapM mkConstructorClassMember constructors
  47         sclasses <- concat <$> mapM mkSelectorClassMember constructors
  48         pure $ ClassD [] (className typename) [PlainTV view] [] (cclasses ++ sclasses)
  49       where
  50         view = mkName "m"
  51
  52         mkConstructorClassMember :: Con -> DecQ
  53         mkConstructorClassMember (NormalC consname fs)
  54             = mkConstructorClassMemberForName consname [t | (_, t)<-fs]
  55         mkConstructorClassMember (RecC consname fs)
  56             = mkConstructorClassMemberForName consname [t | (_, _, t)<-fs]
  57         mkConstructorClassMember t
  58             = fail $ "mkConsClass not supported for types such as: " ++ show t
  59
  60         mkConstructorClassMemberForName :: Name -> [Type] -> DecQ
  61         mkConstructorClassMemberForName consname fs
  62             = pure $ SigD (constructorName consname)
  63                 $ foldr (AppT . AppT ArrowT) resultT
  64                 $ map (AppT $ VarT view) fs
  65
  66         mkSelectorClassMember :: Con -> DecsQ
  67         mkSelectorClassMember (NormalC _ fs)
  68             = mapM (uncurry mkSelectorClassMemberForField)
  69             $ zipWith (\(_, t) i->(mkName $ map toLower (stringName typename) ++ "f" ++ show i, t)) fs [0 :: Int ..]
  70         mkSelectorClassMember (RecC _ fs)
  71             = mapM (uncurry mkSelectorClassMemberForField) [(n, t)|(n, _, t)<-fs]
  72         mkSelectorClassMember t
  73             = fail $ "mkConsClass not supported for types such as: " ++ show t
  74
  75         mkSelectorClassMemberForField :: Name -> Type -> DecQ
  76         mkSelectorClassMemberForField n t = pure
  77             $ SigD (className n)
  78             $ ForallT [] [ConT (mkName "Serialise") `AppT` VarT (getNameTyVarBndr ty) | ty <- tyvars]
  79             $ ArrowT `AppT` resultT `AppT` (VarT view `AppT` t)
  80
  81         resultT = VarT view `AppT` (foldl AppT (ConT typename) $ map (VarT . getNameTyVarBndr) $ tyvars)
  82
  83     mkPrinterInstances :: [TyVarBndr] -> [Con] -> DecQ
  84     mkPrinterInstances _ constructors
  85         =   InstanceD Nothing [] (AppT (ConT $ className typename) (ConT $ mkName "Printer")) . concat
  86         <$> mapM mkPrinterInstance constructors
  87       where
  88         mkPrinterInstance :: Con -> DecsQ
  89         mkPrinterInstance (NormalC name fs)
  90             | null fs = pure [FunD (constructorName name)
  91                     [Clause [] (NormalB $ VarE (mkName "printLit") `AppE` LitE (StringL $ stringName name)) [] ]]
  92             | otherwise =
  93                 let args = map mkName $ numberedArgs fs
  94                 in (:) <$> pure (FunD (constructorName name)
  95                     [Clause
  96                         (map VarP args)
  97                         (NormalB $
  98                             (AppE (VarE (mkName "printCons") `AppE` LitE (StringL $ stringName name))
  99                                 (foldl (\x y->x `pc` pl " " `pc` y) (pl "") $ map VarE args)
 100                             )
 101                         )
 102                         []
 103                     ])
 104                 <*> mapM mkPrinterSelector
 105                     (zipWith (\_ i->map toLower (stringName typename) ++ "f" ++ show i) fs [0 :: Int ..])
 106         mkPrinterInstance (RecC name fs)
 107             = let args = map mkName $ numberedArgs fs
 108             in (:) <$> pure (FunD (constructorName name)
 109                 [Clause
 110                     (map VarP args)
 111                     (NormalB $
 112                         (AppE (VarE (mkName "printRec") `AppE` LitE (StringL $ stringName typename))
 113                             (foldl1 (\x y->x `pc` pl ", " `pc` y)
 114                                 $ zipWith (\x ((Name occ'' _), _, _)->pl (occString occ'') `pc` pl "=" `pc` VarE x) args fs)
 115                         )
 116                     )
 117                     []
 118                 ])
 119             <*> mapM mkPrinterSelector [occString occ | (Name occ _, _, _)<-fs]
 120         mkPrinterInstance t = fail $ "mkConsClass not supported for types such as: " ++ show t
 121
 122         mkPrinterSelector :: String -> Q Dec
 123         mkPrinterSelector n' = do
 124             body <- [|\d->d >> printLit $(pure $ LitE $ StringL ('.':n'))|]
 125             pure $ FunD (selectorName $ mkName n') [Clause [] (NormalB body) []]
 126
 127     mkCompilerInstances :: [TyVarBndr] -> [Con] -> DecQ
 128     mkCompilerInstances _ constructors
 129         =   InstanceD Nothing [] (AppT (ConT $ className typename) (ConT $ mkName "Compiler")) . concat
 130         <$> mapM (uncurry mkCompilerInstance) (zip constructors [0..])
 131       where
 132         mkCompilerInstance :: Con -> Int -> DecsQ
 133         mkCompilerInstance (NormalC name fs) consnum = (:)
 134             <$> mkCompilerInstanceForName name consnum (numberedArgs fs)
 135             <*> mapM (uncurry mkCompilerSelector) (zip [0..] [map toLower (stringName typename) ++ f | f<-numberedArgs fs])
 136         mkCompilerInstance (RecC name fs) consnum = (:)
 137             <$> mkCompilerInstanceForName name consnum [occString occ | (Name occ _, _, _) <- fs]
 138             <*> mapM (uncurry mkCompilerSelector) (zip [0..] [occString occ | (Name occ _, _, _) <- fs])
 139         mkCompilerInstance t _ = fail $ "mkConsClass not supported for types such as: " ++ show t
 140
 141         mkCompilerInstanceForName :: Name -> Int -> [String] -> DecQ
 142         mkCompilerInstanceForName name consnum fs =
 143             let args = map mkName $ numberedArgs fs
 144             in do
 145                 body <- [|instr [Push $(pure $ LitE $ IntegerL $ toInteger $ consnum)] *> $(mkBody $ map VarE args) <* instr [Sth $ $(pure $ LitE $ IntegerL $ toInteger $ length fs) + 1] |]
 146                 pure $ FunD (constructorName name)
 147                     [Clause (map VarP args) (NormalB body) [] ]
 148           where
 149             mkBody :: [Exp] -> Q Exp
 150             mkBody [] = pure $ VarE (mkName "pure") `AppE` ConE name
 151             mkBody (a:as) = foldM (\x y->pure $ ifx "<*>" x y) (ifx "<$>" (ConE name) a) as
 152
 153         mkCompilerSelector :: Int -> String -> DecQ
 154         mkCompilerSelector idx n' = do
 155             body <- [|\d->d >> instr [Ldh $(pure $ LitE (IntegerL (toInteger idx)))] |]
 156             pure $ FunD (selectorName $ mkName n')
 157                 [Clause [] (NormalB body) [] ]
 158
 159 instrE :: Exp -> Exp
 160 instrE e = VarE (mkName "instr") `AppE` ListE [e]
 161
 162 ifx :: String -> Exp -> Exp -> Exp
 163 ifx op a b = InfixE (Just a) (VarE $ mkName op) (Just b)
 164
 165 pc :: Exp -> Exp -> Exp
 166 pc l r = VarE (mkName ">>>") `AppE` l `AppE` r
 167
 168 pl :: String -> Exp
 169 pl s = VarE (mkName "printLit") `AppE` LitE (StringL s)
 170
 171 getNameTyVarBndr :: TyVarBndr -> Name
 172 getNameTyVarBndr (PlainTV name) = name
 173 getNameTyVarBndr (KindedTV name _) = name