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 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 " ++ showPpr i ++ " not " ++ showPpr 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 " ++ show i ++ " in Env or IdInfo."
Just (LAM {}) -> fail $ show i ++ " 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)