module Language.HERMIT.Primitive.Unfold
( externals
, stashDef
, stashApply
, getUnfolding
) where
import GhcPlugins hiding (empty)
import Control.Monad
import Control.Applicative
import Language.HERMIT.Primitive.GHC hiding (externals)
import Language.HERMIT.Primitive.Common
import Language.HERMIT.CoreExtra
import Language.HERMIT.Kure
import Language.HERMIT.Monad
import Language.HERMIT.External
import Language.HERMIT.Context
import Language.HERMIT.GHC
import Prelude hiding (exp)
externals :: [External]
externals =
[ external "remember" stashDef
["Remember the current binding, allowing it to be folded/unfolded in the future."]
, external "unfold" (promoteExprR . stashApply)
["Unfold a remembered definition."]
]
stashDef :: String -> RewriteH Core
stashDef label = sideEffectR $ \ _ core ->
case core of
DefCore def -> saveDef label def
BindCore (NonRec i e) -> saveDef label (Def i e)
_ -> fail "stashDef: not a binding"
stashApply :: String -> RewriteH CoreExpr
stashApply label = setFailMsg "Inlining stashed definition failed: " $
withPatFailMsg (wrongExprForm "Var v") $
do (c, Var v) <- exposeT
constT $ do Def i rhs <- lookupDef label
if idName i == idName v
then ifM (all (inScope c) <$> apply freeVarsT c rhs)
(return rhs)
(fail "some free variables in stashed definition are no longer in scope.")
else fail $ "stashed definition applies to " ++ var2String i ++ " not " ++ var2String v
getUnfolding :: Monad m
=> Bool
-> Bool
-> Id -> Context -> m (CoreExpr, Int)
getUnfolding scrutinee caseBinderOnly i c =
case lookupHermitBinding i c of
Nothing -> case unfoldingInfo (idInfo i) of
CoreUnfolding { uf_tmpl = uft } -> if caseBinderOnly then fail "not a case binder" else return (uft, 0)
_ -> fail $ "cannot find unfolding in Env or IdInfo."
Just (LAM {}) -> fail $ "variable is lambda-bound."
Just (BIND depth _ e') -> if caseBinderOnly then fail "not a case binder." else return (e', depth)
Just (CASE depth s coreAlt) -> return $ if scrutinee
then (s, depth)
else let tys = tyConAppArgs (idType i)
in either (,depth) (,depth+1) (alt2Exp s tys coreAlt)
alt2Exp :: CoreExpr -> [Type] -> (AltCon,[Id]) -> Either CoreExpr CoreExpr
alt2Exp d _ (DEFAULT , _ ) = Left d
alt2Exp _ _ (LitAlt l , _ ) = Right $ Lit l
alt2Exp _ tys (DataAlt dc, as) = Right $ mkCoreConApps dc (map Type tys ++ map Var as)