{-|
  Copyright   :  (C) 2012-2016, University of Twente
                     2021,      QBayLogic B.V.
  License     :  BSD2 (see the file LICENSE)
  Maintainer  :  QBayLogic B.V. <devops@qbaylogic.com>

  Free variable calculations
-}

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}

module Clash.Core.FreeVars
  (-- * Free variable calculation
    typeFreeVars
  , freeIds
  , freeLocalVars
  , freeLocalIds
  , globalIds
  , termFreeTyVars
  -- * occurrence check
  , globalIdOccursIn
  , localVarsDoNotOccurIn
  , countFreeOccurances
  -- * Internal
  , 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 (..), Bind(..))
import Clash.Core.Type                  (Type (..))
import Clash.Core.Var
  (Id, IdScope (..), TyVar, Var (..), isLocalId)
import Clash.Core.VarEnv
  (VarEnv, emptyVarEnv, unionVarEnvWith, unitVarEnv)

-- | Gives the free type-variables in a Type, implemented as a 'Fold'
--
-- The 'Fold' is closed over the types of its variables, so:
--
-- @
-- foldMapOf typeFreeVars unitVarSet ((a:* -> k) Int) = {a, k}
-- @
typeFreeVars :: Fold Type TyVar
typeFreeVars :: (TyVar -> f TyVar) -> Type -> f Type
typeFreeVars = (forall b. Var b -> Bool)
-> IntSet -> (TyVar -> f TyVar) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' (Bool -> Var b -> Bool
forall a b. a -> b -> a
const Bool
True) IntSet
IntSet.empty

-- | Gives the "interesting" free variables in a Type, implemented as a 'Fold'
--
-- The 'Fold' is closed over the types of variables, so:
--
-- @
-- foldMapOf (typeFreeVars' (const True) IntSet.empty) unitVarSet ((a:* -> k) Int) = {a, k}
-- @
--
-- Note [Closing over kind variables]
--
-- Consider the type
--
-- > forall k . b -> k
--
-- where
--
-- > b :: k -> Type
--
-- When we close over the free variables of @forall k . b -> k@, i.e. @b@, then
-- the @k@ in @b :: k -> Type@ is most definitely /not/ the @k@ in
-- @forall k . b -> k@. So when a type variable is free, i.e. not in the inScope
-- set, its kind variables also aren´t; so in order to prevent collisions due to
-- shadowing we close using an empty inScope set.
--
-- See also: https://gitlab.haskell.org/ghc/ghc/-/commit/503514b94f8dc7bd9eab5392206649aee45f140b
typeFreeVars'
  :: (Contravariant f, Applicative f)
  => (forall b . Var b -> Bool)
  -- ^ Predicate telling whether a variable is interesting
  -> IntSet.IntSet
  -- ^ Uniques of the variables in scope, used by 'termFreeVars''
  -> (Var a -> f (Var a))
  -> Type
  -> f Type
typeFreeVars' :: (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
is Var a -> f (Var a)
f = IntSet -> Type -> f Type
go IntSet
is where
  go :: IntSet -> Type -> f Type
go IntSet
inScope = \case
    VarTy TyVar
tv -> f Type
tv1 f Type -> f Type -> f Type
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f a
<* IntSet -> Type -> f Type
go IntSet
inScope1 (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv)
      where
        isInteresting :: Bool
isInteresting = TyVar -> Bool
forall b. Var b -> Bool
interesting TyVar
tv
        tvInScope :: Bool
tvInScope     = TyVar -> Unique
forall a. Var a -> Unique
varUniq TyVar
tv Unique -> IntSet -> Bool
`IntSet.member` IntSet
inScope
        inScope1 :: IntSet
inScope1
          | Bool
tvInScope = IntSet
inScope
          | Bool
otherwise = IntSet
IntSet.empty -- See Note [Closing over type variables]

        tv1 :: f Type
tv1 | Bool
isInteresting
            , Bool -> Bool
not Bool
tvInScope
            = TyVar -> Type
VarTy (TyVar -> Type) -> (Var a -> TyVar) -> Var a -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var a -> TyVar
coerce (Var a -> Type) -> f (Var a) -> f Type
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Var a -> f (Var a)
f (TyVar -> Var a
coerce TyVar
tv)
            | Bool
otherwise
            = Type -> f Type
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (TyVar -> Type
VarTy TyVar
tv)

    ForAllTy TyVar
tv Type
ty -> TyVar -> Type -> Type
ForAllTy (TyVar -> Type -> Type) -> f TyVar -> f (Type -> Type)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> TyVar -> f TyVar
goBndr IntSet
inScope TyVar
tv
                               f (Type -> Type) -> f Type -> f Type
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Type -> f Type
go (Unique -> IntSet -> IntSet
IntSet.insert (TyVar -> Unique
forall a. Var a -> Unique
varUniq TyVar
tv) IntSet
inScope) Type
ty
    AppTy Type
l Type
r -> Type -> Type -> Type
AppTy (Type -> Type -> Type) -> f Type -> f (Type -> Type)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Type -> f Type
go IntSet
inScope Type
l f (Type -> Type) -> f Type -> f Type
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Type -> f Type
go IntSet
inScope Type
r
    Type
ty -> Type -> f Type
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Type
ty

  goBndr :: IntSet -> TyVar -> f TyVar
goBndr IntSet
inScope TyVar
tv = (\Type
t -> TyVar
tv {varType :: Type
varType = Type
t}) (Type -> TyVar) -> f Type -> f TyVar
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Type -> f Type
go IntSet
inScope (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv)

-- | Check whether a set of variables does not occur free in a term
localVarsDoNotOccurIn
  :: [Var a]
  -> Term
  -> Bool
localVarsDoNotOccurIn :: [Var a] -> Term -> Bool
localVarsDoNotOccurIn [Var a]
vs Term
e =
  All -> Bool
getAll (Getting All Term (Var a) -> (Var a -> All) -> Term -> All
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting All Term (Var a)
forall a. Fold Term (Var a)
freeLocalVars (Bool -> All
All (Bool -> All) -> (Var a -> Bool) -> Var a -> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var a -> [Var a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [Var a]
vs)) Term
e)

-- | Check whether a local identifier occurs free in a term
globalIdOccursIn
  :: Id
  -> Term
  -> Bool
globalIdOccursIn :: Id -> Term -> Bool
globalIdOccursIn Id
v Term
e = Any -> Bool
getAny (Getting Any Term Id -> (Id -> Any) -> Term -> Any
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting Any Term Id
Fold Term Id
globalIds (Bool -> Any
Any (Bool -> Any) -> (Id -> Bool) -> Id -> Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
v)) Term
e)

