module HERMIT.Dictionary.Unfold
( externals
, betaReducePlusR
, rememberR
, showStashT
, unfoldR
, unfoldPredR
, unfoldNameR
, unfoldNamesR
, unfoldSaturatedR
, unfoldStashR
, specializeR
) where
import Control.Arrow
import Control.Monad
import Data.List (intercalate)
import qualified Data.Map as Map
import HERMIT.PrettyPrinter.Common (DocH, PrettyH, TransformDocH(..), PrettyC)
import HERMIT.Dictionary.Common
import HERMIT.Dictionary.GHC (substCoreExpr)
import HERMIT.Dictionary.Inline (inlineR)
import HERMIT.Core
import HERMIT.Context
import HERMIT.External
import HERMIT.GHC
import HERMIT.Kure
import HERMIT.Monad
import HERMIT.Name
import Prelude hiding (exp)
import qualified Text.PrettyPrint.MarkedHughesPJ as PP
externals :: [External]
externals =
[ external "beta-reduce-plus" (promoteExprR betaReducePlusR :: RewriteH Core)
[ "Perform one or more beta-reductions."] .+ Eval .+ Shallow
, external "remember" (rememberR :: RememberedName -> RewriteH Core)
[ "Remember the current binding, allowing it to be folded/unfolded in the future." ] .+ Context
, external "unfold-remembered" (promoteExprR . unfoldStashR :: RememberedName -> RewriteH Core)
[ "Unfold a remembered definition." ] .+ Deep .+ Context
, external "unfold" (promoteExprR unfoldR :: RewriteH Core)
[ "In application f x y z, unfold f." ] .+ Deep .+ Context
, external "unfold" (promoteExprR . unfoldNameR . unOccurrenceName :: OccurrenceName -> RewriteH Core)
[ "Inline a definition, and apply the arguments; traditional unfold." ] .+ Deep .+ Context
, external "unfold" (promoteExprR . unfoldNamesR . map unOccurrenceName:: [OccurrenceName] -> RewriteH Core)
[ "Unfold a definition if it is named in the list." ] .+ Deep .+ Context
, external "unfold-saturated" (promoteExprR unfoldSaturatedR :: RewriteH Core)
[ "Unfold a definition only if the function is fully applied." ] .+ Deep .+ Context
, external "specialize" (promoteExprR specializeR :: RewriteH Core)
[ "Specialize an application to its type and coercion arguments." ] .+ Deep .+ Context
, external "show-remembered" (TransformDocH showStashT :: TransformDocH CoreTC)
[ "Display all remembered definitions." ]
]
betaReducePlusR :: MonadCatch m => Rewrite c m CoreExpr
betaReducePlusR = prefixFailMsg "Multi-beta-reduction failed: " $ do
(f,args) <- callT
let (f',args',atLeastOne) = reduceAll f args False
reduceAll (Lam v body) (a:as) _ = reduceAll (substCoreExpr v a body) as True
reduceAll e as b = (e,as,b)
guardMsg atLeastOne "no beta reductions possible."
return $ mkCoreApps f' args'
unfoldR :: forall c m. ( AddBindings c, ExtendPath c Crumb, HasEmptyContext c
, ReadBindings c, ReadPath c Crumb, MonadCatch m )
=> Rewrite c m CoreExpr
unfoldR = prefixFailMsg "unfold failed: " (go >>> tryR betaReducePlusR)
where go :: Rewrite c m CoreExpr
go = appAllR go idR <+ inlineR
unfoldPredR :: ( AddBindings c, ExtendPath c Crumb, HasEmptyContext c, ReadBindings c, ReadPath c Crumb
, MonadCatch m )
=> (Id -> [CoreExpr] -> Bool) -> Rewrite c m CoreExpr
unfoldPredR p = callPredT p >> unfoldR
unfoldNameR :: ( ExtendPath c Crumb, ReadPath c Crumb, AddBindings c, ReadBindings c, HasEmptyContext c
, MonadCatch m )
=> HermitName -> Rewrite c m CoreExpr
unfoldNameR nm = prefixFailMsg ("unfold '" ++ show nm ++ " failed: ") (callNameT nm >> unfoldR)
unfoldNamesR :: ( ExtendPath c Crumb, ReadPath c Crumb, AddBindings c, ReadBindings c, HasEmptyContext c
, MonadCatch m )
=> [HermitName] -> Rewrite c m CoreExpr
unfoldNamesR [] = fail "unfold-names failed: no names given."
unfoldNamesR nms = setFailMsg "unfold-names failed." $ orR (map unfoldNameR nms)
unfoldSaturatedR :: (ExtendPath c Crumb, ReadPath c Crumb, AddBindings c, ReadBindings c, HasEmptyContext c) => Rewrite c HermitM CoreExpr
unfoldSaturatedR = callSaturatedT >> unfoldR
specializeR :: (ExtendPath c Crumb, ReadPath c Crumb, AddBindings c, ReadBindings c, HasEmptyContext c) => Rewrite c HermitM CoreExpr
specializeR = unfoldPredR (const $ all isTyCoArg)
rememberR :: RememberedName -> Rewrite c HermitM Core
rememberR label = sideEffectR $ \ _ -> \case
DefCore def -> saveDef label def
BindCore (NonRec i e) -> saveDef label (Def i e)
_ -> fail "remember failed: not applied to a binding."
unfoldStashR :: ReadBindings c => RememberedName -> Rewrite c HermitM CoreExpr
unfoldStashR label = prefixFailMsg "Inlining stashed definition failed: " $
withPatFailMsg (wrongExprForm "Var v") $
do (c, Var v) <- exposeT
constT $ do Def i rhs <- lookupDef label
dflags <- getDynFlags
if idName i == idName v
then let fvars = varSetElems $ localFreeVarsExpr rhs
in if all (inScope c) fvars
then return rhs
else fail $ "free variables " ++ intercalate "," (map (showPpr dflags) (filter (not . inScope c) fvars)) ++ " in stashed definition are no longer in scope."
else fail $ "stashed definition applies to " ++ unqualifiedName i ++ " not " ++ unqualifiedName v
showStashT :: Injection CoreDef a => PrettyC -> PrettyH a -> Transform c HermitM a DocH
showStashT pctx pp = do
stash <- constT getStash
docs <- forM (Map.toList stash) $ \ (l,d) -> do
dfn <- constT $ applyT (extractT pp) pctx d
return $ PP.text ("[ " ++ show l ++ " ]") PP.$+$ dfn PP.$+$ PP.space
return $ PP.vcat docs