module Language.HERMIT.Primitive.Fold
( externals
, foldR
, stashFoldR
) where
import GhcPlugins hiding (empty)
import Control.Applicative
import Control.Monad
import qualified Data.Map as Map
import Language.HERMIT.Monad
import Language.HERMIT.Context
import Language.HERMIT.External
import Language.HERMIT.Kure
import Language.HERMIT.GHC
import Language.HERMIT.Primitive.GHC hiding (externals)
import Language.HERMIT.Primitive.Unfold hiding (externals)
import qualified Language.Haskell.TH as TH
import Prelude hiding (exp)
externals :: [External]
externals =
[ external "fold" (promoteExprR . foldR)
["fold a definition"
,""
,"double :: Int -> Int"
,"double x = x + x"
,""
,"5 + 5 + 6"
,"any-bu (fold 'double)"
,"double 5 + 6"
,""
,"Note: due to associativity, if you wanted to fold 5 + 6 + 6, "
,"you first need to apply an associativity rewrite."
] .+ Context .+ Deep
, external "fold" (promoteExprR . stashFoldR)
["Fold a remembered definition."] .+ Context .+ Deep
]
stashFoldR :: String -> RewriteH CoreExpr
stashFoldR label = prefixFailMsg "Fold failed: " $
contextfreeT $ \ e -> do
Def i rhs <- lookupDef label
maybe (fail "no match.")
return
(fold i rhs e)
foldR :: TH.Name -> RewriteH CoreExpr
foldR nm = prefixFailMsg "Fold failed: " $
translate $ \ c e -> do
i <- case filter (\i -> nm `cmpTHName2Id` i) $ Map.keys (hermitBindings c) of
[i] -> return i
_ -> fail "cannot find name."
either fail
(\(rhs,_d) -> maybe (fail "no match.")
return
(fold i rhs e))
(getUnfolding False False i c)
fold :: Id -> CoreExpr -> CoreExpr -> Maybe CoreExpr
fold i lam exp = do
let (vs,body) = foldArgs lam
checkEqual :: Maybe CoreExpr -> Maybe CoreExpr -> Maybe CoreExpr
checkEqual m1 m2 = ifM (exprEqual <$> m1 <*> m2) m1 Nothing
al <- foldMatch vs body exp
let m = Map.fromListWith checkEqual [(k,Just v) | (k,v) <- al ]
es <- sequence [ join (Map.lookup v m) | v <- vs ]
return $ mkCoreApps (Var i) es
foldArgs :: CoreExpr -> ([Var], CoreExpr)
foldArgs = go []
where go vs (Lam v e) = go (v:vs) e
go vs e = (reverse vs, e)
foldMatch :: [Var]
-> CoreExpr
-> CoreExpr
-> Maybe [(Var,CoreExpr)]
foldMatch vs (Var i) e | i `elem` vs = return [(i,e)]
| otherwise = case e of
Var i' | i == i' -> return []
_ -> Nothing
foldMatch _ (Lit l) (Lit l') | l == l' = return []
foldMatch vs (App e a) (App e' a') = do
x <- foldMatch vs e e'
y <- foldMatch vs a a'
return (x ++ y)
foldMatch vs (Lam v e) (Lam v' e') | v == v' = foldMatch (filter (==v) vs) e e'
foldMatch vs (Let (NonRec v rhs) e) (Let (NonRec v' rhs') e') | v == v' = do
x <- foldMatch vs rhs rhs'
y <- foldMatch (filter (==v) vs) e e'
return (x ++ y)
foldMatch vs (Let (Rec bnds) e) (Let (Rec bnds') e') | length bnds == length bnds' = do
let vs' = filter (`elem` map fst bnds) vs
bmatch (v,rhs) (v',rhs') | v == v' = foldMatch vs' rhs rhs'
bmatch _ _ = Nothing
x <- zipWithM bmatch bnds bnds'
y <- foldMatch vs' e e'
return (concat x ++ y)
foldMatch vs (Tick t e) (Tick t' e') | t == t' = foldMatch vs e e'
foldMatch vs (Case s b ty alts) (Case s' b' ty' alts')
| (b == b') && (showPpr ty == showPpr ty') && (length alts == length alts') = do
x <- foldMatch vs s s'
let vs' = filter (==b) vs
altMatch (ac, is, e) (ac', is', e') | (ac == ac') && (is == is') =
foldMatch (filter (`elem` is) vs') e e'
altMatch _ _ = Nothing
y <- zipWithM altMatch alts alts'
return (x ++ concat y)
foldMatch vs (Cast e c) (Cast e' c') | showPpr c == showPpr c' = foldMatch vs e e'
foldMatch _ (Type t) (Type t') | showPpr t == showPpr t' = return []
foldMatch _ (Coercion c) (Coercion c') | showPpr c == showPpr c' = return []
foldMatch _ _ _ = Nothing