-- | Calculate the /local/ free variable of an expression: the free type
-- variables and the free identifiers that are not bound in the global
-- environment.
freeLocalVars :: Fold Term (Var a)
freeLocalVars :: (Var a -> f (Var a)) -> Term -> f Term
freeLocalVars = (forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
termFreeVars' forall b. Var b -> Bool
isLocalVar where
  isLocalVar :: Var a -> Bool
isLocalVar (Id {idScope :: forall a. Var a -> IdScope
idScope = IdScope
GlobalId}) = Bool
False
  isLocalVar Var a
_ = Bool
True

-- | Gives the free identifiers of a Term, implemented as a 'Fold'
freeIds :: Fold Term Id
freeIds :: (Id -> f Id) -> Term -> f Term
freeIds = (forall b. Var b -> Bool) -> (Id -> f Id) -> Term -> f Term
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
termFreeVars' forall b. Var b -> Bool
isId where
  isId :: Var a -> Bool
isId (Id {}) = Bool
True
  isId Var a
_       = Bool
False

-- | Calculate the /local/ free identifiers of an expression: the free
-- identifiers that are not bound in the global environment.
freeLocalIds :: Fold Term Id
freeLocalIds :: (Id -> f Id) -> Term -> f Term
freeLocalIds = (forall b. Var b -> Bool) -> (Id -> f Id) -> Term -> f Term
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
termFreeVars' forall b. Var b -> Bool
isLocalId

-- | Calculate the /global/ free identifiers of an expression: the free
-- identifiers that are bound in the global environment.
globalIds :: Fold Term Id
globalIds :: (Id -> f Id) -> Term -> f Term
globalIds = (forall b. Var b -> Bool) -> (Id -> f Id) -> Term -> f Term
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
termFreeVars' forall b. Var b -> Bool
isGlobalId where
  isGlobalId :: Var a -> Bool
isGlobalId (Id {idScope :: forall a. Var a -> IdScope
idScope = IdScope
GlobalId}) = Bool
True
  isGlobalId Var a
_ = Bool
False

-- | Gives the free type-variables of a Term, implemented as a 'Fold'
--
-- The 'Fold' is closed over the types of variables, so:
--
-- @
-- foldMapOf termFreeTyVars unitVarSet (case (x : (a:* -> k) Int)) of {}) = {a, k}
-- @
termFreeTyVars :: Fold Term TyVar
termFreeTyVars :: (TyVar -> f TyVar) -> Term -> f Term
termFreeTyVars = (forall b. Var b -> Bool) -> (TyVar -> f TyVar) -> Term -> f Term
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
termFreeVars' forall b. Var b -> Bool
isTV where
  isTV :: Var a -> Bool
