module CLaSH.Normalize.Util where
import Control.Lens ((%=),(^.),_3)
import qualified Control.Lens as Lens
import Data.Function (on)
import qualified Data.Graph as Graph
import Data.Graph.Inductive (Gr,LNode,lsuc,mkGraph,iDom)
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HashMap
import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set
import qualified Data.Set.Lens as Lens
import Unbound.Generics.LocallyNameless (Fresh, bind, embed, rec)
import SrcLoc (SrcSpan)
import CLaSH.Core.FreeVars (termFreeIds)
import CLaSH.Core.Var (Var (Id))
import CLaSH.Core.Term (Term (..), TmName)
import CLaSH.Core.Type (Type)
import CLaSH.Core.TyCon (TyCon, TyConName)
import CLaSH.Core.Util (collectArgs, isPolyFun)
import CLaSH.Normalize.Types
import CLaSH.Rewrite.Types (bindings,extra)
import CLaSH.Rewrite.Util (specialise)
import CLaSH.Util (curLoc)
alreadyInlined :: TmName
-> TmName
-> NormalizeMonad (Maybe Int)
alreadyInlined f cf = do
inlinedHM <- Lens.use inlineHistory
case HashMap.lookup cf inlinedHM of
Nothing -> return Nothing
Just inlined' -> return (HashMap.lookup f inlined')
addNewInline :: TmName
-> TmName
-> NormalizeMonad ()
addNewInline f cf =
inlineHistory %= HashMap.insertWith
(\_ hm -> HashMap.insertWith (+) f 1 hm)
cf
(HashMap.singleton f 1)
specializeNorm :: NormRewrite
specializeNorm = specialise specialisationCache specialisationHistory specialisationLimit
isClosed :: (Functor m, Fresh m)
=> HashMap TyConName TyCon
-> Term
-> m Bool
isClosed tcm = fmap not . isPolyFun tcm
isConstant :: Term -> Bool
isConstant e = case collectArgs e of
(Data _, args) -> all (either isConstant (const True)) args
(Prim _ _, args) -> all (either isConstant (const True)) args
(Literal _,_) -> True
_ -> False
isRecursiveBndr :: TmName -> NormalizeSession Bool
isRecursiveBndr f = do
cg <- Lens.use (extra.recursiveComponents)
case HashMap.lookup f cg of
Just isR -> return isR
Nothing -> do
bndrs <- Lens.use bindings
let cg' = callGraph [] bndrs f
rcs = concat $ mkRecursiveComponents cg'
isR = f `elem` rcs
cg'' = HashMap.fromList
$ map (\(t,_) -> (t,t `elem` rcs)) cg'
(extra.recursiveComponents) %= HashMap.union cg''
return isR
callGraph :: [TmName]
-> HashMap TmName (Type,SrcSpan,Term)
-> TmName
-> [(TmName,[TmName])]
callGraph visited bindingMap root = node:other
where
rootTm = Maybe.fromMaybe (error $ show root ++ " is not a global binder") $ HashMap.lookup root bindingMap
used = Set.toList $ Lens.setOf termFreeIds (rootTm ^. _3)
node = (root,used)
other = concatMap (callGraph (root:visited) bindingMap) (filter (`notElem` visited) used)
mkRecursiveComponents :: [(TmName,[TmName])]
-> [[TmName]]
mkRecursiveComponents cg = map (List.sortBy (compare `on` (`List.elemIndex` fs)))
. Maybe.catMaybes
. map (\case {Graph.CyclicSCC vs -> Just vs; _ -> Nothing})
. Graph.stronglyConnComp
$ map (\(n,es) -> (n,n,es)) cg
where
fs = map fst cg
lambdaDropPrep :: HashMap TmName (Type,SrcSpan,Term)
-> TmName
-> HashMap TmName (Type,SrcSpan,Term)
lambdaDropPrep bndrs topEntity = bndrs'
where
depGraph = callGraph [] bndrs topEntity
used = HashMap.fromList depGraph
rcs = mkRecursiveComponents depGraph
dropped = map (lambdaDrop bndrs used) rcs
bndrs' = foldr (\(k,v) b -> HashMap.insert k v b) bndrs dropped
lambdaDrop :: HashMap TmName (Type,SrcSpan,Term)
-> HashMap TmName [TmName]
-> [TmName]
-> (TmName,(Type,SrcSpan,Term))
lambdaDrop bndrs depGraph cyc@(root:_) = block
where
doms = dominator depGraph cyc
block = blockSink bndrs doms (0,root)
lambdaDrop _ _ [] = error $ $(curLoc) ++ "Can't lambdadrop empty cycle"
dominator :: HashMap TmName [TmName]
-> [TmName]
-> Gr TmName TmName
dominator cfg cyc = mkGraph nodes (map (\(e,b) -> (b,e,nodesM HashMap.! e)) doms)
where
nodes = zip [0..] cyc
nodesM = HashMap.fromList nodes
nodesI = HashMap.fromList $ zip cyc [0..]
cycEdges = HashMap.map ( map (nodesI HashMap.!)
. filter (`elem` cyc)
)
$ HashMap.filterWithKey (\k _ -> k `elem` cyc) cfg
edges = concatMap (\(i,n) -> zip3 (repeat i) (cycEdges HashMap.! n) (repeat ())
) nodes
graph = mkGraph nodes edges :: Gr TmName ()
doms = iDom graph 0
blockSink :: HashMap TmName (Type,SrcSpan,Term)
-> Gr TmName TmName
-> LNode TmName
-> (TmName,(Type,SrcSpan,Term))
blockSink bndrs doms (nId,tmName) = (tmName,(ty,sp,newTm))
where
(ty,sp,tm) = bndrs HashMap.! tmName
sucTm = lsuc doms nId
tmS = map (blockSink bndrs doms) sucTm
bnds = map (\(tN,(ty',_,tm')) -> (Id tN (embed ty'),embed tm')) tmS
newTm = case sucTm of
[] -> tm
_ -> Letrec (bind (rec bnds) tm)