{-# LANGUAGE TemplateHaskell #-}

module Data.Comp.Trans.DeriveTrans
  (
    deriveTrans
  ) where

import Control.Monad.Trans ( lift )
import Data.Map ( Map )

import Control.Lens ( (^.), _1, _2, (%~), view )
import Language.Haskell.TH

import Data.Comp.Multi ( inj, Cxt(Term), (:&:)(..) )

import Data.Comp.Trans.Util

-- |
-- Creates a functions translating from an ADT
-- to its isomorphic multi-sorted compositional data type
-- 
-- @
-- import qualified Foo as F
-- ...
-- type ArithTerm = Term Arith
-- runCompTrans $ deriveTrans [''Arith, ''Atom, ''Lit] (TH.ConT ''ArithTerm)
-- @
-- 
-- will create,
-- 
-- @
-- 
-- class Trans a l where
--   trans :: a -> ArithTerm l
-- 
-- instance Trans F.Arith ArithL where
--   trans (F.Add x y) = iAdd (trans x) (trans y)
-- 
-- instance Trans F.Atom AtomL where
--   trans (F.Var s)   = iVar s
--   trans (F.Const x) = iConst (trans x)
-- 
-- instance Trans F.Lit LitL where
--   trans (F.Lit n) = iLit n
-- @
--
-- With annotation propagation on, it will instead produce, e.g.:
-- `trans :: F.Arith Ann -> Term (Arith :&: Ann) ArithL`
deriveTrans :: [Name] -> Type -> CompTrans [Dec]
deriveTrans :: [Name] -> Type -> CompTrans [Dec]
deriveTrans names :: [Name]
names term :: Type
term = do
  let classNm :: Name
classNm = String -> Name
mkName "Trans"
  Name
funNm <- ReaderT TransCtx Q Name -> CompTrans Name
forall a. ReaderT TransCtx Q a -> CompTrans a
CompTrans (ReaderT TransCtx Q Name -> CompTrans Name)
-> ReaderT TransCtx Q Name -> CompTrans Name
forall a b. (a -> b) -> a -> b
$ Q Name -> ReaderT TransCtx Q Name
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q Name -> ReaderT TransCtx Q Name)
-> Q Name -> ReaderT TransCtx Q Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
newName "trans"

  Dec
classDec <- Name -> Name -> Type -> CompTrans Dec
mkClass Name
classNm Name
funNm Type
term

  Maybe AnnotationPropInfo
annPropInf <- Getting
  (Maybe AnnotationPropInfo) TransCtx (Maybe AnnotationPropInfo)
-> CompTrans (Maybe AnnotationPropInfo)
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting
  (Maybe AnnotationPropInfo) TransCtx (Maybe AnnotationPropInfo)
forall c. HasTransCtx c => Lens' c (Maybe AnnotationPropInfo)
annotationProp
  TransAlts
transAlts <- case Maybe AnnotationPropInfo
annPropInf of
    Just api :: AnnotationPropInfo
api -> AnnotationPropInfo -> CompTrans TransAlts
mkAnnotationPropTransAlts AnnotationPropInfo
api
    Nothing  -> CompTrans TransAlts
mkNormalTransAlts

  [Dec]
instances <- (Name -> CompTrans Dec) -> [Name] -> CompTrans [Dec]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (TransAlts -> Name -> Name -> Name -> CompTrans Dec
mkInstance TransAlts
transAlts Name
classNm Name
funNm) [Name]
names
  [Dec] -> CompTrans [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> CompTrans [Dec]) -> [Dec] -> CompTrans [Dec]
