{-|
  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 CPP #-}
{-# 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

#if MIN_VERSION_ghc(9,8,4) || (MIN_VERSION_ghc(9,6,7) && !MIN_VERSION_ghc(9,8,0))
#define UNIQUE_IS_WORD64
#endif

import qualified Control.Lens           as Lens
import Control.Lens.Fold                (Fold)
import Control.Lens.Getter              (Contravariant)
import Data.Coerce
#ifdef UNIQUE_IS_WORD64
import qualified GHC.Data.Word64Set     as IntSet
#else
import qualified Data.IntSet            as IntSet
#endif
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 :: Fold Type TyVar
typeFreeVars = (forall b. Var b -> Bool)
-> Word64Set -> (TyVar -> f TyVar) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' (Bool -> Var b -> Bool
forall a b. a -> b -> a
const Bool
True) Word64Set
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
#ifdef UNIQUE_IS_WORD64
  -> IntSet.Word64Set
#else
  -> IntSet.IntSet
#endif
  -- ^ Uniques of the variables in scope, used by 'termFreeVars''
  -> (Var a -> f (Var a))
  -> Type
  -> f Type
typeFreeVars' :: forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' forall b. Var b -> Bool
interesting Word64Set
is Var a -> f (Var a)
f = Word64Set -> Type -> f Type
go Word64Set
is where
  go :: Word64Set -> Type -> f Type
go Word64Set
inScope = \case
    VarTy TyVar
tv -> f Type
tv1 f Type -> f Type -> f Type
forall a b. f a -> f b -> f a
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f a
<* Word64Set -> Type -> f Type
go Word64Set
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 -> Word64Set -> Bool
`IntSet.member` Word64Set
inScope
        inScope1 :: Word64Set
inScope1
          | Bool
tvInScope = Word64Set
inScope
          | Bool
otherwise = Word64Set
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
forall a b. Coercible a b => a -> b
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
forall a b. Coercible a b => a -> b
coerce TyVar
tv)
            | Bool
otherwise
            = Type -> f Type
