module CLaSH.Rewrite.Util where
import Control.DeepSeq
import Control.Exception (throw)
import Control.Lens (Lens', (%=), (+=), (^.),_1,_3)
import qualified Control.Lens as Lens
import qualified Control.Monad as Monad
import qualified Control.Monad.State.Strict as State
import qualified Control.Monad.Writer as Writer
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Lazy as HML
import qualified Data.HashMap.Strict as HMS
import qualified Data.List as List
import qualified Data.Map as Map
import Data.Maybe (catMaybes,isJust,mapMaybe)
import qualified Data.Monoid as Monoid
import qualified Data.Set as Set
import qualified Data.Set.Lens as Lens
import Unbound.Generics.LocallyNameless (Fresh, bind,
embed, makeName, name2String,
rebind, rec, string2Name, unbind,
unembed, unrec)
import Unbound.Generics.LocallyNameless.Unsafe (unsafeUnbind)
import SrcLoc (SrcSpan)
import CLaSH.Core.DataCon (dataConInstArgTys)
import CLaSH.Core.FreeVars (termFreeIds, termFreeTyVars,
typeFreeVars)
import CLaSH.Core.Pretty (showDoc)
import CLaSH.Core.Subst (substTm)
import CLaSH.Core.Term (LetBinding, Pat (..), Term (..),
TmName)
import CLaSH.Core.TyCon (TyCon, TyConName, tyConDataCons)
import CLaSH.Core.Type (KindOrType, Type (..),
TypeView (..), coreView,
typeKind, tyView)
import CLaSH.Core.Util (Delta, Gamma, collectArgs,
mkAbstraction, mkApps, mkId,
mkLams, mkTmApps, mkTyApps,
mkTyLams, mkTyVar, termType)
import CLaSH.Core.Var (Id, TyVar, Var (..))
import CLaSH.Driver.Types (CLaSHException (..))
import CLaSH.Netlist.Util (representableType)
import CLaSH.Rewrite.Types
import CLaSH.Util
zoomExtra :: State.State extra a
-> RewriteMonad extra a
zoomExtra m = R (\_ s -> case State.runState m (s ^. extra) of
(a,s') -> (a,s {_extra = s'},mempty))
apply :: String
-> Rewrite extra
-> Rewrite extra
apply name rewrite ctx expr = do
lvl <- Lens.view dbgLevel
let before = showDoc expr
(expr', anyChanged) <- traceIf (lvl >= DebugAll) ("Trying: " ++ name ++ " on:\n" ++ before) $ Writer.listen $ rewrite ctx expr
let hasChanged = Monoid.getAny anyChanged
Monad.when hasChanged $ transformCounter += 1
let after = showDoc expr'
let expr'' = if hasChanged then expr' else expr
Monad.when (lvl > DebugNone && hasChanged) $ do
tcm <- Lens.view tcCache
beforeTy <- termType tcm expr
let beforeFTV = Lens.setOf termFreeTyVars expr
beforeFV <- Lens.setOf <$> localFreeIds <*> pure expr
afterTy <- termType tcm expr'
let afterFTV = Lens.setOf termFreeTyVars expr
afterFV <- Lens.setOf <$> localFreeIds <*> pure expr'
let newFV = Set.size afterFTV > Set.size beforeFTV ||
Set.size afterFV > Set.size beforeFV
Monad.when newFV $
error ( concat [ $(curLoc)
, "Error when applying rewrite ", name
, " to:\n" , before
, "\nResult:\n" ++ after ++ "\n"
, "Changes free variables from: ", show (beforeFTV,beforeFV)
, "\nto: ", show (afterFTV,afterFV)
]
)
traceIf ( beforeTy /= afterTy)
( concat [ $(curLoc)
, "Error when applying rewrite ", name
, " to:\n" , before
, "\nResult:\n" ++ after ++ "\n"
, "Changes type from:\n", showDoc beforeTy
, "\nto:\n", showDoc afterTy
]
) (return ())
Monad.when (lvl >= DebugApplied && not hasChanged && expr /= expr') $
error $ $(curLoc) ++ "Expression changed without notice(" ++ name ++ "): before" ++ before ++ "\nafter:\n" ++ after
traceIf (lvl >= DebugName && hasChanged) name $
traceIf (lvl >= DebugApplied && hasChanged) ("Changes when applying rewrite to:\n" ++ before ++ "\nResult:\n" ++ after ++ "\n") $
traceIf (lvl >= DebugAll && not hasChanged) ("No changes when applying rewrite " ++ name ++ " to:\n" ++ after ++ "\n") $
return expr''
runRewrite :: String
-> Rewrite extra
-> Term
-> RewriteMonad extra Term
runRewrite name rewrite expr = apply name rewrite [] expr
runRewriteSession :: RewriteEnv
-> RewriteState extra
-> RewriteMonad extra a
-> a
runRewriteSession r s m = traceIf True ("Applied " ++
show (s' ^. transformCounter) ++
" transformations")
a
where
(a,s',_) = runR m r s
setChanged :: RewriteMonad extra ()
setChanged = Writer.tell (Monoid.Any True)
changed :: a -> RewriteMonad extra a
changed val = do
Writer.tell (Monoid.Any True)
return val
contextEnv :: [CoreContext]
-> (Gamma, Delta)
contextEnv = go HML.empty HML.empty
where
go gamma delta [] = (gamma,delta)
go gamma delta (LetBinding _ ids:ctx) = go gamma' delta ctx
where
gamma' = foldl addToGamma gamma ids
go gamma delta (LetBody ids:ctx) = go gamma' delta ctx
where
gamma' = foldl addToGamma gamma ids
go gamma delta (LamBody lId:ctx) = go gamma' delta ctx
where
gamma' = addToGamma gamma lId
go gamma delta (TyLamBody tv:ctx) = go gamma delta' ctx
where
delta' = addToDelta delta tv
go gamma delta (CaseAlt ids:ctx) = go gamma' delta ctx
where
gamma' = foldl addToGamma gamma ids
go gamma delta (_:ctx) = go gamma delta ctx
addToGamma gamma (Id idName ty) = HML.insert idName (unembed ty) gamma
addToGamma _ _ = error $ $(curLoc) ++ "Adding TyVar to Gamma"
addToDelta delta (TyVar tvName ki) = HML.insert tvName (unembed ki) delta
addToDelta _ _ = error $ $(curLoc) ++ "Adding Id to Delta"
closestLetBinder :: [CoreContext] -> Maybe Id
closestLetBinder [] = Nothing
closestLetBinder (LetBinding id_ _:_) = Just id_
closestLetBinder (_:ctx) = closestLetBinder ctx
mkDerivedName :: [CoreContext] -> String -> String
mkDerivedName ctx sf = case closestLetBinder ctx of
Just id_ -> ((++ ('_':sf)) . name2String . varName) id_
_ -> sf
mkEnv :: [CoreContext]
-> RewriteMonad extra (Gamma, Delta)
mkEnv ctx = do
let (gamma,delta) = contextEnv ctx
tsMap <- fmap (HML.map (^. _1)) $ Lens.use bindings
let gamma' = tsMap `HML.union` gamma
return (gamma',delta)
mkTmBinderFor :: (Functor m, Fresh m, MonadUnique m)
=> HashMap TyConName TyCon
-> String
-> Term
-> m (Id, Term)
mkTmBinderFor tcm name e = do
(Left r) <- mkBinderFor tcm name (Left e)
return r
mkBinderFor :: (Functor m, Monad m, MonadUnique m, Fresh m)
=> HashMap TyConName TyCon
-> String
-> Either Term Type
-> m (Either (Id,Term) (TyVar,Type))
mkBinderFor tcm name (Left term) =
Left <$> (mkInternalVar name =<< termType tcm term)
mkBinderFor tcm name (Right ty) = do
name' <- fmap (makeName name . toInteger) getUniqueM
let kind = typeKind tcm ty
return $ Right (TyVar name' (embed kind), VarTy kind name')
mkInternalVar :: (Functor m, Monad m, MonadUnique m)
=> String
-> KindOrType
-> m (Id,Term)
mkInternalVar name ty = do
name' <- fmap (makeName name . toInteger) getUniqueM
return (Id name' (embed ty),Var ty name')
inlineBinders :: (Term -> LetBinding -> RewriteMonad extra Bool)
-> Rewrite extra
inlineBinders condition _ expr@(Letrec b) = do
(xes,res) <- unbind b
let expr' = Letrec (bind xes res)
(replace,others) <- partitionM (condition expr') (unrec xes)
case replace of
[] -> return expr
_ -> do
let (others',res') = substituteBinders replace others res
newExpr = case others' of
[] -> res'
_ -> Letrec (bind (rec others') res')
changed newExpr
inlineBinders _ _ e = return e
isJoinPointIn :: Id
-> Term
-> Bool
isJoinPointIn id_ e = case tailCalls id_ e of
Just n | n > 1 -> True
_ -> False
tailCalls :: Id
-> Term
-> Maybe Int
tailCalls id_ expr = case expr of
Var _ nm | varName id_ == nm -> Just 1
| otherwise -> Just 0
Lam b -> let (_,expr') = unsafeUnbind b
in tailCalls id_ expr'
TyLam b -> let (_,expr') = unsafeUnbind b
in tailCalls id_ expr'
App l r -> case tailCalls id_ r of
Just 0 -> tailCalls id_ l
_ -> Nothing
TyApp l _ -> tailCalls id_ l
Letrec b ->
let (bsR,expr') = unsafeUnbind b
(bsIds,bsExprs) = unzip (unrec bsR)
bsTls = map (tailCalls id_ . unembed) bsExprs
bsIdsUsed = mapMaybe (\(l,r) -> pure l <* r) (zip bsIds bsTls)
bsIdsTls = map (`tailCalls` expr') bsIdsUsed
bsCount = pure . sum $ catMaybes bsTls
in case (all isJust bsTls) of
False -> Nothing
True -> case (all (==0) $ catMaybes bsTls) of
False -> case all isJust bsIdsTls of
False -> Nothing
True -> (+) <$> bsCount <*> tailCalls id_ expr'
True -> tailCalls id_ expr'
Case scrut _ alts ->
let scrutTl = tailCalls id_ scrut
altsTl = map (tailCalls id_ . snd . unsafeUnbind) alts
in case scrutTl of
Just 0 | all (/= Nothing) altsTl -> Just (sum (catMaybes altsTl))
_ -> Nothing
_ -> Just 0
substituteBinders :: [LetBinding]
-> [LetBinding]
-> Term
-> ([LetBinding],Term)
substituteBinders [] others res = (others,res)
substituteBinders ((bndr,valE):rest) others res = substituteBinders rest' others' res'
where
val = unembed valE
bndrName = varName bndr
selfRef = bndrName `elem` Lens.toListOf termFreeIds val
(res',rest',others') = if selfRef
then (res,rest,(bndr,valE):others)
else ( substTm (varName bndr) val res
, map (second ( embed
. substTm bndrName val
. unembed)
) rest
, map (second ( embed
. substTm bndrName val
. unembed)
) others
)
localFreeIds :: (Applicative f, Lens.Contravariant f)
=> RewriteMonad extra ((TmName -> f TmName) -> Term -> f Term)
localFreeIds = do
globalBndrs <- Lens.use bindings
return ((termFreeIds . Lens.filtered (not . (`HML.member` globalBndrs))))
inlineOrLiftBinders :: (LetBinding -> RewriteMonad extra Bool)
-> (Term -> LetBinding -> RewriteMonad extra Bool)
-> Rewrite extra
inlineOrLiftBinders condition inlineOrLift ctx expr@(Letrec b) = do
(xesR,res) <- unbind b
let xes = unrec xesR
(replace,others) <- partitionM condition xes
case replace of
[] -> return expr
_ -> do
let expr' = Letrec (bind xesR res)
(doInline,doLift) <- partitionM (inlineOrLift expr') replace
let (others',res') = substituteBinders doInline (doLift ++ others) res
(doLift',others'') = splitAt (length doLift) others'
(gamma,delta) <- mkEnv (LetBinding undefined (map fst xes) : ctx)
doLift'' <- mapM (liftBinding gamma delta) doLift'
let (others3,res'') = substituteBinders doLift'' others'' res'
newExpr = case others3 of
[] -> res''
_ -> Letrec (bind (rec others3) res'')
changed newExpr
inlineOrLiftBinders _ _ _ e = return e
liftBinding :: Gamma
-> Delta
-> LetBinding
-> RewriteMonad extra LetBinding
liftBinding gamma delta (Id idName tyE,eE) = do
let e = unembed eE
let localFTVs = List.nub $ Lens.toListOf termFreeTyVars e
localFVs <- List.nub <$> (Lens.toListOf <$> localFreeIds <*> pure e)
let localFTVkinds = map (\k -> HML.lookupDefault (error $ $(curLoc) ++ show k ++ " not found") k delta) localFTVs
localFVs' = filter (/= idName) localFVs
localFVtys' = map (\k -> HML.lookupDefault (error $ $(curLoc) ++ show k ++ " not found") k gamma) localFVs'
boundFTVs = zipWith mkTyVar localFTVkinds localFTVs
boundFVs = zipWith mkId localFVtys' localFVs'
tcm <- Lens.view tcCache
newBodyTy <- termType tcm $ mkTyLams (mkLams e boundFVs) boundFTVs
(cf,sp) <- Lens.use curFun
newBodyId <- fmap (makeName (name2String cf ++ "_" ++ name2String idName) . toInteger) getUniqueM
let newExpr = mkTmApps
(mkTyApps (Var newBodyTy newBodyId)
(zipWith VarTy localFTVkinds localFTVs))
(zipWith Var localFVtys' localFVs')
e' = substTm idName newExpr e
newBody = mkTyLams (mkLams e' boundFVs) boundFTVs
aeqExisting <- (HMS.toList . HMS.filter ((== newBody) . (^. _3))) <$> Lens.use bindings
case aeqExisting of
[] -> do
bindings %= HMS.insert newBodyId (newBodyTy,sp,newBody)
return (Id idName tyE, embed newExpr)
((k,(aeqTy,_,_)):_) ->
let newExpr' = mkTmApps
(mkTyApps (Var aeqTy k)
(zipWith VarTy localFTVkinds localFTVs))
(zipWith Var localFVtys' localFVs')
in return (Id idName tyE, embed newExpr')
liftBinding _ _ _ = error $ $(curLoc) ++ "liftBinding: invalid core, expr bound to tyvar"
mkFunction :: TmName
-> SrcSpan
-> Term
-> RewriteMonad extra (TmName,Type)
mkFunction bndr sp body = do
tcm <- Lens.view tcCache
bodyTy <- termType tcm body
bodyId <- cloneVar bndr
addGlobalBind bodyId bodyTy sp body
return (bodyId,bodyTy)
addGlobalBind :: TmName
-> Type
-> SrcSpan
-> Term
-> RewriteMonad extra ()
addGlobalBind vId ty sp body = (ty,body) `deepseq` bindings %= HMS.insert vId (ty,sp,body)
cloneVar :: TmName
-> RewriteMonad extra TmName
cloneVar name = fmap (makeName (name2String name) . toInteger) getUniqueM
isLocalVar :: Term
-> RewriteMonad extra Bool
isLocalVar (Var _ name)
= fmap (not . HML.member name)
$ Lens.use bindings
isLocalVar _ = return False
isUntranslatable :: Term
-> RewriteMonad extra Bool
isUntranslatable tm = do
tcm <- Lens.view tcCache
not <$> (representableType <$> Lens.view typeTranslator
<*> pure tcm
<*> termType tcm tm)
isUntranslatableType :: Type
-> RewriteMonad extra Bool
isUntranslatableType ty =
not <$> (representableType <$> Lens.view typeTranslator
<*> Lens.view tcCache
<*> pure ty)
isLambdaBodyCtx :: CoreContext
-> Bool
isLambdaBodyCtx (LamBody _) = True
isLambdaBodyCtx _ = False
mkWildValBinder :: (Functor m, Monad m, MonadUnique m)
=> Type
-> m Id
mkWildValBinder = fmap fst . mkInternalVar "wild"
mkSelectorCase :: (Functor m, Monad m, MonadUnique m, Fresh m)
=> String
-> HashMap TyConName TyCon
-> Term
-> Int
-> Int
-> m Term
mkSelectorCase caller tcm scrut dcI fieldI = do
scrutTy <- termType tcm scrut
go scrutTy
where
go (coreView tcm -> Just ty') = go ty'
go scrutTy@(tyView -> TyConApp tc args) =
case tyConDataCons (tcm HMS.! tc) of
[] -> cantCreate $(curLoc) ("TyCon has no DataCons: " ++ show tc ++ " " ++ showDoc tc) scrutTy
dcs | dcI > length dcs -> cantCreate $(curLoc) "DC index exceeds max" scrutTy
| otherwise -> do
let dc = indexNote ($(curLoc) ++ "No DC with tag: " ++ show (dcI1)) dcs (dcI1)
let (Just fieldTys) = dataConInstArgTys dc args
if fieldI >= length fieldTys
then cantCreate $(curLoc) "Field index exceed max" scrutTy
else do
wildBndrs <- mapM mkWildValBinder fieldTys
let ty = indexNote ($(curLoc) ++ "No DC field#: " ++ show fieldI) fieldTys fieldI
selBndr <- mkInternalVar "sel" ty
let bndrs = take fieldI wildBndrs ++ [fst selBndr] ++ drop (fieldI+1) wildBndrs
pat = DataPat (embed dc) (rebind [] bndrs)
retVal = Case scrut ty [ bind pat (snd selBndr) ]
return retVal
go scrutTy = cantCreate $(curLoc) ("Type of subject is not a datatype: " ++ showDoc scrutTy) scrutTy
cantCreate loc info scrutTy = error $ loc ++ "Can't create selector " ++ show (caller,dcI,fieldI) ++ " for: (" ++ showDoc scrut ++ " :: " ++ showDoc scrutTy ++ ")\nAdditional info: " ++ info
specialise :: Lens' extra (Map.Map (TmName, Int, Either Term Type) (TmName,Type))
-> Lens' extra (HashMap TmName Int)
-> Lens' extra Int
-> Rewrite extra
specialise specMapLbl specHistLbl specLimitLbl ctx e = case e of
(TyApp e1 ty) -> specialise' specMapLbl specHistLbl specLimitLbl ctx e (collectArgs e1) (Right ty)
(App e1 e2) -> specialise' specMapLbl specHistLbl specLimitLbl ctx e (collectArgs e1) (Left e2)
_ -> return e
specialise' :: Lens' extra (Map.Map (TmName, Int, Either Term Type) (TmName,Type))
-> Lens' extra (HashMap TmName Int)
-> Lens' extra Int
-> [CoreContext]
-> Term
-> (Term, [Either Term Type])
-> Either Term Type
-> RewriteMonad extra Term
specialise' specMapLbl specHistLbl specLimitLbl ctx e (Var _ f, args) specArg = do
lvl <- Lens.view dbgLevel
(specBndrs,specVars) <- specArgBndrsAndVars ctx specArg
let argLen = length args
specAbs = either (Left . (`mkAbstraction` specBndrs)) (Right . id) specArg
specM <- Map.lookup (f,argLen,specAbs) <$> Lens.use (extra.specMapLbl)
case specM of
Just (fname,fty) ->
traceIf (lvl >= DebugApplied) ("Using previous specialization of " ++ showDoc f ++ " on " ++ (either showDoc showDoc) specAbs ++ ": " ++ showDoc fname) $
changed $ mkApps (Var fty fname) (args ++ specVars)
Nothing -> do
bodyMaybe <- fmap (HML.lookup f) $ Lens.use bindings
case bodyMaybe of
Just (_,sp,bodyTm) -> do
specHistM <- HML.lookup f <$> Lens.use (extra.specHistLbl)
specLim <- Lens.use (extra . specLimitLbl)
if maybe False (> specLim) specHistM
then throw (CLaSHException
sp
(unlines [ "Hit specialisation limit " ++ show specLim ++ " on function `" ++ showDoc f ++ "'.\n"
, "The function `" ++ showDoc f ++ "' is most likely recursive, and looks like it is being indefinitely specialized on a growing argument.\n"
, "Body of `" ++ showDoc f ++ "':\n" ++ showDoc bodyTm ++ "\n"
, "Argument (in position: " ++ show argLen ++ ") that triggered termination:\n" ++ (either showDoc showDoc) specArg
, "Run with '-clash-spec-limit=N' to increase the specialisation limit to N."
])
Nothing)
else do
tcm <- Lens.view tcCache
(boundArgs,argVars) <- fmap (unzip . map (either (Left *** Left) (Right *** Right))) $
mapM (mkBinderFor tcm "pTS") args
let newBody = mkAbstraction (mkApps bodyTm (argVars ++ [specArg])) (boundArgs ++ specBndrs)
newf <- mkFunction f sp newBody
(extra.specHistLbl) %= HML.insertWith (+) f 1
(extra.specMapLbl) %= Map.insert (f,argLen,specAbs) newf
let newExpr = mkApps ((uncurry . flip) Var newf) (args ++ specVars)
newf `deepseq` changed newExpr
Nothing -> return e
specialise' _ _ _ ctx _ (appE,args) (Left specArg) = do
(specBndrs,specVars) <- specArgBndrsAndVars ctx (Left specArg)
let newBody = mkAbstraction specArg specBndrs
existing <- HML.filter ((== newBody) . (^. _3)) <$> Lens.use bindings
newf <- case HML.toList existing of
[] -> do (cf,sp) <- Lens.use curFun
mkFunction (string2Name (name2String cf ++ "_" ++ "specF")) sp newBody
((k,(kTy,_,_)):_) -> return (k,kTy)
let newArg = Left $ mkApps ((uncurry . flip) Var newf) specVars
let newExpr = mkApps appE (args ++ [newArg])
changed newExpr
specialise' _ _ _ _ e _ _ = return e
specArgBndrsAndVars :: [CoreContext]
-> Either Term Type
-> RewriteMonad extra ([Either Id TyVar],[Either Term Type])
specArgBndrsAndVars ctx specArg = do
let specFTVs = List.nub $ either (Lens.toListOf termFreeTyVars) (Lens.toListOf typeFreeVars) specArg
specFVs <- List.nub <$> either ((Lens.toListOf <$> localFreeIds <*>) . pure) (const (pure [])) specArg
(gamma,delta) <- mkEnv ctx
let (specTyBndrs,specTyVars) = unzip
$ map (\tv -> let ki = HML.lookupDefault (error $ $(curLoc) ++ show tv ++ " not found") tv delta
in (Right $ TyVar tv (embed ki), Right $ VarTy ki tv)) specFTVs
(specTmBndrs,specTmVars) = unzip
$ map (\tm -> let ty = HML.lookupDefault (error $ $(curLoc) ++ show tm ++ " not found") tm gamma
in (Left $ Id tm (embed ty), Left $ Var ty tm)) specFVs
return (specTyBndrs ++ specTmBndrs,specTyVars ++ specTmVars)