{-# OPTIONS_GHC -fglasgow-exts #-}

module UnivAlg (declareClass) where

import Language.Haskell.TH
import Data.List (union)
import Char


{-  declareClass is a TH function which generates classes from datatypes.
    
    Usage:                  $(declareClass ''T)     for some datatype T with constructors C1 .. Cn
    declares a class:       TClass
    with member functions:  c1 .. cn
    an instance:            TClass T
    and a function:         fromT :: TClass a => T -> a
    
    Class member functions are named by decapitalising the corresponding constructor name, or changing : to % for operators.
    Parametrised datatypes work in the obvious way by applying variables where necessary.
    Only datatypes which can be reifyed will work, so no existentials or GADTs.
-}

declareClass :: Name -> Q [Dec]
declareClass t = do
    reifyT <- reify t
    case reifyT of
        TyConI (DataD cxt _t args cons _) -> do
            let className = mkName (nameBase t ++ "Class")
            p <- newName (drpfst t)
            let classDec = ClassD [] className [p] [] (map (classMember p t args cxt) cons)                 -- the declaration for the class
            let instanceDec = InstanceD [] (ConT className `AppT` ConT t) (map instanceMember cons)         -- the declaration for the instance
            let fromName = mkName ("from" ++ nameBase t)
            let appArgs tp = foldl AppT tp (map VarT args)                                                  -- arguments to reduce tp to kind '*'
            let fromSig = SigD fromName (ForallT (p:args) ((ConT className `AppT` VarT p):cxt)
                                            ((ArrowT `AppT` appArgs (ConT t)) `AppT` appArgs (VarT p)))     -- the signature for the function
            rhs <- mapM (fromMember fromName t) cons                                                        -- the clauses of the declaration for the function
            let fromDec = FunD fromName rhs
            return [classDec,instanceDec,fromSig,fromDec]
        _ -> report True (nameBase t ++ " is not a datatype in 'declareClass "++ nameBase t ++"'") >> return []
        

-- drpfst generates a var name from a con name by changing the first character

drpfst :: Name -> String
drpfst n    | isUpper c = toLower c:cs
            | c == ':'  = '%':cs
                            where (c:cs) = nameBase n
mkdrp = mkName . drpfst


-- subst p t u      replaces all instances of ConT t in u by p

subst :: Type -> Name -> Type -> Type
subst p t (ForallT vs cxt v)
                    | t `elem` vs   = ForallT vs cxt v
                    | otherwise     = ForallT vs cxt (subst p t v)
subst p t (VarT n)                  = VarT n
subst p t (ConT n)  | t == n        = p
                    | otherwise     = ConT n
subst p t (TupleT j)                = TupleT j
subst p t ArrowT                    = ArrowT
subst p t ListT                     = ListT
subst p t (AppT v w)                = AppT (subst p t v) (subst p t w)


-- why isn't this function in the prelude?

thd3 (x,y,z) = z


{-  classMember p t args cxt con    generates the declaration for the class member function corresponding to constructor con
                                    p is the type-class parameter, t is the datatype, args are the arguments to t,
                                    cxt is the context to t.
-}

classMember :: Name -> Name -> [Name] -> Cxt -> Con -> Dec
classMember p t args cxt con    | args == []    = member con
                                | otherwise     = let (SigD dc typ) = member con in SigD dc (ForallT args cxt typ) where
    t0  = foldl AppT (VarT p) (map VarT args)
    member (NormalC c ts)           = SigD (mkdrp c) (foldr (AppT . AppT ArrowT) t0 (map (subst (VarT p) t . snd) ts))
    member (RecC c ts)              = SigD (mkdrp c) (foldr (AppT . AppT ArrowT) t0 (map (subst (VarT p) t . thd3) ts))
    member (InfixC (_,u) c (_,v))   = SigD (mkdrp c) ((ArrowT `AppT` (subst (VarT p) t) u) `AppT` ((ArrowT `AppT` (subst (VarT p) t) v) `AppT` t0))
    member (ForallC vs cxt' con')   = let (SigD dc typ) = member con' in SigD dc (ForallT vs cxt' typ) 


{-  instanceMember con              generates the declaration for the instance member function corresponding to constructor con
-}
    
instanceMember :: Con -> Dec
instanceMember (NormalC c ts)       = ValD (VarP (mkdrp c)) (NormalB (ConE c)) []
instanceMember (RecC c ts)          = ValD (VarP (mkdrp c)) (NormalB (ConE c)) []
instanceMember (InfixC u c v)       = ValD (VarP (mkdrp c)) (NormalB (ConE c)) []


{-  translation j f t type          construct a function of type 'type -> type[a/t]' from 'f::t -> a', if necessary
                                    the integer j records how many arguments are needed to reach kind '*'
-}

fmapList = VarE (mkName "fmap") : [VarE (mkName ('f' : show j ++ "map")) | j <- [2..]]

translation :: Int -> Name -> Name -> Type -> Maybe Exp
translation j f t (ConT n) | t == n = Just (VarE f)
translation j f t (AppT u v)        = case (translation (j+1) f t u, translation 0 f t v) of
                                        (Nothing,Nothing)   -> Nothing
                                        (Nothing,Just x)    -> Just ((fmapList!!j) `AppE` x)
                                        (Just x,Nothing)    -> Just x
                                        (Just x,Just y)     -> Just ((VarE (mkName ".") `AppE` x) `AppE` ((fmapList!!j) `AppE` y))
translation j f t _                 = Nothing


{-  fromType f t type               create a new variable, and apply the translation function to it if necessary
-}

fromType :: Name -> Name -> Type -> Q (Name,Exp)
fromType f t typ = do
    v <- newName "x"
    case (translation 0 f t typ) of
        Nothing -> return (v, VarE v)
        Just x  -> return (v, AppE x (VarE v))


{-  fromMember f t con              generates the clause of the function which pattern matches on constructor con
                                    f is the name of the function, t is the datatype on which the definition must be recursive.
-}

fromMember :: Name -> Name -> Con -> Q Clause
fromMember f t (NormalC c ts)   = do
    vws <- mapM (fromType f t . snd) ts
    let (vs,ws) = unzip vws
    return (Clause [ConP c (map VarP vs)] (NormalB (foldl AppE (VarE (mkdrp c)) ws)) [])
fromMember f t (RecC c ts)      = do
    vws <- mapM (fromType f t . thd3) ts
    let (vs,ws) = unzip vws
    return (Clause [ConP c (map VarP vs)] (NormalB (foldl AppE (VarE (mkdrp c)) ws)) [])
fromMember f t (InfixC (_,u) c (_,v))   = do
    (ux,uy) <- fromType f t u
    (vx,vy) <- fromType f t v
    return (Clause [InfixP (VarP ux) c (VarP vx)] (NormalB (InfixE (Just uy) (VarE (mkdrp c)) (Just vy))) [])