forall a. a -> f a
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
<$> Word64Set -> TyVar -> f TyVar
goBndr Word64Set
inScope TyVar
tv
                               f (Type -> Type) -> f Type -> f Type
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Type -> f Type
go (Unique -> Word64Set -> Word64Set
IntSet.insert (TyVar -> Unique
forall a. Var a -> Unique
varUniq TyVar
tv) Word64Set
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
<$> Word64Set -> Type -> f Type
go Word64Set
inScope Type
l f (Type -> Type) -> f Type -> f Type
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Type -> f Type
go Word64Set
inScope Type
r
    Type
ty -> Type -> f Type
forall a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Type
ty

  goBndr :: Word64Set -> TyVar -> f TyVar
goBndr Word64Set
inScope TyVar
tv = (\Type
t -> TyVar
tv {varType = t}) (Type -> TyVar) -> f Type -> f TyVar
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Word64Set -> Type -> f Type
go Word64Set
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 :: forall a. [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 (f :: Type -> Type).
(Contravariant f, Applicative f) =>
(Var a -> f (Var a)) -> Term -> f Term
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 :: forall a (f :: Type -> Type).
(Contravariant f, Applicative f) =>
(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' Var b -> Bool
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 :: Fold Term Id
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' Var b -> Bool
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 :: Fold Term Id
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' Var b -> Bool
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 :: Fold Term Id
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' Var b -> Bool
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 :: Fold Term TyVar
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' Var b -> Bool
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 (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
interesting Var a -> f (Var a)
f = Word64Set -> Term -> f Term
go Word64Set
IntSet.empty where
  go :: Word64Set -> Term -> f Term
go Word64Set
inLocalScope = \case
    Var Id
v -> f Term
v1 f Term -> f Type -> f Term
forall a b. f a -> f b -> f a
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f a
<* (forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' Var b -> Bool
forall b. Var b -> Bool
interesting Word64Set
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 -> Word64Set -> Bool
`IntSet.member` Word64Set
inLocalScope
        inLocalScope1 :: Word64Set
inLocalScope1
          | Bool
vInScope  = Word64Set
inLocalScope
          | Bool
otherwise = Word64Set
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
forall a b. Coercible a b => a -> b
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
forall a b. Coercible a b => a -> b
coerce Id
v)
           | Bool
otherwise
           = Term -> f Term
forall a. a -> f a
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
<$> Word64Set -> Id -> f Id
forall {a}. Word64Set -> Var a -> f (Var a)
goBndr Word64Set
inLocalScope Id
id_
          f (Term -> Term) -> f Term -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go (Unique -> Word64Set -> Word64Set
IntSet.insert (Id -> Unique
forall a. Var a -> Unique
varUniq Id
id_) Word64Set
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
<$> Word64Set -> TyVar -> f TyVar
forall {a}. Word64Set -> Var a -> f (Var a)
goBndr Word64Set
inLocalScope TyVar
tv
            f (Term -> Term) -> f Term -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go (Unique -> Word64Set -> Word64Set
IntSet.insert (TyVar -> Unique
forall a. Var a -> Unique
varUniq TyVar
tv) Word64Set
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
<$> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
l f (Term -> Term) -> f Term -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go Word64Set
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
<$> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
l
            f (Type -> Term) -> f Type -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' Var b -> Bool
forall b. Var b -> Bool
interesting Word64Set
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
<$> Word64Set -> Id -> f Id
forall {a}. Word64Set -> Var a -> f (Var a)
goBndr Word64Set
inLocalScope Id
i f (Term -> Bind Term) -> f Term -> f (Bind Term)
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
x)
          f (Term -> Term) -> f Term -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go (Unique -> Word64Set -> Word64Set
IntSet.insert (Id -> Unique
forall a. Var a -> Unique
varUniq Id
i) Word64Set
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)
forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Word64Set -> (Id, Term) -> f (Id, Term)
goBind Word64Set
inLocalScope') [(Id, Term)]
bs)
          f (Term -> Term) -> f Term -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go Word64Set
inLocalScope' Term
e
     where
      inLocalScope' :: Word64Set
inLocalScope' = ((Id, Term) -> Word64Set -> Word64Set)
-> Word64Set -> [(Id, Term)] -> Word64Set
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Unique -> Word64Set -> Word64Set
IntSet.insert (Unique -> Word64Set -> Word64Set)
-> ((Id, Term) -> Unique) -> (Id, Term) -> Word64Set -> Word64Set
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) Word64Set
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
<$> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
subj
           f (Type -> [Alt] -> Term) -> f Type -> f ([Alt] -> Term)
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' Var b -> Bool
forall b. Var b -> Bool
interesting Word64Set
inLocalScope Var a -> f (Var a)
f Type
ty
           f ([Alt] -> Term) -> f [Alt] -> f Term
forall a b. f (a -> b) -> f a -> f b
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)
forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Word64Set -> Alt -> f Alt
goAlt Word64Set
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
<$> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
tm
           f (Type -> Type -> Term) -> f Type -> f (Type -> Term)
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' Var b -> Bool
forall b. Var b -> Bool
interesting Word64Set
inLocalScope Var a -> f (Var a)
f Type
t1
           f (Type -> Term) -> f Type -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' Var b -> Bool
forall b. Var b -> Bool
interesting Word64Set
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
<$> Word64Set -> TickInfo -> f TickInfo
goTick Word64Set
inLocalScope TickInfo
tick
      f (Term -> Term) -> f Term -> f Term
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
tm

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

  goBndr :: Word64Set -> Var a -> f (Var a)
goBndr Word64Set
inLocalScope Var a
v =
    (\Type
t -> Var a
v  {varType = 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)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' Var b -> Bool
forall b. Var b -> Bool
interesting Word64Set
inLocalScope Var a -> f (Var a)
f (Var a -> Type
forall a. Var a -> Type
varType Var a
v)

  goBind :: Word64Set -> (Id, Term) -> f (Id, Term)
goBind Word64Set
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
<$> Word64Set -> Id -> f Id
forall {a}. Word64Set -> Var a -> f (Var a)
goBndr Word64Set
inLocalScope Id
l f (Term -> (Id, Term)) -> f Term -> f (Id, Term)
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
r

  goAlt :: Word64Set -> Alt -> f Alt
goAlt Word64Set
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 a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure DataCon
dc
                                           f ([TyVar] -> [Id] -> Pat) -> f [TyVar] -> f ([Id] -> Pat)
forall a b. f (a -> b) -> f a -> f b
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)
forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Word64Set -> TyVar -> f TyVar
forall {a}. Word64Set -> Var a -> f (Var a)
goBndr Word64Set
inLocalScope') [TyVar]
tvs
                                           f ([Id] -> Pat) -> f [Id] -> f Pat
forall a b. f (a -> b) -> f a -> f b
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)
forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Word64Set -> Id -> f Id
forall {a}. Word64Set -> Var a -> f (Var a)
goBndr Word64Set
inLocalScope') [Id]
ids)
                              f (Term -> Alt) -> f Term -> f Alt
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go Word64Set
inLocalScope' Term
alt
      where
        inLocalScope' :: Word64Set
inLocalScope' = (Unique -> Word64Set -> Word64Set)
-> Word64Set -> [Unique] -> Word64Set
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Unique -> Word64Set -> Word64Set
IntSet.insert
                         ((Unique -> Word64Set -> Word64Set)
-> Word64Set -> [Unique] -> Word64Set
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Unique -> Word64Set -> Word64Set
IntSet.insert Word64Set
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 a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Pat
pat f (Term -> Alt) -> f Term -> f Alt
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Word64Set -> Term -> f Term
go Word64Set
inLocalScope Term
alt

  goTick :: Word64Set -> TickInfo -> f TickInfo
goTick Word64Set
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)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
forall (f :: Type -> Type) a.
(Contravariant f, Applicative f) =>
(forall b. Var b -> Bool)
-> Word64Set -> (Var a -> f (Var a)) -> Type -> f Type
typeFreeVars' Var b -> Bool
forall b. Var b -> Bool
interesting Word64Set
inLocalScope Var a -> f (Var a)
f Type
ty
    TickInfo
tick         -> TickInfo -> f TickInfo
forall a. a -> f a
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 Int
countFreeOccurances =
  Fold Term Id
-> (VarEnv Int -> VarEnv Int -> VarEnv Int)
-> VarEnv Int
-> (Id -> VarEnv Int)
-> Term
-> VarEnv Int
forall s a r. Fold s a -> (r -> r -> r) -> r -> (a -> r) -> s -> r
Lens.foldMapByOf (Id -> f Id) -> Term -> f Term
Fold Term Id
freeLocalIds ((Int -> Int -> Int) -> VarEnv Int -> VarEnv Int -> VarEnv Int
forall a. (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
unionVarEnvWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)) VarEnv Int
forall a. VarEnv a
emptyVarEnv
                   (Id -> Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a
`unitVarEnv` (Int
1 :: Int))