forall a b. (a -> b) -> a -> b
$ [Dec
classDec] [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
instances

data TransAlts = TransAlts {
                             TransAlts -> Name -> Name -> [(Name, Type)] -> Body
makeTransRhs :: Name -> Name -> [(Name, Type)] -> Body -- Fun nm, constructor, variables, types
                           }

mkAnnotationPropTransAlts :: AnnotationPropInfo -> CompTrans TransAlts
mkAnnotationPropTransAlts :: AnnotationPropInfo -> CompTrans TransAlts
mkAnnotationPropTransAlts api :: AnnotationPropInfo
api = do Map Name Type
substs <- Getting (Map Name Type) TransCtx (Map Name Type)
-> CompTrans (Map Name Type)
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (Map Name Type) TransCtx (Map Name Type)
forall c. HasTransCtx c => Lens' c (Map Name Type)
substitutions
                                   TransAlts -> CompTrans TransAlts
forall (m :: * -> *) a. Monad m => a -> m a
return (TransAlts -> CompTrans TransAlts)
-> TransAlts -> CompTrans TransAlts
forall a b. (a -> b) -> a -> b
$ TransAlts :: (Name -> Name -> [(Name, Type)] -> Body) -> TransAlts
TransAlts { makeTransRhs :: Name -> Name -> [(Name, Type)] -> Body
makeTransRhs = Map Name Type
-> AnnotationPropInfo -> Name -> Name -> [(Name, Type)] -> Body
makeTransRhsPropAnn Map Name Type
substs AnnotationPropInfo
api}

mkNormalTransAlts :: CompTrans TransAlts
mkNormalTransAlts :: CompTrans TransAlts
mkNormalTransAlts = TransAlts -> CompTrans TransAlts
forall (m :: * -> *) a. Monad m => a -> m a
return (TransAlts -> CompTrans TransAlts)
-> TransAlts -> CompTrans TransAlts
forall a b. (a -> b) -> a -> b
$ TransAlts :: (Name -> Name -> [(Name, Type)] -> Body) -> TransAlts
TransAlts { makeTransRhs :: Name -> Name -> [(Name, Type)] -> Body
makeTransRhs = Name -> Name -> [(Name, Type)] -> Body
makeTransRhsNormal}

-- |
-- Example:
-- 
-- @
-- translate :: J.CompilationUnit -> JavaTerm CompilationUnitL
-- translate = trans
-- @
mkFunc :: Type -> Name -> Type -> CompTrans [Dec]
mkFunc :: Type -> Name -> Type -> CompTrans [Dec]
mkFunc typ :: Type
typ funNm :: Name
funNm term :: Type
term = do
  Type
srcTyp <- Type -> CompTrans Type
forall x. Data x => x -> CompTrans x
applyCurSubstitutions Type
typ
  Type -> Bool
isAnn <- CompTrans (Type -> Bool)
getIsAnn
  Type
lab <- (Type -> Bool) -> Type -> CompTrans Type
getLab Type -> Bool
isAnn Type
srcTyp
  [Dec] -> CompTrans [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return [ Name -> Type -> Dec
SigD Name
translate (Type -> Type -> Type
AppT (Type -> Type -> Type
AppT Type
ArrowT Type
srcTyp) (Type -> Type -> Type
AppT Type
term Type
lab))
         , Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP Name
translate) (Exp -> Body
NormalB Exp
funNm') []
         ]
  where
    translate :: Name
translate = String -> Name
mkName "translate"
    funNm' :: Exp
funNm' = Name -> Exp
VarE Name
funNm

-- |
-- Example:
-- 
-- @
-- class Trans a l where
--   trans a -> JavaTerm l
-- @
mkClass :: Name -> Name -> Type -> CompTrans Dec
mkClass :: Name -> Name -> Type -> CompTrans Dec
mkClass classNm :: Name
classNm funNm :: Name
funNm term :: Type
term = do Name
a <- ReaderT TransCtx Q Name -> CompTrans Name
forall a. ReaderT TransCtx Q a -> CompTrans a
CompTrans (ReaderT TransCtx Q Name -> CompTrans Name)
-> ReaderT TransCtx Q Name -> CompTrans Name
forall a b. (a -> b) -> a -> b
$ Q Name -> ReaderT TransCtx Q Name
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q Name -> ReaderT TransCtx Q Name)
-> Q Name -> ReaderT TransCtx Q Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
newName "a"
                                Name
i <- ReaderT TransCtx Q Name -> CompTrans Name
forall a. ReaderT TransCtx Q a -> CompTrans a
CompTrans (ReaderT TransCtx Q Name -> CompTrans Name)
-> ReaderT TransCtx Q Name -> CompTrans Name
forall a b. (a -> b) -> a -> b
$ Q Name -> ReaderT TransCtx Q Name
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q Name -> ReaderT TransCtx Q Name)
-> Q Name -> ReaderT TransCtx Q Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
newName "i"
                                let transDec :: Dec
