FreeVars.hs

{-|
  Copyright   :  (C) 2012-2016, University of Twente
                     2021-2026, 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'
  , termFreeIds'
  )
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 = typeFreeVars' (const True) 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' 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 -- See Note [Closing over type variables]

        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)

-- | Check whether a set of variables does not occur free in a term
localVarsDoNotOccurIn
  :: [Var a]
  -> Term
  -> Bool
localVarsDoNotOccurIn vs e =
  getAll (Lens.foldMapOf freeLocalVars (All . (`notElem` vs)) e)

-- | Check whether a local identifier occurs free in a term
globalIdOccursIn
  :: Id
  -> Term
  -> Bool
globalIdOccursIn v e = getAny (Lens.foldMapOf globalIds (Any . (== v)) 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 = termFreeVars' isLocalVar where
  isLocalVar (Id {idScope = GlobalId}) = False
  isLocalVar _ = True

-- | Gives the free identifiers of a Term, implemented as a 'Fold'
freeIds :: Fold Term Id
freeIds = termFreeIds' isId where
  isId (Id {}) = True
  isId _       = 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 = termFreeIds' isLocalId

-- | Calculate the /global/ free identifiers of an expression: the free
-- identifiers that are bound in the global environment.
globalIds :: Fold Term Id
globalIds = termFreeIds' isGlobalId where
  isGlobalId (Id {idScope = GlobalId}) = True
  isGlobalId _ = 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 = termFreeVars' isTV where
  isTV (TyVar {}) = True
  isTV _          = 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' = termFreeVarsWorker True

-- | Gives the "interesting" free identifiers in a Term, implemented as a
-- 'Fold'. Like 'termFreeVars'', but fold does not descend into the types of
-- variables and binders.
termFreeIds'
  :: (Contravariant f, Applicative f)
  => (forall b . Var b -> Bool)
  -- ^ Predicate telling whether an identifier is interesting
  -> (Var a -> f (Var a))
  -> Term
  -> f Term
termFreeIds' = termFreeVarsWorker False

-- | Worker for 'termFreeVars'' and 'termFreeIds''. Has option that selects
-- whether to descend into types or not.
termFreeVarsWorker
  :: forall a f
   . (Contravariant f, Applicative f)
  => Bool
  -- ^ Descend into the types of variables and binders?
  -> (forall b . Var b -> Bool)
  -- ^ Predicate telling whether a variable is interesting
  -> (Var a -> f (Var a))
  -> Term
  -> f Term
termFreeVarsWorker descendIntoTypes interesting f = go IntSet.empty where
  -- Fold over a type only when descending into types; otherwise leave it as-is.
  onType inScope ty
    | descendIntoTypes = typeFreeVars' interesting inScope f ty
    | otherwise        = pure ty

  go inLocalScope = \case
    Var v
      | descendIntoTypes -> v1 <* typeFreeVars' interesting inLocalScope1 f (varType v)
      | otherwise        -> v1
      where
        isInteresting = interesting v
        vInScope      = isLocalId v && varUniq v `IntSet.member` inLocalScope
        inLocalScope1
          | vInScope  = inLocalScope
          | otherwise = IntSet.empty -- See Note [Closing over type variables]

        v1 | isInteresting
           , not vInScope
           = Var . coerce <$> f (coerce v)
           | otherwise
           = pure (Var v)

    Lam id_ tm ->
      Lam <$> goBndr inLocalScope id_
          <*> go (IntSet.insert (varUniq id_) inLocalScope) tm
    TyLam tv tm ->
      TyLam <$> goBndr inLocalScope tv
            <*> go (IntSet.insert (varUniq tv) inLocalScope) tm

    App l r ->
      App <$> go inLocalScope l <*> go inLocalScope r

    TyApp l r ->
      TyApp <$> go inLocalScope l
            <*> onType inLocalScope r

    Let (NonRec i x) e ->
      Let <$> (NonRec <$> goBndr inLocalScope i <*> go inLocalScope x)
          <*> go (IntSet.insert (varUniq i) inLocalScope) e

    Let (Rec bs) e ->
      Let <$> (Rec <$> traverse (goBind inLocalScope') bs)
          <*> go inLocalScope' e
     where
      inLocalScope' = foldr (IntSet.insert . varUniq . fst) inLocalScope bs

    Case subj ty alts ->
      Case <$> go inLocalScope subj
           <*> onType inLocalScope ty
           <*> traverse (goAlt inLocalScope) alts

    Cast tm t1 t2 ->
      Cast <$> go inLocalScope tm
           <*> onType inLocalScope t1
           <*> onType inLocalScope t2

    Tick tick tm ->
      Tick <$> goTick inLocalScope tick
      <*> go inLocalScope tm

    tm -> pure tm

  -- Fold over a binder's type only when descending; otherwise leave it as-is.
  goBndr inLocalScope v
    | descendIntoTypes = (\t -> v {varType = t}) <$> onType inLocalScope (varType v)
    | otherwise        = pure v

  goBind inLocalScope (l,r) = (,) <$> goBndr inLocalScope l <*> go inLocalScope r

  goAlt inLocalScope (pat,alt) = case pat of
    DataPat dc tvs ids -> (,) <$> (DataPat dc
                                     <$> traverse (goBndr inLocalScope') tvs
                                     <*> traverse (goBndr inLocalScope') ids)
                              <*> go inLocalScope' alt
      where
        inLocalScope' = foldr IntSet.insert
                         (foldr IntSet.insert inLocalScope (map varUniq tvs))
                         (map varUniq ids)
    _ -> (,) pat <$> go inLocalScope alt

  goTick inLocalScope = \case
    NameMod m ty -> NameMod m <$> onType inLocalScope ty
    Attributes ty tm -> Attributes <$> onType inLocalScope ty
                                   <*> go inLocalScope tm
    tick         -> pure tick
{-# INLINE termFreeVarsWorker #-}

-- | Get the free variables of an expression and count the number of occurrences
countFreeOccurances
  :: Term
  -> VarEnv Int
countFreeOccurances =
  Lens.foldMapByOf freeLocalIds (unionVarEnvWith (+)) emptyVarEnv
                   (`unitVarEnv` (1 :: Int))