module E.Annotate where
import Control.Monad.Reader
import Data.Monoid
import qualified Data.Traversable as T
import E.E
import E.Program
import E.Subst
import GenUtil
import Info.Info(Info)
import Name.Id
import Util.HasSize
import Util.SetLike
annotateCombs :: forall m . Monad m =>
(IdMap (Maybe E))
-> (Id -> Info -> m Info)
-> (E -> Info -> m Info)
-> (E -> Info -> m Info)
-> [Comb]
-> m [Comb]
annotateCombs imap idann letann lamann cs = do
cs <- forM cs $ \comb -> do
nfo <- letann (combBody comb) (tvrInfo $ combHead comb)
nt <- annotate imap idann letann lamann (tvrType $ combHead comb)
return $ combHead_u (tvrInfo_s nfo . tvrType_s nt) comb
let nimap = fromList [ (combIdent c, Just . EVar $ combHead c) | c <- cs ]
`mappend` imap
f :: (IdMap (Maybe E)) -> E -> m E
f ni e = annotate ni idann letann lamann e
let mrule :: Rule -> m Rule
mrule r = do
let g tvr = do
nfo <- idann (tvrIdent tvr) (tvrInfo tvr)
let ntvr = tvr { tvrInfo = nfo }
return (ntvr,minsert (tvrIdent tvr) (Just $ EVar ntvr))
bs <- mapM g $ ruleBinds r
let nnimap = (foldr (.) id $ snds bs) nimap :: IdMap (Maybe E)
args <- mapM (f nnimap) (ruleArgs r)
body <- (f nnimap) (ruleBody r)
return r { ruleBinds = fsts bs, ruleBody = body, ruleArgs = args }
forM cs $ \comb -> do
rs <- mapM mrule (combRules comb)
nb <- f nimap (combBody comb)
return . combRules_s rs . combBody_s nb $ comb
annotateDs :: Monad m =>
(IdMap (Maybe E))
-> (Id -> Info -> m Info)
-> (E -> Info -> m Info)
-> (E -> Info -> m Info)
-> [(TVr,E)]
-> m [(TVr,E)]
annotateDs imap idann letann lamann ds = do
ELetRec { eDefs = ds', eBody = Unknown } <- annotate imap idann letann lamann (ELetRec ds Unknown)
return ds'
annotateProgram :: Monad m =>
(IdMap (Maybe E))
-> (Id -> Info -> m Info)
-> (E -> Info -> m Info)
-> (E -> Info -> m Info)
-> Program
-> m Program
annotateProgram imap idann letann lamann prog = do
ds <- annotateCombs imap idann letann lamann (progCombinators prog)
return $ programUpdate $ prog { progCombinators = ds }
type AM m = ReaderT (IdMap (Maybe E)) m
annotate :: Monad m =>
(IdMap (Maybe E))
-> (Id -> Info -> m Info)
-> (E -> Info -> m Info)
-> (E -> Info -> m Info)
-> E
-> m E
annotate imap idann letann lamann e = runReaderT (f e) imap where
f eo@(EVar tvr@(TVr { tvrIdent = i, tvrType = t })) = do
mp <- ask
case mlookup i mp of
Just (Just v) -> return v
_ -> return eo
f (ELam tvr e) = lp ELam tvr e
f (EPi tvr e) = lp EPi tvr e
f (EAp a b) = liftM2 EAp (f a) (f b)
f (EError x e) = liftM (EError x) (f e)
f (EPrim x es e) = liftM2 (EPrim x) (mapM f es) (f e)
f ELetRec { eDefs = dl, eBody = e } = do
dl' <- flip mapM dl $ \ (t,e) -> do
nfo <- lift $ letann e (tvrInfo t)
return t { tvrInfo = nfo }
(as,rs) <- liftM unzip $ mapMntvr dl'
local (foldr (.) id rs) $ do
ds <- mapM f (snds dl)
e' <- f e
return $ ELetRec (zip as ds) e'
f (ELit l) = liftM ELit $ litSMapM f l
f Unknown = return Unknown
f e@(ESort {}) = return e
f ec@(ECase {}) = do
e' <- f $ eCaseScrutinee ec
let caseBind = eCaseBind ec
(b',r) <- ntvr [] caseBind
d <- local r $ T.mapM f $ eCaseDefault ec
let da (Alt lc@LitCons { litName = s, litArgs = vs, litType = t } e) = do
t' <- f t
(as,rs) <- liftM unzip $ mapMntvr vs
e' <- local (foldr (.) id rs) $ f e
return $ Alt lc { litArgs = as, litType = t' } e'
da (Alt l e) = do
l' <- T.mapM f l
e' <- f e
return $ Alt l' e'
alts <- local r (mapM da $ eCaseAlts ec)
t' <- f (eCaseType ec)
return $ caseUpdate ECase { eCaseAllFV = error "no eCaseAllFV needed", eCaseScrutinee = e', eCaseType = t', eCaseDefault = d, eCaseBind = b', eCaseAlts = alts }
lp lam tvr@(TVr { tvrIdent = n, tvrType = t}) e | n == emptyId = do
t' <- f t
nfo <- lift $ lamann e (tvrInfo tvr)
nfo <- lift $ idann n nfo
e' <- local (minsert n Nothing) $ f e
return $ lam (tvr { tvrIdent = emptyId, tvrType = t', tvrInfo = nfo}) e'
lp lam tvr e = do
nfo <- lift $ lamann e (tvrInfo tvr)
(tv,r) <- ntvr [] tvr { tvrInfo = nfo }
e' <- local r $ f e
return $ lam tv e'
mapMntvr ts = f ts [] where
f [] xs = return $ reverse xs
f (t:ts) rs = do
(t',r) <- ntvr vs t
local r $ f ts ((t',r):rs)
vs = [ tvrIdent x | x <- ts ]
ntvr xs tvr@(TVr { tvrIdent = n, tvrType = t}) | n == emptyId = do
t' <- f t
nfo <- lift $ idann emptyId (tvrInfo tvr)
let nvr = (tvr { tvrType = t', tvrInfo = nfo})
return (nvr,id)
ntvr xs tvr@(TVr {tvrIdent = i, tvrType = t}) = do
t' <- f t
ss <- ask
nfo' <- lift $ idann i (tvrInfo tvr)
let i' = mnv xs i ss
let nvr = (tvr { tvrIdent = i', tvrType = t', tvrInfo = nfo'})
case i == i' of
True -> return (nvr,minsert i (Just $ EVar nvr))
False -> return (nvr,minsert i (Just $ EVar nvr) . minsert i' Nothing)
mnv xs i ss
| isInvalidId i || i `member` ss = newId (size ss) isOkay
| otherwise = i
where isOkay i = (i `notMember` ss) && (i `notElem` xs)