transDec = Name -> Type -> Dec
SigD Name
funNm ((Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT Type
ArrowT [Name -> Type
VarT Name
a, Type -> Type -> Type
AppT Type
term (Name -> Type
VarT Name
i)])
                                Dec -> CompTrans Dec
forall (m :: * -> *) a. Monad m => a -> m a
return (Dec -> CompTrans Dec) -> Dec -> CompTrans Dec
forall a b. (a -> b) -> a -> b
$ [Type] -> Name -> [TyVarBndr] -> [FunDep] -> [Dec] -> Dec
ClassD [] Name
classNm [Name -> TyVarBndr
PlainTV Name
a, Name -> TyVarBndr
PlainTV Name
i] [] [Dec
transDec]

-- |
-- Example:
-- 
-- @
-- instance Trans J.CompilationUnit CompilationUnitL where
--   trans (J.CompilationUnit x y z) = iCompilationUnit (trans x) (trans y) (trans z)
-- @
mkInstance :: TransAlts -> Name -> Name -> Name -> CompTrans Dec
mkInstance :: TransAlts -> Name -> Name -> Name -> CompTrans Dec
mkInstance transAlts :: TransAlts
transAlts classNm :: Name
classNm funNm :: Name
funNm typNm :: Name
typNm = do
  Info
inf <- ReaderT TransCtx Q Info -> CompTrans Info
forall a. ReaderT TransCtx Q a -> CompTrans a
CompTrans (ReaderT TransCtx Q Info -> CompTrans Info)
-> ReaderT TransCtx Q Info -> CompTrans Info
forall a b. (a -> b) -> a -> b
$ Q Info -> ReaderT TransCtx Q Info
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q Info -> ReaderT TransCtx Q Info)
-> Q Info -> ReaderT TransCtx Q Info
forall a b. (a -> b) -> a -> b
$ Name -> Q Info
reify Name
typNm
  Type
srcTyp <- Name -> CompTrans Type
getFullyAppliedType Name
typNm
  let nmTyps :: [(Name, [Type])]
nmTyps = Info -> [(Name, [Type])]
simplifyDataInf Info
inf
  [Clause]
clauses <- ((Name, [Type]) -> CompTrans Clause)
-> [(Name, [Type])] -> CompTrans [Clause]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Name -> [Type] -> CompTrans Clause)
-> (Name, [Type]) -> CompTrans Clause
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Name -> [Type] -> CompTrans Clause)
 -> (Name, [Type]) -> CompTrans Clause)
-> (Name -> [Type] -> CompTrans Clause)
-> (Name, [Type])
-> CompTrans Clause
forall a b. (a -> b) -> a -> b
$ TransAlts -> Name -> Name -> [Type] -> CompTrans Clause
mkClause TransAlts
transAlts Name
funNm) [(Name, [Type])]
nmTyps
  let targNm :: Name
targNm = Name -> Name
nameLab Name
typNm
  Dec -> CompTrans Dec
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
InstanceD
                   Maybe Overlap
forall a. Maybe a
Nothing
                   []
                   (Type -> Type -> Type
AppT (Type -> Type -> Type
AppT (Name -> Type
ConT Name
classNm) Type
srcTyp) (Name -> Type
ConT Name
targNm))
                   [Name -> [Clause] -> Dec
FunD Name
funNm [Clause]
clauses])


atom :: Name -> (Name, Type) -> Exp
atom :: Name -> (Name, Type) -> Exp
atom _     (x :: Name
x, t :: Type
t) | Type -> [Type] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Type
t [Type]
baseTypes = Name -> Exp
VarE Name
x
atom funNm :: Name
funNm (x :: Name
x, _)                    = Exp -> Exp -> Exp
AppE (Name -> Exp
VarE Name
funNm) (Name -> Exp
VarE Name
x)

makeTransRhsPropAnn :: Map Name Type -> AnnotationPropInfo -> Name -> Name -> [(Name, Type)] -> Body
makeTransRhsPropAnn :: Map Name Type
-> AnnotationPropInfo -> Name -> Name -> [(Name, Type)] -> Body
makeTransRhsPropAnn substs :: Map Name Type
substs annPropInf :: AnnotationPropInfo
annPropInf funNm :: Name
funNm con :: Name
con nmTps :: [(Name, Type)]
nmTps = Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Term) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Exp -> Exp -> Exp
AppE (Name -> Exp
ConE '(:&:)) Exp
nodeExp) Exp
annExp
  where
    nmTps' :: [(Name, Type)]
    nmTps' :: [(Name, Type)]
nmTps' = ((Name, Type) -> (Name, Type)) -> [(Name, Type)] -> [(Name, Type)]
forall a b. (a -> b) -> [a] -> [b]
map ((Type -> Identity Type) -> (Name, Type) -> Identity (Name, Type)
forall s t a b. Field2 s t a b => Lens s t a b
_2 ((Type -> Identity Type) -> (Name, Type) -> Identity (Name, Type))
-> (Type -> Type) -> (Name, Type) -> (Name, Type)
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (Map Name Type -> Type -> Type
forall x. Data x => Map Name Type -> x -> x
applySubsts Map Name Type
substs)) [(Name, Type)]
nmTps

    annVar :: (a, Type) -> Bool
    annVar :: (a, Type) -> Bool
annVar (_, t :: Type
t) = (AnnotationPropInfo
annPropInf AnnotationPropInfo
-> Getting (Type -> Bool) AnnotationPropInfo (Type -> Bool)
-> Type
-> Bool
forall s a. s -> Getting a s a -> a
^. Getting (Type -> Bool) AnnotationPropInfo (Type -> Bool)
forall c. HasAnnotationPropInfo c => Lens' c (Type -> Bool)
isAnnotation) Type
t

    nodeExp :: Exp
