{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ViewPatterns #-}
module Clash.Core.FreeVars
(
typeFreeVars
, termFreeVarsX
, freeIds
, freeLocalVars
, freeLocalIds
, globalIds
, termFreeTyVars
, tyFVsOfTypes
, localFVsOfTerms
, hasLocalFreeVars
, noFreeVarsOfType
, localIdOccursIn
, globalIdOccursIn
, localIdDoesNotOccurIn
, localIdsDoNotOccurIn
, localVarsDoNotOccurIn
, typeFreeVars'
, termFreeVars'
)
where
import qualified Control.Lens as Lens
import Control.Lens.Fold (Fold)
import Control.Lens.Getter (Contravariant)
import Data.Coerce
import qualified Data.IntSet as IntSet
import Data.Monoid (All (..), Any (..))
import Clash.Core.Term (Pat (..), Term (..), TickInfo (..))
import Clash.Core.Type (Type (..))
import Clash.Core.Var (Id, IdScope (..), TyVar, Var (..))
import Clash.Core.VarEnv (VarSet, unitVarSet)
typeFreeVars :: Fold Type TyVar
typeFreeVars = typeFreeVars' (const True) IntSet.empty
typeFreeVars'
:: (Contravariant f, Applicative f)
=> (forall b . Var b -> Bool)
-> IntSet.IntSet
-> (Var a -> f (Var a))
-> Type
-> f Type
typeFreeVars' interesting is f = go is where
go inScope = \case
VarTy tv -> tv1 <* go inScope1 (varType tv)
where
isInteresting = interesting tv
tvInScope = varUniq tv `IntSet.member` inScope
inScope1
| tvInScope = inScope
| otherwise = IntSet.empty
tv1 | isInteresting
, not tvInScope
= VarTy . coerce <$> f (coerce tv)
| otherwise
= pure (VarTy tv)
ForAllTy tv ty -> ForAllTy <$> goBndr inScope tv
<*> go (IntSet.insert (varUniq tv) inScope) ty
AppTy l r -> AppTy <$> go inScope l <*> go inScope r
ty -> pure ty
goBndr inScope tv = (\t -> tv {varType = t}) <$> go inScope (varType tv)
localIdDoesNotOccurIn
:: Id
-> Term
-> Bool
localIdDoesNotOccurIn v e = getAll (Lens.foldMapOf freeLocalIds (All . (/= v)) e)
localIdsDoNotOccurIn
:: [Id]
-> Term
-> Bool
localIdsDoNotOccurIn vs e =
getAll (Lens.foldMapOf freeLocalIds (All . (`notElem` vs)) e)
localVarsDoNotOccurIn
:: [Var a]
-> Term
-> Bool
localVarsDoNotOccurIn vs e =
getAll (Lens.foldMapOf freeLocalVars (All . (`notElem` vs)) e)
localIdOccursIn
:: Id
-> Term
-> Bool
localIdOccursIn v e = getAny (Lens.foldMapOf freeLocalIds (Any . (== v)) e)
globalIdOccursIn
:: Id
-> Term
-> Bool
globalIdOccursIn v e = getAny (Lens.foldMapOf globalIds (Any . (== v)) e)
freeLocalVars :: Fold Term (Var a)
freeLocalVars = termFreeVars' isLocalVar where
isLocalVar (Id {idScope = GlobalId}) = False
isLocalVar _ = True
freeIds :: Fold Term Id
freeIds = termFreeVars' isId where
isId (Id {}) = True
isId _ = False
freeLocalIds :: Fold Term Id
freeLocalIds = termFreeVars' isLocalId where
isLocalId (Id {idScope = LocalId}) = True
isLocalId _ = False
globalIds :: Fold Term Id
globalIds = termFreeVars' isGlobalId where
isGlobalId (Id {idScope = GlobalId}) = True
isGlobalId _ = False
hasLocalFreeVars :: Term -> Bool
hasLocalFreeVars = Lens.notNullOf freeLocalVars
termFreeTyVars :: Fold Term TyVar
termFreeTyVars = termFreeVars' isTV where
isTV (TyVar {}) = True
isTV _ = False
termFreeVarsX :: Fold Term (Var a)
termFreeVarsX = termFreeVars' (const True)
termFreeVars'
:: (Contravariant f, Applicative f)
=> (forall b . Var b -> Bool)
-> (Var a -> f (Var a))
-> Term
-> f Term
termFreeVars' interesting f = go IntSet.empty where
go inScope = \case
Var v -> v1 <* typeFreeVars' interesting inScope1 f (varType v)
where
isInteresting = interesting v
vInScope = varUniq v `IntSet.member` inScope
inScope1
| vInScope = inScope
| otherwise = IntSet.empty
v1 | isInteresting
, not vInScope
= Var . coerce <$> f (coerce v)
| otherwise
= pure (Var v)
Lam id_ tm -> Lam <$> goBndr inScope id_
<*> go (IntSet.insert (varUniq id_) inScope) tm
TyLam tv tm -> TyLam <$> goBndr inScope tv
<*> go (IntSet.insert (varUniq tv) inScope) tm
App l r -> App <$> go inScope l <*> go inScope r
TyApp l r -> TyApp <$> go inScope l
<*> typeFreeVars' interesting inScope f r
Letrec bs e -> Letrec <$> traverse (goBind inScope') bs <*> go inScope' e
where inScope' = foldr IntSet.insert inScope (map (varUniq.fst) bs)
Case subj ty alts -> Case <$> go inScope subj
<*> typeFreeVars' interesting inScope f ty
<*> traverse (goAlt inScope) alts
Cast tm t1 t2 -> Cast <$> go inScope tm
<*> typeFreeVars' interesting inScope f t1
<*> typeFreeVars' interesting inScope f t2
Tick tick tm -> Tick <$> goTick inScope tick <*> go inScope tm
tm -> pure tm
goBndr inScope v =
(\t -> v {varType = t}) <$> typeFreeVars' interesting inScope f (varType v)
goBind inScope (l,r) = (,) <$> goBndr inScope l <*> go inScope r
goAlt inScope (pat,alt) = case pat of
DataPat dc tvs ids -> (,) <$> (DataPat <$> pure dc
<*> traverse (goBndr inScope') tvs
<*> traverse (goBndr inScope') ids)
<*> go inScope' alt
where
inScope' = foldr IntSet.insert
(foldr IntSet.insert inScope (map varUniq tvs))
(map varUniq ids)
_ -> (,) <$> pure pat <*> go inScope alt
goTick inScope = \case
NameMod m ty -> NameMod m <$> typeFreeVars' interesting inScope f ty
tick -> pure tick
noFreeVarsOfType
:: Type
-> Bool
noFreeVarsOfType ty = case ty of
VarTy {} -> False
ForAllTy {} -> getAll (Lens.foldMapOf typeFreeVars (const (All False)) ty)
AppTy l r -> noFreeVarsOfType l && noFreeVarsOfType r
_ -> True
tyFVsOfTypes
:: Foldable f
=> f Type
-> VarSet
tyFVsOfTypes = foldMap go
where
go = Lens.foldMapOf typeFreeVars unitVarSet
localFVsOfTerms
:: Foldable f
=> f Term
-> VarSet
localFVsOfTerms = foldMap go
where
go = Lens.foldMapOf freeLocalVars unitVarSet