module DDC.Core.Transform.Rewrite
( RewriteRule(..)
, rewriteModule
, rewriteX)
where
import DDC.Base.Pretty
import DDC.Core.Exp
import DDC.Core.Module
import Data.Map (Map)
import DDC.Core.Simplifier.Base (TransformResult(..), TransformInfo(..))
import qualified DDC.Core.Compounds as X
import qualified DDC.Core.Transform.AnonymizeX as A
import qualified DDC.Core.Transform.Rewrite.Disjoint as RD
import qualified DDC.Core.Transform.Rewrite.Env as RE
import qualified DDC.Core.Transform.Rewrite.Match as RM
import DDC.Core.Transform.Rewrite.Rule
import qualified DDC.Core.Transform.SubstituteXX as S
import qualified DDC.Type.Transform.SubstituteT as S
import qualified DDC.Core.Transform.Trim as Trim
import qualified DDC.Core.Transform.LiftX as L
import qualified DDC.Type.Compounds as T
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Maybe
import Control.Monad
import Control.Monad.Writer (tell, runWriter)
import Data.List
import Data.Typeable
data RewriteInfo
= RewriteInfo [RewriteLog]
deriving Typeable
data RewriteLog
= LogRewrite String
| LogUnfold String
deriving Typeable
instance Pretty RewriteInfo where
ppr (RewriteInfo rules)
= text "Rules fired:"
<$> indent 4 (vcat $ map ppr rules)
instance Pretty RewriteLog where
ppr (LogRewrite name) = text "Rewrite: " <> text name
ppr (LogUnfold name) = text "Unfold: " <> text name
isProgress = not . null
rewriteModule
:: (Show a, Show n, Ord n, Pretty n)
=> [NamedRewriteRule a n]
-> Module a n
-> Module a n
rewriteModule rules mm
= mm { moduleBody = result $ rewriteX' True rules $ moduleBody mm }
rewriteX
:: (Show a, Show n, Ord n, Pretty n)
=> [NamedRewriteRule a n]
-> Exp a n
-> TransformResult (Exp a n)
rewriteX = rewriteX' False
rewriteX'
:: (Show a, Show n, Ord n, Pretty n)
=> Bool
-> [NamedRewriteRule a n]
-> Exp a n
-> TransformResult (Exp a n)
rewriteX' ignore_toplevel rules x0
=
let (x,info) = runWriter $ go RE.empty x0 [] ignore_toplevel
in TransformResult
{ result = x
, resultAgain = isProgress info
, resultProgress = isProgress info
, resultInfo = TransformInfo (RewriteInfo info) }
where
rewrites env f args
= rewrites' rules env f args
rewrites' [] _ f args
= return $ X.makeXAppsWithAnnots f args
rewrites' ((n, rule) : rs) env f args
= case rewriteWithX rule env f args of
Nothing -> rewrites' rs env f args
Just x -> tell [LogRewrite n] >> go env x [] False
down env x
= go env x [] False
go env (XApp a f arg) args _toplevel
= do arg' <- down env arg
go env f ((arg',a):args) False
go env x@XVar{} args _toplevel
= rewrites env x args
go env x@XCon{} args _toplevel
= rewrites env x args
go env (XLAM a b e) args _toplevel
= do e' <- down (RE.lift b env) e
rewrites env (XLAM a b e') args
go env (XLam a b e) args _toplevel
= do e' <- down (RE.extend b env) e
rewrites env (XLam a b e') args
go env (XLet a l@(LRec _) e) args toplevel
= do
let env' = if toplevel
then env
else RE.extendLets l env
l' <- goLets l env'
e' <- down env' e
rewrites env' (XLet a l' e') args
go env (XLet a l e) args _toplevel
= do l' <- goLets l env
dh <- goDefHoles rules a l' e env down
rewrites env dh args
go env (XCase a e alts) args _toplevel
= do e' <- down env e
alts' <- mapM (goAlts env) alts
rewrites env (XCase a e' alts') args
go env (XCast a c e) args _toplevel
= do e' <- down env e
rewrites env (XCast a c e') args
go env x@(XType{}) args _toplevel
= rewrites env x args
go env x@(XWitness{}) args _toplevel
= rewrites env x args
goLets (LLet b e) ws
= do e' <- down ws e
return $ LLet b e'
goLets (LRec bs) ws
= do bs' <- mapM (down ws) $ map snd bs
return $ LRec $ zip (map fst bs) bs'
goLets l _
= return $ l
goAlts ws (AAlt p e)
= do e' <- down ws e
return $ AAlt p e'
goDefHoles rules a l@(LLet let_bind def) e env down
| (((sub, []), name, RewriteRule { ruleBinds = bs, ruleLeft = hole }):_)
<- checkHoles rules def env
= let
bs' = filter (isBMValue . fst) bs
bas' = lookupFromSubst bs' sub
isUIx x = case x of
XVar _ (UIx _) -> True
XVar _ (UPrim _ _) -> True
_ -> False
already_done
= all isUIx $ map snd bas'
bsK' = filter ((== BMSpec) . fst) bs
basK = lookupFromSubst bsK' sub
basK' = concatMap (\(b,x) -> case X.takeXType x of
Just t -> [(b,t)]
Nothing-> []) basK
values = map (\(b,v) -> (BAnon (S.substituteTs basK' $ T.typeOfBind b), v))
(reverse bas')
anons = zipWith (\(b,_) i -> (b, XVar a (UIx i))) bas' [0..]
lets = map (\(b,v) -> LLet b v) values
def' = S.substituteXArgs basK
$ S.substituteXArgs anons hole
let_bind' = S.substituteTs basK' let_bind
lets' = lets ++ [LLet let_bind' def']
depth = case let_bind of
BAnon _ -> 1
_ -> 0
e' = L.liftAtDepthX (length bas') depth e
env' = foldl (flip RE.extendLets) env lets'
in if already_done
then do
e'' <- down (RE.extendLets l env) e
return $ XLet a l e''
else do
tell [LogUnfold name]
e'' <- down env' e'
return $ X.xLets a lets' e''
| otherwise
= do e' <- down (RE.extendLets l env) e
return $ XLet a l e'
goDefHoles _rules a l e env down
= do e' <- down (RE.extendLets l env) e
return $ XLet a l e'
checkHoles
:: (Show n, Show a, Ord n, Pretty n)
=> [NamedRewriteRule a n]
-> Exp a n
-> RE.RewriteEnv a n
-> [ ( (RM.SubstInfo a n, [(Exp a n, a)])
, String
, RewriteRule a n) ]
checkHoles rules def ws
= let rules' = catMaybes $ map holeRule rules
(f,args) = X.takeXAppsWithAnnots def
in catMaybes
$ map (\(name,r) -> fmap (\s -> (s,name,r))
$ matchWithRule r ws f args RM.emptySubstInfo)
rules'
holeRule (name, rule@RewriteRule { ruleLeftHole = Just hole })
= Just ( name
, rule { ruleLeft = hole
, ruleLeftHole = Nothing })
holeRule _ = Nothing
rewriteWithX
:: (Show n, Show a, Ord n, Pretty n)
=> RewriteRule a n
-> RE.RewriteEnv a n
-> Exp a n
-> [(Exp a n, a)]
-> Maybe (Exp a n)
rewriteWithX rule env f args
= do
let RewriteRule
{ ruleBinds = binds
, ruleConstraints = constrs
, ruleRight = rhs
, ruleWeakEff = eff
, ruleWeakClo = clo }
= rule
(m, rest) <- matchWithRule rule env f args RM.emptySubstInfo
let Just a = X.takeAnnotOfExp f
let bas2 = lookupFromSubst binds m
let rhs2 = A.anonymizeX rhs
let (bas3,lets) = wrapLets a binds bas2
let rhs3 = L.liftX (length lets) rhs2
let eff' = liftM (substT bas3) eff
let clo' = Trim.trimClosures a
$ map (S.substituteXArgs bas3) clo
let constrs' = map (substT bas3) constrs
when (not $ all (satisfiedContraint env) constrs')
$ Nothing
let x' = X.xLets a lets
$ weakeff a eff'
$ weakclo a clo'
$ S.substituteXArgs bas3 rhs3
return $ X.makeXAppsWithAnnots x' rest
satisfiedContraint :: (Ord n, Show n) => RE.RewriteEnv a n -> Type n -> Bool
satisfiedContraint env c
= RE.containsWitness c env
|| RD.checkDisjoint c env
|| RD.checkDistinct c env
weakeff :: Ord n
=> a -> Maybe (Effect n)
-> Exp a n -> Exp a n
weakeff a meff x
= maybe x (\e -> XCast a (CastWeakenEffect e) x) meff
weakclo :: Ord n
=> a -> [Exp a n]
-> Exp a n -> Exp a n
weakclo a clos x
= case clos of
[] -> x
_ -> XCast a (CastWeakenClosure clos) x
wrapLets
:: Ord n
=> a
-> [(BindMode, Bind n)]
-> [(Bind n, Exp a n)]
-> ( [(Bind n, Exp a n)]
, [Lets a n])
wrapLets a binds bas
= let isMkLet (_, (BMValue i, _)) = i /= 1
isMkLet _ = False
(as, bs'') = partition isMkLet (bas `zip` binds)
as' = map fst as
bs' = map fst bs''
anons = zipWith (\(b,_) i -> (b, XVar a (UIx i))) as' [0..]
values = map (\(b,v) -> (BAnon (substT bs' $ T.typeOfBind b), v))
(reverse as')
lets = map (\(b,v) -> LLet b v) values
in (bs' ++ anons, lets)
substT :: Ord n => [(Bind n, Exp a n)] -> Type n -> Type n
substT bas x
= let sub = [(b, t) | (b, XType t) <- bas ]
in S.substituteTs sub x
matchWithRule
:: (Show n, Show a, Ord n, Pretty n)
=> RewriteRule a n
-> RE.RewriteEnv a n
-> Exp a n
-> [(Exp a n, a)]
-> RM.SubstInfo a n
-> Maybe ( RM.SubstInfo a n
, [(Exp a n, a)])
matchWithRule
RewriteRule
{ ruleBinds = binds
, ruleLeft = lhs
, ruleLeftHole = Nothing
, ruleFreeVars = free }
env f args sub
= do
when (any (flip RE.hasDef env) free)
$ Nothing
let Just vs
= liftM Set.fromList
$ sequence
$ map T.takeNameOfBind
$ map snd binds
l:ls <- return $ X.takeXAppsAsList lhs
sub' <- RM.match sub vs l f
let go m [] rest
= do return $ (m, rest)
go m (l':ls') ((r,_):rs)
= do m' <- RM.match m vs l' r
go m' ls' rs
go _ _ _
= Nothing
go sub' ls args
matchWithRule
rule@(RewriteRule { ruleLeftHole = Just hole })
env f args sub
| Just a <- X.takeAnnotOfExp f
, lhs_full <- XApp a (ruleLeft rule) hole
, rule_full <- rule { ruleLeft = lhs_full, ruleLeftHole = Nothing}
, Just subst <- matchWithRule rule_full env f args sub
= Just subst
| rule_some <- rule { ruleLeftHole = Nothing }
, Just (sub', (XVar _ b, _) : as)
<- matchWithRule rule_some env f args sub
, Just d' <- RE.getDef b env
, (fd, ad) <- X.takeXAppsWithAnnots d'
, rule_hole <- rule { ruleLeft = hole, ruleLeftHole = Nothing }
, Just (subd, asd) <- matchWithRule rule_hole env fd ad sub'
= Just (subd, asd ++ as)
| otherwise
= Nothing
lookupFromSubst :: Ord n
=> [(BindMode,Bind n)]
-> (Map n (Exp a n), Map n (Type n))
-> [(Bind n, Exp a n)]
lookupFromSubst bs m
= let bas = catMaybes $ map (lookupX m) bs
in map (\(b, a) -> (A.anonymizeX b, A.anonymizeX a)) bas
where lookupX (xs,_) (BMValue _, b@(BName n _))
| Just x <- Map.lookup n xs
= Just (b, x)
lookupX (_,tys) (BMSpec, b@(BName n _))
| Just t <- Map.lookup n tys
= Just (b, XType t)
lookupX _ _ = Nothing