nodeExp = Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'inj) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Exp -> Exp
AppE (Name -> Exp
ConE (Name -> Name
transName Name
con)) (((Name, Type) -> Exp) -> [(Name, Type)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> (Name, Type) -> Exp
atom Name
funNm) ([(Name, Type)] -> [Exp]) -> [(Name, Type)] -> [Exp]
forall a b. (a -> b) -> a -> b
$ ((Name, Type) -> Bool) -> [(Name, Type)] -> [(Name, Type)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not(Bool -> Bool) -> ((Name, Type) -> Bool) -> (Name, Type) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Name, Type) -> Bool
forall a. (a, Type) -> Bool
annVar) [(Name, Type)]
nmTps')

    annExp :: Exp
annExp = (AnnotationPropInfo
annPropInf AnnotationPropInfo
-> Getting
     ([(Exp, Type)] -> Exp) AnnotationPropInfo ([(Exp, Type)] -> Exp)
-> [(Exp, Type)]
-> Exp
forall s a. s -> Getting a s a -> a
^. Getting
  ([(Exp, Type)] -> Exp) AnnotationPropInfo ([(Exp, Type)] -> Exp)
forall c. HasAnnotationPropInfo c => Lens' c ([(Exp, Type)] -> Exp)
propAnn) (((Name, Type) -> (Exp, Type)) -> [(Name, Type)] -> [(Exp, Type)]
forall a b. (a -> b) -> [a] -> [b]
map ((Name -> Identity Exp) -> (Name, Type) -> Identity (Exp, Type)
forall s t a b. Field1 s t a b => Lens s t a b
_1 ((Name -> Identity Exp) -> (Name, Type) -> Identity (Exp, Type))
-> (Name -> Exp) -> (Name, Type) -> (Exp, Type)
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Name -> Exp
VarE) ([(Name, Type)] -> [(Exp, Type)])
-> [(Name, Type)] -> [(Exp, Type)]
forall a b. (a -> b) -> a -> b
$ ((Name, Type) -> Bool) -> [(Name, Type)] -> [(Name, Type)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Name, Type) -> Bool
forall a. (a, Type) -> Bool
annVar [(Name, Type)]
nmTps')

makeTransRhsNormal :: Name -> Name -> [(Name, Type)] -> Body
makeTransRhsNormal :: Name -> Name -> [(Name, Type)] -> Body
makeTransRhsNormal funNm :: Name
funNm con :: Name
con nmTps :: [(Name, Type)]
nmTps = Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Name
smartConstrName Name
con)) (((Name, Type) -> Exp) -> [(Name, Type)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> (Name, Type) -> Exp
atom Name
funNm) [(Name, Type)]
nmTps)

mkClause :: TransAlts -> Name -> Name -> [Type] -> CompTrans Clause
mkClause :: TransAlts -> Name -> Name -> [Type] -> CompTrans Clause
mkClause transAlts :: TransAlts
transAlts funNm :: Name
funNm con :: Name
con tps :: [Type]
tps = do [Name]
nms <- ReaderT TransCtx Q [Name] -> CompTrans [Name]
forall a. ReaderT TransCtx Q a -> CompTrans a
CompTrans (ReaderT TransCtx Q [Name] -> CompTrans [Name])
-> ReaderT TransCtx Q [Name] -> CompTrans [Name]
forall a b. (a -> b) -> a -> b
$ Q [Name] -> ReaderT TransCtx Q [Name]
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q [Name] -> ReaderT TransCtx Q [Name])
-> Q [Name] -> ReaderT TransCtx Q [Name]
forall a b. (a -> b) -> a -> b
$ (Type -> Q Name) -> [Type] -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Q Name -> Type -> Q Name
forall a b. a -> b -> a
const (Q Name -> Type -> Q Name) -> Q Name -> Type -> Q Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
newName "x") [Type]
tps
                                      Clause -> CompTrans Clause
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause -> CompTrans Clause) -> Clause -> CompTrans Clause
forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [[Name] -> Pat
pat [Name]
nms] (TransAlts -> Name -> Name -> [(Name, Type)] -> Body
makeTransRhs TransAlts
transAlts Name
funNm Name
con ([(Name, Type)] -> Body) -> [(Name, Type)] -> Body
forall a b. (a -> b) -> a -> b
$ [Name] -> [Type] -> [(Name, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
nms [Type]
tps) []
  where
    pat :: [Name] -> Pat
pat nms :: [Name]
nms = Name -> [Pat] -> Pat
ConP Name
con ((Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
nms)