species-0.3.2: Math/Combinatorics/Species/TH.hs
{-# LANGUAGE NoImplicitPrelude
, TemplateHaskell
, FlexibleInstances
, TypeSynonymInstances
, TypeFamilies
, PatternGuards
, DeriveDataTypeable
#-}
{- Refactoring plan:
* need function to compute a (default) species from a Struct.
- currently have structToSp :: Struct -> Q Exp.
- [X] refactor it into two pieces, Struct -> SpeciesAST and SpeciesAST -> Q Exp.
* should really go through and add some comments to things!
Unfortunately I wasn't good about that when I wrote the code... =P
* Maybe need to do a similar refactoring of the structToTy stuff?
* make version of deriveSpecies that takes a SpeciesAST as an argument,
and use Struct -> SpeciesAST to generate default
* deriveSpecies should pass the SpeciesAST to... other things that
currently just destruct the Struct to decide what to do. Will have to
pattern-match on both the species and the Struct now and make sure
that they match, which is a bit annoying, but can't really be helped.
-}
-----------------------------------------------------------------------------
-- |
-- Module : Math.Combinatorics.Species.CycleIndex
-- Copyright : (c) Brent Yorgey 2010
-- License : BSD-style (see LICENSE)
-- Maintainer : byorgey@cis.upenn.edu
-- Stability : experimental
--
-- Use Template Haskell to automatically derive species instances for
-- user-defined data types.
--
-----------------------------------------------------------------------------
module Math.Combinatorics.Species.TH
( deriveDefaultSpecies
, deriveSpecies
) where
import NumericPrelude
import PreludeBase hiding (cycle)
import Math.Combinatorics.Species.Class
import Math.Combinatorics.Species.Enumerate
import Math.Combinatorics.Species.Structures
import Math.Combinatorics.Species.AST
import Math.Combinatorics.Species.AST.Instances () -- only import instances
import Control.Arrow (first, second, (***))
import Control.Monad (zipWithM, liftM2, mapM, ap)
import Control.Applicative (Applicative(..), (<$>), (<*>))
import Data.Char (toLower)
import Data.Maybe (isJust)
import Data.Typeable
import Language.Haskell.TH
import Language.Haskell.TH.Syntax (lift)
------------------------------------------------------------
-- Preliminaries -----------------------------------------
------------------------------------------------------------
instance Applicative Q where
pure = return
(<*>) = ap
-- | Report a fatal error and stop processing in the 'Q' monad.
errorQ :: String -> Q a
errorQ msg = report True msg >> error msg
------------------------------------------------------------
-- Parsing type declarations -----------------------------
------------------------------------------------------------
-- XXX possible improvement: add special cases to Struct for things
-- like Bool, Either, and (,)
-- | A data structure to represent data type declarations.
data Struct = SId
| SList
| SConst Type -- ^ for types of kind *
| SEnum Type -- ^ for Enumerable type constructors of kind (* -> *)
| SSumProd [(Name, [Struct])] -- ^ sum-of-products
| SComp Struct Struct -- ^ composition
| SSelf -- ^ recursive occurrence
deriving Show
-- | Extract the relevant information about a type constructor into a
-- 'Struct'.
nameToStruct :: Name -> Q Struct
nameToStruct nm = reify nm >>= infoToStruct
where infoToStruct (TyConI d) = decToStruct nm d
infoToStruct _ = errorQ (show nm ++ " is not a type constructor.")
-- XXX do something with contexts? Later extension...
-- | Extract the relevant information about a data type declaration
-- into a 'Struct', given the name of the type and the declaraion.
decToStruct :: Name -> Dec -> Q Struct
decToStruct _ (DataD _ nm [bndr] cons _)
= SSumProd <$> mapM (conToStruct nm (tyVarNm bndr)) cons
decToStruct _ (NewtypeD _ nm [bndr] con _)
= SSumProd . (:[]) <$> conToStruct nm (tyVarNm bndr) con
decToStruct _ (TySynD nm [bndr] ty)
= tyToStruct nm (tyVarNm bndr) ty
decToStruct nm _
= errorQ $ "Processing " ++ show nm ++ ": Only type constructors of kind * -> * are supported."
-- | Throw away kind annotations to extract the type variable name.
tyVarNm :: TyVarBndr -> Name
tyVarNm (PlainTV n) = n
tyVarNm (KindedTV n _) = n
-- | Extract relevant information about a data constructor. The first
-- two arguments are the name of the type constructor, and the name
-- of its type argument. Returns the name of the data constructor
-- and a list of descriptions of its arguments.
conToStruct :: Name -> Name -> Con -> Q (Name, [Struct])
conToStruct nm var (NormalC cnm tys)
= (,) cnm <$> mapM (tyToStruct nm var) (map snd tys)
conToStruct nm var (RecC cnm tys)
= (,) cnm <$> mapM (tyToStruct nm var) (map thrd tys)
where thrd (_,_,t) = t
conToStruct nm var (InfixC ty1 cnm ty2)
= (,) cnm <$> mapM (tyToStruct nm var) [snd ty1, snd ty2]
-- XXX do something with ForallC?
-- XXX check this...
-- | Extract a 'Struct' describing an arbitrary type.
tyToStruct :: Name -> Name -> Type -> Q Struct
tyToStruct nm var (VarT v) | v == var = return SId
| otherwise = errorQ $ "Unknown variable " ++ show v
tyToStruct nm var ListT = return SList
tyToStruct nm var t@(ConT b)
| b == ''[] = return SList
| otherwise = return $ SConst t
tyToStruct nm var (AppT t (VarT v)) -- F `o` TX === F
| v == var && t == (ConT nm) = return $ SSelf -- recursive occurrence
| v == var = return $ SEnum t -- t had better be Enumerable
| otherwise = errorQ $ "Unknown variable " ++ show v
tyToStruct nm var (AppT t1 t2@(AppT _ _)) -- composition
= SComp <$> tyToStruct nm var t1 <*> tyToStruct nm var t2
tyToStruct nm vars t@(AppT _ _)
= return $ SConst t
-- XXX add something to deal with tuples?
-- XXX add something to deal with things that are actually OK like Either a [a]
-- and so on
-- XXX deal with arrow types?
------------------------------------------------------------
-- Misc Struct utilities ---------------------------------
------------------------------------------------------------
-- | Decide whether a type is recursively defined, given its
-- description.
isRecursive :: Struct -> Bool
isRecursive (SSumProd cons) = any isRecursive (concatMap snd cons)
isRecursive (SComp s1 s2) = isRecursive s1 || isRecursive s2
isRecursive SSelf = True
isRecursive _ = False
------------------------------------------------------------
-- Generating default species ----------------------------
------------------------------------------------------------
-- | Convert a 'Struct' into a default corresponding species.
structToSp :: Struct -> SpeciesAST
structToSp SId = X
structToSp SList = L
structToSp (SConst (ConT t))
| t == ''Bool = N 2
| otherwise = error $ "structToSp: unrecognized type " ++ show t ++ " in SConst"
structToSp (SEnum t) = error "SEnum in structToSp"
structToSp (SSumProd []) = Zero
structToSp (SSumProd ss) = foldl1 (+) $ map conToSp ss
structToSp (SComp s1 s2) = structToSp s1 `o` structToSp s2
structToSp SSelf = Omega
-- | Convert a data constructor and its arguments into a default
-- species.
conToSp :: (Name, [Struct]) -> SpeciesAST
conToSp (_,[]) = One
conToSp (_,ps) = foldl1 (*) $ map structToSp ps
------------------------------------------------------------
-- Generating things from species ------------------------
------------------------------------------------------------
-- | Given a name to use in recursive occurrences, convert a species
-- AST into an actual splice-able expression of type Species s => s.
spToExp :: Name -> SpeciesAST -> Q Exp
spToExp self = spToExp'
where
spToExp' Zero = [| 0 |]
spToExp' One = [| 1 |]
spToExp' (N n) = lift n
spToExp' X = [| singleton |]
spToExp' E = [| set |]
spToExp' C = [| cycle |]
spToExp' L = [| linOrd |]
spToExp' Subset = [| subset |]
spToExp' (KSubset k) = [| ksubset $(lift k) |]
spToExp' Elt = [| element |]
spToExp' (f :+: g) = [| $(spToExp' f) + $(spToExp' g) |]
spToExp' (f :*: g) = [| $(spToExp' f) * $(spToExp' g) |]
spToExp' (f :.: g) = [| $(spToExp' f) `o` $(spToExp' g) |]
spToExp' (f :><: g) = [| $(spToExp' f) >< $(spToExp' g) |]
spToExp' (f :@: g) = [| $(spToExp' f) @@ $(spToExp' g) |]
spToExp' (Der f) = [| oneHole $(spToExp' f) |]
spToExp' (OfSize _ _) = error "Can't reify general size predicate into code"
spToExp' (OfSizeExactly f k) = [| $(spToExp' f) `ofSizeExactly` $(lift k) |]
spToExp' (NonEmpty f) = [| nonEmpty $(spToExp' f) |]
spToExp' (Rec _) = [| wrap $(varE self) |]
spToExp' Omega = [| wrap $(varE self) |]
-- | Generate the structure type for a given species.
spToTy :: Name -> SpeciesAST -> Q Type
spToTy self = spToTy'
where
spToTy' Zero = [t| Void |]
spToTy' One = [t| Unit |]
spToTy' (N n) = [t| Const Integer |] -- was finTy n, but that
-- doesn't match up with the
-- type annotation on TSpeciesAST
spToTy' X = [t| Id |]
spToTy' E = [t| Set |]
spToTy' C = [t| Cycle |]
spToTy' L = [t| [] |]
spToTy' Subset = [t| Set |]
spToTy' (KSubset _) = [t| Set |]
spToTy' Elt = [t| Id |]
spToTy' (f :+: g) = [t| Sum $(spToTy' f) $(spToTy' g) |]
spToTy' (f :*: g) = [t| Prod $(spToTy' f) $(spToTy' g) |]
spToTy' (f :.: g) = [t| Comp $(spToTy' f) $(spToTy' g) |]
spToTy' (f :><: g) = [t| Prod $(spToTy' f) $(spToTy' g) |]
spToTy' (f :@: g) = [t| Comp $(spToTy' f) $(spToTy' g) |]
spToTy' (Der f) = [t| Star $(spToTy' f) |]
spToTy' (OfSize f _) = spToTy' f
spToTy' (OfSizeExactly f _) = spToTy' f
spToTy' (NonEmpty f) = spToTy' f
spToTy' (Rec _) = varT self
spToTy' Omega = varT self
{-
-- | Generate a finite type of a given size, using a binary scheme.
finTy :: Integer -> Q Type
finTy 0 = [t| Void |]
finTy 1 = [t| Unit |]
finTy 2 = [t| Const Bool |]
finTy n | even n = [t| Prod (Const Bool) $(finTy $ n `div` 2) |]
| otherwise = [t| Sum Unit $(finTy $ pred n) |]
-}
------------------------------------------------------------
-- Code generation ---------------------------------------
------------------------------------------------------------
-- Enumerable ----------------
-- | Generate an instance of the Enumerable type class, i.e. an
-- isomorphism from the user's data type and the structure type
-- corresponding to the chosen species (or to the default species if
-- the user did not specify one).
--
-- If the third argument is @Nothing@, generate a normal
-- non-recursive instance. If the third argument is @Just code@,
-- then the instance is for a recursive type with the given code.
mkEnumerableInst :: Name -> SpeciesAST -> Struct -> Maybe Name -> Q Dec
mkEnumerableInst nm sp st code = do
clauses <- mkIsoClauses (isJust code) sp st
let stTy = case code of
Just cd -> [t| Mu $(conT cd) |]
Nothing -> spToTy undefined sp -- undefined is OK, it isn't recursive
-- so won't use that argument
instanceD (return []) (appT (conT ''Enumerable) (conT nm))
[ tySynInstD ''StructTy [conT nm] stTy
, return $ FunD 'iso clauses
]
-- | Generate the clauses for the definition of the 'iso' method in
-- the 'Enumerable' instance, which translates from the structure
-- type of the species to the user's data type. The first argument
-- indicates whether the type is recursive.
mkIsoClauses :: Bool -> SpeciesAST -> Struct -> Q [Clause]
mkIsoClauses isRec sp st = (fmap.map) (mkClause isRec) (mkIsoMatches sp st)
where mkClause False (pat, exp) = Clause [pat] (NormalB $ exp) []
mkClause True (pat, exp) = Clause [ConP 'Mu [pat]] (NormalB $ exp) []
mkIsoMatches :: SpeciesAST -> Struct -> Q [(Pat, Exp)]
mkIsoMatches _ SId = newName "x" >>= \x ->
return [(ConP 'Id [VarP x], VarE x)]
mkIsoMatches _ (SConst t)
| t == ConT ''Bool = return [(ConP 'Const [LitP $ IntegerL 1], ConE 'False)
,(ConP 'Const [LitP $ IntegerL 2], ConE 'True)]
| otherwise = error "mkIsoMatches: unrecognized type in SConst case"
mkIsoMatches _ (SEnum t) = newName "x" >>= \x ->
return [(VarP x, AppE (VarE 'iso) (VarE x))]
mkIsoMatches _ (SSumProd []) = return []
mkIsoMatches sp (SSumProd [con]) = mkIsoConMatches sp con
mkIsoMatches sp (SSumProd cons) = addInjs 0 <$> zipWithM mkIsoConMatches (terms sp) cons
where terms (f :+: g) = terms f ++ [g]
terms f = [f]
addInjs :: Int -> [[(Pat, Exp)]] -> [(Pat, Exp)]
addInjs n [ps] = map (addInj (n-1) 'Inr) ps
addInjs n (ps:pss) = map (addInj n 'Inl) ps ++ addInjs (n+1) pss
addInj 0 c = first (ConP c . (:[]))
addInj n c = first (ConP 'Inr . (:[])) . addInj (n-1) c
-- XXX the below is not correct...
-- should really do iso1 . fmap iso2 where iso1 = ... iso2 = ...
-- which are obtained from recursive calls.
mkIsoMatches _ (SComp s1 s2) = newName "x" >>= \x ->
return [ (ConP 'Comp [VarP x]
, AppE (VarE 'iso) (AppE (AppE (VarE 'fmap) (VarE 'iso)) (VarE x))) ]
mkIsoMatches _ SSelf = newName "s" >>= \s ->
return [(VarP s, AppE (VarE 'iso) (VarE s))]
mkIsoConMatches :: SpeciesAST -> (Name, [Struct]) -> Q [(Pat, Exp)]
mkIsoConMatches _ (cnm, []) = return [(ConP 'Unit [], ConE cnm)]
mkIsoConMatches sp (cnm, ps) = map mkProd . sequence <$> zipWithM mkIsoMatches (factors sp) ps
where factors (f :*: g) = factors f ++ [g]
factors f = [f]
mkProd :: [(Pat, Exp)] -> (Pat, Exp)
mkProd = (foldl1 (\x y -> (ConP 'Prod [x, y])) *** foldl AppE (ConE cnm))
. unzip
-- Species definition --------
-- | Given a name n, generate the declaration
--
-- > n :: Species s => s
--
mkSpeciesSig :: Name -> Q Dec
mkSpeciesSig nm = sigD nm [t| Species s => s |]
-- XXX can this use quasiquoting?
-- | Given a name n and a species, generate a declaration for it of
-- that name. The third parameter indicates whether the species is
-- recursive, and if so what the name of the code is.
mkSpecies :: Name -> SpeciesAST -> Maybe Name -> Q Dec
mkSpecies nm sp (Just code) = valD (varP nm) (normalB (appE (varE 'rec) (conE code))) []
mkSpecies nm sp Nothing = valD (varP nm) (normalB (spToExp undefined sp)) []
{-
structToSpAST :: Name -> Struct -> Q Exp
structToSpAST _ SId = [| TX |]
structToSpAST _ (SConst t) = error "SConst in structToSpAST?"
structToSpAST self (SEnum t) = typeToSpAST self t
structToSpAST _ (SSumProd []) = [| TZero |]
structToSpAST self (SSumProd ss) = foldl1 (\x y -> [| annI $x :+: annI $y |])
$ map (conToSpAST self) ss
structToSpAST self (SComp s1 s2) = [| annI $(structToSpAST self s1) :.: annI $(structToSpAST self s2) |]
structToSpAST self SSelf = varE self
conToSpAST :: Name -> (Name, [Struct]) -> Q Exp
conToSpAST _ (_,[]) = [| TOne |]
conToSpAST self (_,ps) = foldl1 (\x y -> [| annI $x :*: annI $y |]) $ map (structToSpAST self) ps
typeToSpAST :: Name -> Type -> Q Exp
typeToSpAST _ ListT = [| TL |]
typeToSpAST self (ConT c) | c == ''[] = [| TL |]
| otherwise = nameToStruct c >>= structToSpAST self -- XXX this is wrong! Need to do something else for recursive types?
typeToSpAST _ _ = error "non-constructor in typeToSpAST?"
-}
------------------------------------------------------------
-- Putting it all together -------------------------------
------------------------------------------------------------
-- XXX need to add something to check whether the type and given
-- species are compatible.
-- | Generate default species declarations for the given user-defined
-- data type. To use it:
--
-- > {-# LANGUAGE TemplateHaskell,
-- > TypeFamilies,
-- > DeriveDataTypeable,
-- > FlexibleInstances,
-- > UndecidableInstances #-}
-- >
-- > data MyType = ...
-- >
-- > $(deriveDefaultSpecies ''MyType)
--
-- Yes, you really do need all those extensions. And don't panic
-- about the @UndecidableInstances@; the instances generated
-- actually are decidable, but GHC just can't tell.
--
-- This is what you get:
--
-- * An 'Enumerable' instance for @MyType@ (and various other
-- supporting things like a code and an 'ASTFunctor' instance if
-- your data type is recursive)
--
-- * A declaration of @myType :: Species s => s@ (the same name as
-- the type constructor but with the first letter lowercased)
--
-- You can then use @myType@ in any species expression, or as input
-- to any function expecting a species. For example, to count your
-- data type's distinct shapes, you can do
--
-- > take 10 . unlabeled $ myType
--
deriveDefaultSpecies :: Name -> Q [Dec]
deriveDefaultSpecies nm = do
st <- nameToStruct nm
deriveSpecies nm (structToSp st)
-- | Like 'deriveDefaultSpecies', except that you specify the species
-- expression that your data type should be isomorphic to. Note: this
-- is currently experimental (read: bug-ridden).
deriveSpecies :: Name -> SpeciesAST -> Q [Dec]
deriveSpecies nm sp = do
st <- nameToStruct nm
let spNm = mkName . map toLower . nameBase $ nm
if (isRecursive st)
then mkEnumerableRec nm spNm st sp
else mkEnumerableNonrec nm spNm st sp
where
mkEnumerableRec nm spNm st sp = do
codeNm <- newName (nameBase nm)
self <- newName "self"
let declCode = DataD [] codeNm [] [NormalC codeNm []] [''Typeable]
[showCode] <- [d| instance Show $(conT codeNm) where
show _ = $(lift (nameBase nm))
|]
[interpCode] <- [d| type instance Interp $(conT codeNm) $(varT self)
= $(spToTy self sp)
|]
applyBody <- NormalB <$> [| unwrap $(spToExp self sp) |]
let astFunctorInst = InstanceD [] (AppT (ConT ''ASTFunctor) (ConT codeNm))
[FunD 'apply [Clause [WildP, VarP self] applyBody []]]
[showMu] <- [d| instance Show a => Show (Mu $(conT codeNm) a) where
show = show . unMu
|]
enum <- mkEnumerableInst nm sp st (Just codeNm)
sig <- mkSpeciesSig spNm
spD <- mkSpecies spNm sp (Just codeNm)
return $ [ declCode
, showCode
, interpCode
, astFunctorInst
, showMu
, enum
, sig
, spD
]
mkEnumerableNonrec nm spNm st sp =
sequence
[ mkEnumerableInst nm sp st Nothing
, mkSpeciesSig spNm
, mkSpecies spNm sp Nothing
]