isTV (TyVar {}) = Bool
True
  isTV Var a
_          = Bool
False

-- | Gives the "interesting" free variables in a Term, implemented as a 'Fold'
--
-- The 'Fold' is closed over the types of variables, so:
--
-- @
-- foldMapOf (termFreeVars' (const True)) unitVarSet (case (x : (a:* -> k) Int)) of {}) = {x, a, k}
-- @
--
-- Note [Closing over type variables]
--
-- Consider the term
--
-- > /\(k :: Type) -> \(b :: k) -> a
--
-- where
--
-- > a :: k
--
-- When we close over the free variables of @/\k -> \(b :: k) -> (a :: k)@, i.e.
-- @a@, then the @k@ in @a :: k@ is most definitely /not/ the @k@ in introduced
-- by the @/\k ->@. So when a term variable is free, i.e. not in the inScope
-- set, its type variables also aren´t; so in order to prevent collisions due to
-- shadowing we close using an empty inScope set.
--
-- See also: https://gitlab.haskell.org/ghc/ghc/-/commit/503514b94f8dc7bd9eab5392206649aee45f140b
termFreeVars'
  :: (Contravariant f, Applicative f)
  => (forall b . Var b -> Bool)
  -- ^ Predicate telling whether a variable is interesting
  -> (Var a -> f (Var a))
  -> Term
  -> f Term
termFreeVars' :: (forall b. Var b -> Bool) -> (Var a -> f (Var a)) -> Term -> f Term
termFreeVars' forall b. Var b -> Bool
interesting Var a -> f (Var a)
f = IntSet -> Term -> f Term
go IntSet
IntSet.empty where
  go :: IntSet -> Term -> f Term
go IntSet
inLocalScope = \case
    Var Id
v -> f Term
v1 f Term -> f Type -> f Term
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f a
<* (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
inLocalScope1 Var a -> f (Var a)
f (Id -> Type
forall a. Var a -> Type
varType Id
v)
      where
        isInteresting :: Bool
isInteresting = Id -> Bool
forall b. Var b -> Bool
interesting Id
v
        vInScope :: Bool
vInScope      = Id -> Bool
forall b. Var b -> Bool
isLocalId Id
v Bool -> Bool -> Bool
&& Id -> Unique
forall a. Var a -> Unique
varUniq Id
v Unique -> IntSet -> Bool
`IntSet.member` IntSet
inLocalScope
        inLocalScope1 :: IntSet
inLocalScope1
          | Bool
vInScope  = IntSet
inLocalScope
          | Bool
otherwise = IntSet
IntSet.empty -- See Note [Closing over type variables]

        v1 :: f Term
v1 | Bool
isInteresting
           , Bool -> Bool
not Bool
vInScope
           = Id -> Term
Var (Id -> Term) -> (Var a -> Id) -> Var a -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var a -> Id
coerce (Var a -> Term) -> f (Var a) -> f Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Var a -> f (Var a)
f (Id -> Var a
coerce Id
v)
           | Bool
otherwise
           = Term -> f Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Id -> Term
Var Id
v)

    Lam Id
id_ Term
tm ->
      Id -> Term -> Term
Lam (Id -> Term -> Term) -> f Id -> f (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Id -> f Id
forall a. IntSet -> Var a -> f (Var a)
goBndr IntSet
inLocalScope Id
id_
          f (Term -> Term) -> f Term -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go (Unique -> IntSet -> IntSet
IntSet.insert (Id -> Unique
forall a. Var a -> Unique
varUniq Id
id_) IntSet
inLocalScope) Term
tm
    TyLam TyVar
tv Term
tm ->
      TyVar -> Term -> Term
TyLam (TyVar -> Term -> Term) -> f TyVar -> f (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> TyVar -> f TyVar
forall a. IntSet -> Var a -> f (Var a)
goBndr IntSet
inLocalScope TyVar
tv
            f (Term -> Term) -> f Term -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go (Unique -> IntSet -> IntSet
IntSet.insert (TyVar -> Unique
forall a. Var a -> Unique
varUniq TyVar
tv) IntSet
inLocalScope) Term
tm

    App Term
l Term
r ->
      Term -> Term -> Term
App (Term -> Term -> Term) -> f Term -> f (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
l f (Term -> Term) -> f Term -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
r

    TyApp Term
l Type
r ->
      Term -> Type -> Term
TyApp (Term -> Type -> Term) -> f Term -> f (Type -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
l
            f (Type -> Term) -> f Type -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
inLocalScope Var a -> f (Var a)
f Type
r

    Let (NonRec Id
i Term
x) Term
e ->
      Bind Term -> Term -> Term
Let (Bind Term -> Term -> Term) -> f (Bind Term) -> f (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Id -> Term -> Bind Term
forall a. Id -> a -> Bind a
NonRec (Id -> Term -> Bind Term) -> f Id -> f (Term -> Bind Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Id -> f Id
forall a. IntSet -> Var a -> f (Var a)
goBndr IntSet
inLocalScope Id
i f (Term -> Bind Term) -> f Term -> f (Bind Term)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
x)
          f (Term -> Term) -> f Term -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go (Unique -> IntSet -> IntSet
IntSet.insert (Id -> Unique
forall a. Var a -> Unique
varUniq Id
i) IntSet
inLocalScope) Term
e

    Let (Rec [(Id, Term)]
bs) Term
e ->
      Bind Term -> Term -> Term
Let (Bind Term -> Term -> Term) -> f (Bind Term) -> f (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(Id, Term)] -> Bind Term
forall a. [(Id, a)] -> Bind a
Rec ([(Id, Term)] -> Bind Term) -> f [(Id, Term)] -> f (Bind Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Id, Term) -> f (Id, Term)) -> [(Id, Term)] -> f [(Id, Term)]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (IntSet -> (Id, Term) -> f (Id, Term)
goBind IntSet
inLocalScope') [(Id, Term)]
bs)
          f (Term -> Term) -> f Term -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go IntSet
inLocalScope' Term
e
     where
      inLocalScope' :: IntSet
inLocalScope' = ((Id, Term) -> IntSet -> IntSet)
-> IntSet -> [(Id, Term)] -> IntSet
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Unique -> IntSet -> IntSet
IntSet.insert (Unique -> IntSet -> IntSet)
-> ((Id, Term) -> Unique) -> (Id, Term) -> IntSet -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Unique
forall a. Var a -> Unique
varUniq (Id -> Unique) -> ((Id, Term) -> Id) -> (Id, Term) -> Unique
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, Term) -> Id
forall a b. (a, b) -> a
fst) IntSet
inLocalScope [(Id, Term)]
bs

    Case Term
subj Type
ty [Alt]
alts ->
      Term -> Type -> [Alt] -> Term
Case (Term -> Type -> [Alt] -> Term)
-> f Term -> f (Type -> [Alt] -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
subj
           f (Type -> [Alt] -> Term) -> f Type -> f ([Alt] -> Term)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
inLocalScope Var a -> f (Var a)
f Type
ty
           f ([Alt] -> Term) -> f [Alt] -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (Alt -> f Alt) -> [Alt] -> f [Alt]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (IntSet -> Alt -> f Alt
goAlt IntSet
inLocalScope) [Alt]
alts

    Cast Term
tm Type
t1 Type
t2 ->
      Term -> Type -> Type -> Term
Cast (Term -> Type -> Type -> Term)
-> f Term -> f (Type -> Type -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
tm
           f (Type -> Type -> Term) -> f Type -> f (Type -> Term)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
inLocalScope Var a -> f (Var a)
f Type
t1
           f (Type -> Term) -> f Type -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
inLocalScope Var a -> f (Var a)
f Type
t2

    Tick TickInfo
tick Term
tm ->
      TickInfo -> Term -> Term
Tick (TickInfo -> Term -> Term) -> f TickInfo -> f (Term -> Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> TickInfo -> f TickInfo
goTick IntSet
inLocalScope TickInfo
tick
      f (Term -> Term) -> f Term -> f Term
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
tm

    Term
tm -> Term -> f Term
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Term
tm

  goBndr :: IntSet -> Var a -> f (Var a)
goBndr IntSet
inLocalScope Var a
v =
    (\Type
t -> Var a
v  {varType :: Type
varType = Type
t}) (Type -> Var a) -> f Type -> f (Var a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
inLocalScope Var a -> f (Var a)
f (Var a -> Type
forall a. Var a -> Type
varType Var a
v)

  goBind :: IntSet -> (Id, Term) -> f (Id, Term)
goBind IntSet
inLocalScope (Id
l,Term
r) = (,) (Id -> Term -> (Id, Term)) -> f Id -> f (Term -> (Id, Term))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> Id -> f Id
forall a. IntSet -> Var a -> f (Var a)
goBndr IntSet
inLocalScope Id
l f (Term -> (Id, Term)) -> f Term -> f (Id, Term)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
r

  goAlt :: IntSet -> Alt -> f Alt
goAlt IntSet
inLocalScope (Pat
pat,Term
alt) = case Pat
pat of
    DataPat DataCon
dc [TyVar]
tvs [Id]
ids -> (,) (Pat -> Term -> Alt) -> f Pat -> f (Term -> Alt)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (DataCon -> [TyVar] -> [Id] -> Pat
DataPat (DataCon -> [TyVar] -> [Id] -> Pat)
-> f DataCon -> f ([TyVar] -> [Id] -> Pat)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> DataCon -> f DataCon
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure DataCon
dc
                                           f ([TyVar] -> [Id] -> Pat) -> f [TyVar] -> f ([Id] -> Pat)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (TyVar -> f TyVar) -> [TyVar] -> f [TyVar]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (IntSet -> TyVar -> f TyVar
forall a. IntSet -> Var a -> f (Var a)
goBndr IntSet
inLocalScope') [TyVar]
tvs
                                           f ([Id] -> Pat) -> f [Id] -> f Pat
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (Id -> f Id) -> [Id] -> f [Id]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (IntSet -> Id -> f Id
forall a. IntSet -> Var a -> f (Var a)
goBndr IntSet
inLocalScope') [Id]
ids)
                              f (Term -> Alt) -> f Term -> f Alt
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go IntSet
inLocalScope' Term
alt
      where
        inLocalScope' :: IntSet
inLocalScope' = (Unique -> IntSet -> IntSet) -> IntSet -> [Unique] -> IntSet
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Unique -> IntSet -> IntSet
IntSet.insert
                         ((Unique -> IntSet -> IntSet) -> IntSet -> [Unique] -> IntSet
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Unique -> IntSet -> IntSet
IntSet.insert IntSet
inLocalScope ((TyVar -> Unique) -> [TyVar] -> [Unique]
forall a b. (a -> b) -> [a] -> [b]
map TyVar -> Unique
forall a. Var a -> Unique
varUniq [TyVar]
tvs))
                         ((Id -> Unique) -> [Id] -> [Unique]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Unique
forall a. Var a -> Unique
varUniq [Id]
ids)
    Pat
_ -> (,) (Pat -> Term -> Alt) -> f Pat -> f (Term -> Alt)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Pat -> f Pat
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Pat
pat f (Term -> Alt) -> f Term -> f Alt
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> IntSet -> Term -> f Term
go IntSet
inLocalScope Term
alt

  goTick :: IntSet -> TickInfo -> f TickInfo
goTick IntSet
inLocalScope = \case
    NameMod NameMod
m Type
ty -> NameMod -> Type -> TickInfo
NameMod NameMod
m (Type -> TickInfo) -> f Type -> f TickInfo
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> IntSet -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting IntSet
inLocalScope Var a -> f (Var a)
f Type
ty
    TickInfo
tick         -> TickInfo -> f TickInfo
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure TickInfo
tick

-- | Get the free variables of an expression and count the number of occurrences
countFreeOccurances
  :: Term
  -> VarEnv Int
countFreeOccurances :: Term -> VarEnv Unique
countFreeOccurances =
  Fold Term Id
-> (VarEnv Unique -> VarEnv Unique -> VarEnv Unique)
-> VarEnv Unique
-> (Id -> VarEnv Unique)
-> Term
-> VarEnv Unique
forall s a r. Fold s a -> (r -> r -> r) -> r -> (a -> r) -> s -> r
Lens.foldMapByOf Fold Term Id
freeLocalIds ((Unique -> Unique -> Unique)
-> VarEnv Unique -> VarEnv Unique -> VarEnv Unique
forall a. (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
unionVarEnvWith Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
(+)) VarEnv Unique
forall a. VarEnv a
emptyVarEnv
                   (Id -> Unique -> VarEnv Unique
forall b a. Var b -> a -> VarEnv a
`unitVarEnv` (Unique
1 :: Int))