{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}

module Clash.Core.VarEnv
  ( -- * Environment with variables as keys
    VarEnv
    -- ** Accessors
    -- *** Size information
  , nullVarEnv
    -- ** Indexing
  , lookupVarEnv
  , lookupVarEnv'
  , lookupVarEnvDirectly
    -- ** Construction
  , emptyVarEnv
  , unitVarEnv
  , mkVarEnv
    -- ** Modification
  , extendVarEnv
  , extendVarEnvList
  , extendVarEnvWith
  , delVarEnv
  , delVarEnvList
  , unionVarEnv
  , unionVarEnvWith
  , differenceVarEnv
    -- ** Element-wise operations
    -- *** Mapping
  , mapVarEnv
  , mapMaybeVarEnv
    -- ** Folding
  , foldlWithUniqueVarEnv'
    -- ** Working with predicates
    -- *** Searching
  , elemVarEnv
  , notElemVarEnv
    -- ** Conversions
    -- *** Lists
  , eltsVarEnv
    -- * Sets of variables
  , VarSet
    -- ** Construction
  , emptyVarSet
  , unitVarSet
    -- ** Modification
  , delVarSetByKey
  , unionVarSet
  , differenceVarSet
    -- ** Working with predicates
  , nullVarSet
    -- *** Searching
  , elemVarSet
  , notElemVarSet
  , subsetVarSet
  , disjointVarSet
    -- ** Conversions
    -- *** Lists
  , mkVarSet
  , eltsVarSet
    -- * In-scope sets
  , InScopeSet
    -- ** Accessors
    -- *** Size information
  , emptyInScopeSet
    -- *** Indexing
  , lookupInScope
    -- ** Construction
  , mkInScopeSet
    -- ** Modification
  , extendInScopeSet
  , extendInScopeSetList
  , unionInScope
    -- ** Working with predicates
    -- *** Searching
  , elemInScopeSet
  , elemUniqInScopeSet
  , notElemInScopeSet
  , varSetInScope
    -- ** Unique generation
  , uniqAway
  , uniqAway'
    -- * Dual renaming
  , RnEnv
    -- ** Construction
  , mkRnEnv
    -- ** Renaming
  , rnTmBndr
  , rnTyBndr
  , rnTmBndrs
  , rnTyBndrs
  , rnOccLId
  , rnOccRId
  , rnOccLTy
  , rnOccRTy
  )
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           Control.DeepSeq           (NFData)
import           Data.Binary               (Binary)
import           Data.Coerce               (coerce)
import qualified Data.List                 as List
import qualified Data.List.Extra           as List
import           Data.Maybe                (fromMaybe)
#ifdef UNIQUE_IS_WORD64
import           Data.Word                 (Word64)
#endif

#if MIN_VERSION_prettyprinter(1,7,0)
import           Prettyprinter
#else
import           Data.Text.Prettyprint.Doc
#endif

import           GHC.Exts                  (Any)
import           GHC.Generics              (Generic)
import           GHC.Stack                 (HasCallStack)

import           Clash.Core.Pretty         ()
import           Clash.Core.Var
import           Clash.Data.UniqMap        (UniqMap)
import qualified Clash.Data.UniqMap as UniqMap
import           Clash.Debug               (debugIsOn)
import           Clash.Unique
import           Clash.Util
import           Clash.Pretty

-- * VarEnv

-- | Map indexed by variables
type VarEnv a = UniqMap a

-- | Empty map
emptyVarEnv
  :: VarEnv a
emptyVarEnv :: forall a. VarEnv a
emptyVarEnv = UniqMap a
forall a. VarEnv a
UniqMap.empty

-- | Environment containing a single variable-value pair
unitVarEnv
  :: Var b
  -> a
  -> VarEnv a
unitVarEnv :: forall b a. Var b -> a -> VarEnv a
unitVarEnv = Var b -> a -> UniqMap a
forall a b. Uniquable a => a -> b -> UniqMap b
UniqMap.singleton

-- | Look up a value based on the variable
lookupVarEnv
  :: Var b
  -> VarEnv a
  -> Maybe a
lookupVarEnv :: forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv = Var b -> UniqMap a -> Maybe a
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup

-- | Lookup a value based on the unique of a variable
lookupVarEnvDirectly
  :: Unique
  -> VarEnv a
  -> Maybe a
lookupVarEnvDirectly :: forall a. Unique -> VarEnv a -> Maybe a
lookupVarEnvDirectly = Unique -> UniqMap a -> Maybe a
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup

-- | Lookup a value based on the variable
--
-- Errors out when the variable is not present
lookupVarEnv'
  :: HasCallStack
  => VarEnv a
  -> Var b
  -> a
lookupVarEnv' :: forall a b. HasCallStack => VarEnv a -> Var b -> a
lookupVarEnv' = (Var b -> VarEnv a -> a) -> VarEnv a -> Var b -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Var b -> VarEnv a -> a
forall a b. Uniquable a => a -> UniqMap b -> b
UniqMap.find

-- | Remove a variable-value pair from the environment
delVarEnv
  :: VarEnv a
  -> Var b
  -> VarEnv a
delVarEnv :: forall a b. VarEnv a -> Var b -> VarEnv a
delVarEnv = (Var b -> VarEnv a -> VarEnv a) -> VarEnv a -> Var b -> VarEnv a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Var b -> VarEnv a -> VarEnv a
forall a b. Uniquable a => a -> UniqMap b -> UniqMap b
UniqMap.delete

-- | Remove a list of variable-value pairs from the environment
delVarEnvList
  :: VarEnv a
  -> [Var b]
  -> VarEnv a
delVarEnvList :: forall a b. VarEnv a -> [Var b] -> VarEnv a
delVarEnvList = ([Var b] -> VarEnv a -> VarEnv a)
-> VarEnv a -> [Var b] -> VarEnv a
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Var b] -> VarEnv a -> VarEnv a
forall a b. Uniquable a => [a] -> UniqMap b -> UniqMap b
UniqMap.deleteMany

-- | Add a variable-value pair to the environment; overwrites the value if the
-- variable already exists
extendVarEnv
  :: Var b
  -> a
  -> VarEnv a
  -> VarEnv a
extendVarEnv :: forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv = Var b -> a -> UniqMap a -> UniqMap a
forall a b. Uniquable a => a -> b -> UniqMap b -> UniqMap b
UniqMap.insert

-- | Add a variable-value pair to the environment; if the variable already
-- exists, the two values are merged with the given function
extendVarEnvWith
  :: Var b
  -> a
  -> (a -> a -> a)
  -> VarEnv a
  -> VarEnv a
extendVarEnvWith :: forall b a. Var b -> a -> (a -> a -> a) -> VarEnv a -> VarEnv a
extendVarEnvWith Var b
k a
v a -> a -> a
f =
  (a -> a -> a) -> Var b -> a -> UniqMap a -> UniqMap a
forall a b.
Uniquable a =>
(b -> b -> b) -> a -> b -> UniqMap b -> UniqMap b
UniqMap.insertWith a -> a -> a
f Var b
k a
v

-- | Add a list of variable-value pairs; the values of existing keys will be
-- overwritten
extendVarEnvList
  :: VarEnv a
  -> [(Var b, a)]
  -> VarEnv a
extendVarEnvList :: forall a b. VarEnv a -> [(Var b, a)] -> VarEnv a
extendVarEnvList = ([(Var b, a)] -> VarEnv a -> VarEnv a)
-> VarEnv a -> [(Var b, a)] -> VarEnv a
forall a b c. (a -> b -> c) -> b -> a -> c
flip [(Var b, a)] -> VarEnv a -> VarEnv a
forall a b. Uniquable a => [(a, b)] -> UniqMap b -> UniqMap b
UniqMap.insertMany

-- | Is the environment empty
nullVarEnv
  :: VarEnv a
  -> Bool
nullVarEnv :: forall a. VarEnv a -> Bool
nullVarEnv = UniqMap a -> Bool
forall a. VarEnv a -> Bool
UniqMap.null

-- | Get the (left-biased) union of two environments
unionVarEnv
  :: VarEnv a
  -> VarEnv a
  -> VarEnv a
unionVarEnv :: forall a. VarEnv a -> VarEnv a -> VarEnv a
unionVarEnv = VarEnv a -> VarEnv a -> VarEnv a
forall a. Semigroup a => a -> a -> a
(<>)

-- | Get the union of two environments, mapped values existing in both
-- environments will be merged with the given function.
unionVarEnvWith
  :: (a -> a -> a)
  -> VarEnv a
  -> VarEnv a
  -> VarEnv a
unionVarEnvWith :: forall a. (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
unionVarEnvWith = (a -> a -> a) -> UniqMap a -> UniqMap a -> UniqMap a
forall a. (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
UniqMap.unionWith

-- | Filter the first varenv to only contain keys which are not in the second varenv.
differenceVarEnv
  :: VarEnv a
  -> VarEnv a
  -> VarEnv a
differenceVarEnv :: forall a. VarEnv a -> VarEnv a -> VarEnv a
differenceVarEnv = UniqMap a -> UniqMap a -> UniqMap a
forall a. VarEnv a -> VarEnv a -> VarEnv a
UniqMap.difference

-- | Create an environment given a list of var-value pairs
mkVarEnv
  :: [(Var a,b)]
  -> VarEnv b
mkVarEnv :: forall a b. [(Var a, b)] -> VarEnv b
mkVarEnv = [(Var a, b)] -> UniqMap b
forall a b. Uniquable a => [(a, b)] -> UniqMap b
UniqMap.fromList

-- | Apply a function to every element in the environment
mapVarEnv
  :: (a -> b)
  -> VarEnv a
  -> VarEnv b
mapVarEnv :: forall a b. (a -> b) -> VarEnv a -> VarEnv b
mapVarEnv = (a -> b) -> UniqMap a -> UniqMap b
forall a b. (a -> b) -> VarEnv a -> VarEnv b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap

-- | Apply a function to every element in the environment; values for which the
-- function returns 'Nothing' are removed from the environment
mapMaybeVarEnv
  :: (a -> Maybe b)
  -> VarEnv a
  -> VarEnv b
mapMaybeVarEnv :: forall a b. (a -> Maybe b) -> VarEnv a -> VarEnv b
mapMaybeVarEnv = (a -> Maybe b) -> UniqMap a -> UniqMap b
forall a b. (a -> Maybe b) -> VarEnv a -> VarEnv b
UniqMap.mapMaybe

-- | Strict left-fold over an environment using both the unique of the
-- the variable and the value
foldlWithUniqueVarEnv'
  :: (a -> Unique -> b -> a)
  -> a
  -> VarEnv b
  -> a
foldlWithUniqueVarEnv' :: forall a b. (a -> Unique -> b -> a) -> a -> VarEnv b -> a
foldlWithUniqueVarEnv' = (a -> Unique -> b -> a) -> a -> UniqMap b -> a
forall a b. (a -> Unique -> b -> a) -> a -> VarEnv b -> a
UniqMap.foldlWithUnique'

-- | Extract the elements
eltsVarEnv
  :: VarEnv a
  -> [a]
eltsVarEnv :: forall a. VarEnv a -> [a]
eltsVarEnv = UniqMap a -> [a]
forall a. VarEnv a -> [a]
UniqMap.elems

-- | Does the variable exist in the environment
elemVarEnv
  :: Var a
  -> VarEnv b
  -> Bool
elemVarEnv :: forall a b. Var a -> VarEnv b -> Bool
elemVarEnv = Var a -> UniqMap b -> Bool
forall a b. Uniquable a => a -> UniqMap b -> Bool
UniqMap.elem

-- | Does the variable not exist in the environment
notElemVarEnv
  :: Var a
  -> VarEnv b
  -> Bool
notElemVarEnv :: forall a b. Var a -> VarEnv b -> Bool
notElemVarEnv = Var a -> UniqMap b -> Bool
forall a b. Uniquable a => a -> UniqMap b -> Bool
UniqMap.notElem

-- * VarSet

-- | Set of variables
type VarSet = UniqMap (Var Any)

-- | The empty set
emptyVarSet
  :: VarSet
emptyVarSet :: UniqMap (Var Any)
emptyVarSet = UniqMap (Var Any)
forall a. VarEnv a
UniqMap.empty

-- | The set of a single variable
unitVarSet
  :: Var a
  -> VarSet
unitVarSet :: forall a. Var a -> UniqMap (Var Any)
unitVarSet Var a
v = Var Any -> UniqMap (Var Any)
forall a. Uniquable a => a -> UniqMap a
UniqMap.singletonUnique (Var a -> Var Any
forall a b. Coercible a b => a -> b
coerce Var a
v)

-- | Add a variable to the set
extendVarSet
  :: VarSet
  -> Var a
  -> VarSet
extendVarSet :: forall a. UniqMap (Var Any) -> Var a -> UniqMap (Var Any)
extendVarSet UniqMap (Var Any)
env Var a
v = Var Any -> UniqMap (Var Any) -> UniqMap (Var Any)
forall a. Uniquable a => a -> UniqMap a -> UniqMap a
UniqMap.insertUnique (Var a -> Var Any
forall a b. Coercible a b => a -> b
coerce Var a
v) UniqMap (Var Any)
env

-- | Union two sets
unionVarSet
  :: VarSet
  -> VarSet
  -> VarSet
unionVarSet :: UniqMap (Var Any) -> UniqMap (Var Any) -> UniqMap (Var Any)
unionVarSet = UniqMap (Var Any) -> UniqMap (Var Any) -> UniqMap (Var Any)
forall a. Semigroup a => a -> a -> a
(<>)

-- | Take the difference of two sets
differenceVarSet
  :: VarSet
  -> VarSet
  -> VarSet
differenceVarSet :: UniqMap (Var Any) -> UniqMap (Var Any) -> UniqMap (Var Any)
differenceVarSet = UniqMap (Var Any) -> UniqMap (Var Any) -> UniqMap (Var Any)
forall a. VarEnv a -> VarEnv a -> VarEnv a
UniqMap.difference

-- | Is the variable an element in the set
elemVarSet
  :: Var a
  -> VarSet
  -> Bool
elemVarSet :: forall a. Var a -> UniqMap (Var Any) -> Bool
elemVarSet Var a
v = Unique -> UniqMap (Var Any) -> Bool
forall a b. Uniquable a => a -> UniqMap b -> Bool
UniqMap.elem (Var a -> Unique
forall a. Uniquable a => a -> Unique
getUnique Var a
v)

-- | Is the variable not an element in the set
notElemVarSet
  :: Var a
  -> VarSet
  -> Bool
notElemVarSet :: forall a. Var a -> UniqMap (Var Any) -> Bool
notElemVarSet Var a
v = Unique -> UniqMap (Var Any) -> Bool
forall a b. Uniquable a => a -> UniqMap b -> Bool
UniqMap.notElem (Var a -> Unique
forall a. Uniquable a => a -> Unique
getUnique Var a
v)

-- | Is the set of variables A a subset of the variables B
subsetVarSet
  :: VarSet
  -- ^ Set of variables A
  -> VarSet
  -- ^ Set of variables B
  -> Bool
subsetVarSet :: UniqMap (Var Any) -> UniqMap (Var Any) -> Bool
subsetVarSet = UniqMap (Var Any) -> UniqMap (Var Any) -> Bool
forall b. UniqMap b -> UniqMap b -> Bool
UniqMap.submap

-- | Are the sets of variables disjoint
disjointVarSet
  :: VarSet
  -> VarSet
  -> Bool
disjointVarSet :: UniqMap (Var Any) -> UniqMap (Var Any) -> Bool
disjointVarSet = UniqMap (Var Any) -> UniqMap (Var Any) -> Bool
forall b. UniqMap b -> UniqMap b -> Bool
UniqMap.disjoint

-- | Check whether a varset is empty
nullVarSet
  :: VarSet
  -> Bool
nullVarSet :: UniqMap (Var Any) -> Bool
nullVarSet = UniqMap (Var Any) -> Bool
forall a. VarEnv a -> Bool
UniqMap.null

-- | Look up a variable in the set, returns it if it exists
lookupVarSet
  :: Var a
  -> VarSet
  -> Maybe (Var Any)
lookupVarSet :: forall a. Var a -> UniqMap (Var Any) -> Maybe (Var Any)
lookupVarSet = Var a -> UniqMap (Var Any) -> Maybe (Var Any)
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
UniqMap.lookup

-- | Remove a variable from the set based on its 'Unique'
delVarSetByKey
  :: Unique
  -> VarSet
  -> VarSet
delVarSetByKey :: Unique -> UniqMap (Var Any) -> UniqMap (Var Any)
delVarSetByKey = Unique -> UniqMap (Var Any) -> UniqMap (Var Any)
forall a b. Uniquable a => a -> UniqMap b -> UniqMap b
UniqMap.delete

-- | Create a set from a list of variables
mkVarSet
  :: [Var a]
  -> VarSet
mkVarSet :: forall a. [Var a] -> UniqMap (Var Any)
mkVarSet [Var a]
xs = [(Unique, Var Any)] -> UniqMap (Var Any)
forall a b. Uniquable a => [(a, b)] -> UniqMap b
UniqMap.fromList ([(Unique, Var Any)] -> UniqMap (Var Any))
-> [(Unique, Var Any)] -> UniqMap (Var Any)
forall a b. (a -> b) -> a -> b
$ (Var a -> (Unique, Var Any)) -> [Var a] -> [(Unique, Var Any)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Var a
x -> (Var a -> Unique
forall a. Uniquable a => a -> Unique
getUnique Var a
x, Var a -> Var Any
forall a b. Coercible a b => a -> b
coerce Var a
x)) [Var a]
xs

eltsVarSet
  :: VarSet
  -> [Var Any]
eltsVarSet :: UniqMap (Var Any) -> [Var Any]
eltsVarSet = UniqMap (Var Any) -> [Var Any]
forall a. VarEnv a -> [a]
UniqMap.elems

-- * InScopeSet

type Seed
#ifdef UNIQUE_IS_WORD64
  = Word64
#else
  = Int
#endif

-- | Set of variables that is in scope at some point
--
-- The 'Seed' is a kind of hash-value used to generate new uniques. It should
-- never be zero
--
-- See "Secrets of the Glasgow Haskell Compiler inliner" Section 3.2 for the
-- motivation
data InScopeSet = InScopeSet VarSet {-# UNPACK #-} !Seed
  deriving ((forall x. InScopeSet -> Rep InScopeSet x)
-> (forall x. Rep InScopeSet x -> InScopeSet) -> Generic InScopeSet
forall x. Rep InScopeSet x -> InScopeSet
forall x. InScopeSet -> Rep InScopeSet x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. InScopeSet -> Rep InScopeSet x
from :: forall x. InScopeSet -> Rep InScopeSet x
$cto :: forall x. Rep InScopeSet x -> InScopeSet
to :: forall x. Rep InScopeSet x -> InScopeSet
Generic, InScopeSet -> ()
(InScopeSet -> ()) -> NFData InScopeSet
forall a. (a -> ()) -> NFData a
$crnf :: InScopeSet -> ()
rnf :: InScopeSet -> ()
NFData, Get InScopeSet
[InScopeSet] -> Put
InScopeSet -> Put
(InScopeSet -> Put)
-> Get InScopeSet -> ([InScopeSet] -> Put) -> Binary InScopeSet
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
$cput :: InScopeSet -> Put
put :: InScopeSet -> Put
$cget :: Get InScopeSet
get :: Get InScopeSet
$cputList :: [InScopeSet] -> Put
putList :: [InScopeSet] -> Put
Binary)

instance ClashPretty InScopeSet where
  clashPretty :: InScopeSet -> Doc ()
clashPretty (InScopeSet UniqMap (Var Any)
s Unique
_) = UniqMap (Var Any) -> Doc ()
forall a. ClashPretty a => a -> Doc ()
clashPretty UniqMap (Var Any)
s

-- | The empty set
extendInScopeSet
  :: InScopeSet
  -> Var a
  -> InScopeSet
extendInScopeSet :: forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet (InScopeSet UniqMap (Var Any)
inScope Unique
n) Var a
v =
  UniqMap (Var Any) -> Unique -> InScopeSet
InScopeSet (UniqMap (Var Any) -> Var a -> UniqMap (Var Any)
forall a. UniqMap (Var Any) -> Var a -> UniqMap (Var Any)
extendVarSet UniqMap (Var Any)
inScope Var a
v) (Unique
n Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
+ Unique
1)

-- | Add a list of variables in scope
extendInScopeSetList
  :: InScopeSet
  -> [Var a]
  -> InScopeSet
extendInScopeSetList :: forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList (InScopeSet UniqMap (Var Any)
inScope Unique
n) [Var a]
vs =
  UniqMap (Var Any) -> Unique -> InScopeSet
InScopeSet ((UniqMap (Var Any) -> Var a -> UniqMap (Var Any))
-> UniqMap (Var Any) -> [Var a] -> UniqMap (Var Any)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' UniqMap (Var Any) -> Var a -> UniqMap (Var Any)
forall a. UniqMap (Var Any) -> Var a -> UniqMap (Var Any)
extendVarSet UniqMap (Var Any)
inScope [Var a]
vs) (Unique
n Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
+ Int -> Unique
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Var a] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Var a]
vs))

-- | Union two sets of in scope variables
unionInScope
  :: InScopeSet
  -> InScopeSet
  -> InScopeSet
unionInScope :: InScopeSet -> InScopeSet -> InScopeSet
unionInScope (InScopeSet UniqMap (Var Any)
s1 Unique
_) (InScopeSet UniqMap (Var Any)
s2 Unique
n2)
  = UniqMap (Var Any) -> Unique -> InScopeSet
InScopeSet (UniqMap (Var Any)
s1 UniqMap (Var Any) -> UniqMap (Var Any) -> UniqMap (Var Any)
`unionVarSet` UniqMap (Var Any)
s2) Unique
n2

-- | Is the set of variables in scope
varSetInScope
  :: VarSet
  -> InScopeSet
  -> Bool
varSetInScope :: UniqMap (Var Any) -> InScopeSet -> Bool
varSetInScope UniqMap (Var Any)
vars (InScopeSet UniqMap (Var Any)
s1 Unique
_)
  = UniqMap (Var Any)
vars UniqMap (Var Any) -> UniqMap (Var Any) -> Bool
`subsetVarSet` UniqMap (Var Any)
s1

-- | Look up a variable in the 'InScopeSet'. This gives you the canonical
-- version of the variable
lookupInScope
  :: InScopeSet
  -> Var a
  -> Maybe (Var Any)
lookupInScope :: forall a. InScopeSet -> Var a -> Maybe (Var Any)
lookupInScope (InScopeSet UniqMap (Var Any)
s Unique
_) Var a
v = Var a -> UniqMap (Var Any) -> Maybe (Var Any)
forall a. Var a -> UniqMap (Var Any) -> Maybe (Var Any)
lookupVarSet Var a
v UniqMap (Var Any)
s

-- | Is the variable in scope
elemInScopeSet
  :: Var a
  -> InScopeSet
  -> Bool
elemInScopeSet :: forall a. Var a -> InScopeSet -> Bool
elemInScopeSet Var a
v (InScopeSet UniqMap (Var Any)
s Unique
_) = Var a -> UniqMap (Var Any) -> Bool
forall a. Var a -> UniqMap (Var Any) -> Bool
elemVarSet Var a
v UniqMap (Var Any)
s

-- | Check whether an element exists in the set based on the `Unique` contained
-- in that element
elemUniqInScopeSet
  :: Unique
  -> InScopeSet
  -> Bool
elemUniqInScopeSet :: Unique -> InScopeSet -> Bool
elemUniqInScopeSet Unique
u (InScopeSet UniqMap (Var Any)
s Unique
_) = Unique -> UniqMap (Var Any) -> Bool
forall a b. Uniquable a => a -> UniqMap b -> Bool
UniqMap.elem Unique
u UniqMap (Var Any)
s

-- | Is the variable not in scope
notElemInScopeSet
  :: Var a
  -> InScopeSet
  -> Bool
notElemInScopeSet :: forall a. Var a -> InScopeSet -> Bool
notElemInScopeSet Var a
v (InScopeSet UniqMap (Var Any)
s Unique
_) = Var a -> UniqMap (Var Any) -> Bool
forall a. Var a -> UniqMap (Var Any) -> Bool
notElemVarSet Var a
v UniqMap (Var Any)
s

-- | Create a set of variables in scope
mkInScopeSet
  :: VarSet
  -> InScopeSet
mkInScopeSet :: UniqMap (Var Any) -> InScopeSet
mkInScopeSet UniqMap (Var Any)
is = UniqMap (Var Any) -> Unique -> InScopeSet
InScopeSet UniqMap (Var Any)
is Unique
1

-- | The empty set
emptyInScopeSet
  :: InScopeSet
emptyInScopeSet :: InScopeSet
emptyInScopeSet = UniqMap (Var Any) -> InScopeSet
mkInScopeSet UniqMap (Var Any)
emptyVarSet

-- | Ensure that the 'Unique' of a variable does not occur in the 'InScopeSet'
uniqAway
  :: (Uniquable a, ClashPretty a)
  => InScopeSet
  -> a
  -> a
uniqAway :: forall a. (Uniquable a, ClashPretty a) => InScopeSet -> a -> a
uniqAway (InScopeSet UniqMap (Var Any)
set Unique
n) a
a =
  (Unique -> Bool) -> Unique -> a -> a
forall a.
(Uniquable a, ClashPretty a) =>
(Unique -> Bool) -> Unique -> a -> a
uniqAway' (Unique -> UniqMap (Var Any) -> Bool
forall a b. Uniquable a => a -> UniqMap b -> Bool
`UniqMap.elem` UniqMap (Var Any)
set) Unique
n a
a

uniqAway'
  :: (Uniquable a, ClashPretty a)
  => (Unique -> Bool)
  -- ^ Unique in scope test
  -> Seed
  -- ^ Seed
  -> a
  -> a
uniqAway' :: forall a.
(Uniquable a, ClashPretty a) =>
(Unique -> Bool) -> Unique -> a -> a
uniqAway' Unique -> Bool
inScopeTest Unique
n a
u =
  if Unique -> Bool
inScopeTest (a -> Unique
forall a. Uniquable a => a -> Unique
getUnique a
u) then
    Unique -> a
try Unique
1
  else
    a
u
 where
  origUniq :: Unique
origUniq = a -> Unique
forall a. Uniquable a => a -> Unique
getUnique a
u
  try :: Unique -> a
try Unique
k
    | Bool
debugIsOn Bool -> Bool -> Bool
&& Unique
k Unique -> Unique -> Bool
forall a. Ord a => a -> a -> Bool
> Unique
1000
    = String -> Doc () -> a
forall ann a. String -> Doc ann -> a
pprPanic String
"uniqAway loop:" Doc ()
msg
    | Unique -> Bool
inScopeTest Unique
uniq
    = Unique -> a
try (Unique
k Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
+ Unique
1)
    | Unique
k Unique -> Unique -> Bool
forall a. Ord a => a -> a -> Bool
> Unique
3
    = String -> Doc () -> a -> a
forall ann a. String -> Doc ann -> a -> a
pprTraceDebug String
"uniqAway:" Doc ()
msg (a -> Unique -> a
forall a. Uniquable a => a -> Unique -> a
setUnique a
u Unique
uniq)
    | Bool
otherwise
    = a -> Unique -> a
forall a. Uniquable a => a -> Unique -> a
setUnique a
u Unique
uniq
    where
      msg :: Doc ()
msg  = Unique -> Doc ()
forall a. Pretty a => a -> Doc ()
fromPretty Unique
k Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
"tries" Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc ()
forall a. ClashPretty a => a -> Doc ()
clashPretty a
u Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Unique -> Doc ()
forall a. Pretty a => a -> Doc ()
fromPretty Unique
n
      uniq :: Unique
uniq = Unique -> Unique -> Unique
deriveUnique Unique
origUniq (Unique
n Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
* Unique
k)

deriveUnique
  :: Unique
  -> Seed
  -> Unique
deriveUnique :: Unique -> Unique -> Unique
deriveUnique Unique
i Unique
delta = Unique
i Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
+ Unique
delta

-- * RnEnv

-- | Rename environment for e.g. alpha equivalence
--
-- When going under binders for e.g.
--
-- @
-- \x -> e1  `aeq` \y -> e2
-- @
--
-- We want to rename @[x -> y]@  or @[y -> x]@, but we have to pick a binder
-- that is neither free in @e1@ nor @e2@ or we risk accidental capture.
--
-- So we must maintain:
--
--   1. A renaming for the left term
--
--   2. A renaming for the right term
--
--   3. A set of in scope variables
data RnEnv
  = RnEnv
  { RnEnv -> VarEnv TyVar
rn_envLTy  :: VarEnv TyVar
    -- ^ Type renaming for the left term
  , RnEnv -> VarEnv Id
rn_envLTm  :: VarEnv Id
    -- ^ Term renaming for the left term
  , RnEnv -> VarEnv TyVar
rn_envRTy  :: VarEnv TyVar
    -- ^ Type renaming for the right term
  , RnEnv -> VarEnv Id
rn_envRTm  :: VarEnv Id
    -- ^ Term renaming for the right term
  , RnEnv -> InScopeSet
rn_inScope :: InScopeSet
    -- ^ In scope in left or right terms
  }

-- | Create an empty renaming environment
mkRnEnv
  :: InScopeSet -> RnEnv
mkRnEnv :: InScopeSet -> RnEnv
mkRnEnv InScopeSet
vars
  = RnEnv
  { rn_envLTy :: VarEnv TyVar
rn_envLTy  = VarEnv TyVar
forall a. VarEnv a
emptyVarEnv
  , rn_envLTm :: VarEnv Id
rn_envLTm  = VarEnv Id
forall a. VarEnv a
emptyVarEnv
  , rn_envRTy :: VarEnv TyVar
rn_envRTy  = VarEnv TyVar
forall a. VarEnv a
emptyVarEnv
  , rn_envRTm :: VarEnv Id
rn_envRTm  = VarEnv Id
forall a. VarEnv a
emptyVarEnv
  , rn_inScope :: InScopeSet
rn_inScope = InScopeSet
vars
  }

-- | Look up the renaming of an type-variable occurrence in the left term
rnOccLTy
  :: RnEnv -> TyVar -> TyVar
rnOccLTy :: RnEnv -> TyVar -> TyVar
rnOccLTy RnEnv
rn TyVar
v = TyVar -> Maybe TyVar -> TyVar
forall a. a -> Maybe a -> a
fromMaybe TyVar
v (TyVar -> VarEnv TyVar -> Maybe TyVar
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv TyVar
v (RnEnv -> VarEnv TyVar
rn_envLTy RnEnv
rn))

-- | Look up the renaming of an type-variable occurrence in the right term
rnOccRTy
  :: RnEnv -> TyVar -> TyVar
rnOccRTy :: RnEnv -> TyVar -> TyVar
rnOccRTy RnEnv
rn TyVar
v = TyVar -> Maybe TyVar -> TyVar
forall a. a -> Maybe a -> a
fromMaybe TyVar
v (TyVar -> VarEnv TyVar -> Maybe TyVar
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv TyVar
v (RnEnv -> VarEnv TyVar
rn_envRTy RnEnv
rn))

-- | Simultaneously go under the type-variable binder /bTvL/ and type-variable
-- binder /bTvR/, finds a new binder /newTvB/, and return an environment mapping
-- @[bTvL -> newB]@ and @[bTvR -> newB]@
rnTyBndr
  :: RnEnv -> TyVar -> TyVar -> RnEnv
rnTyBndr :: RnEnv -> TyVar -> TyVar -> RnEnv
rnTyBndr rv :: RnEnv
rv@(RnEnv {rn_envLTy :: RnEnv -> VarEnv TyVar
rn_envLTy = VarEnv TyVar
lenv, rn_envRTy :: RnEnv -> VarEnv TyVar
rn_envRTy = VarEnv TyVar
renv, rn_inScope :: RnEnv -> InScopeSet
rn_inScope = InScopeSet
inScope}) TyVar
bL TyVar
bR =
  RnEnv
rv { rn_envLTy = extendVarEnv bL newB lenv -- See Note [Rebinding and shadowing]
     , rn_envRTy = extendVarEnv bR newB renv
     , rn_inScope = extendInScopeSet inScope newB }
 where
  -- Find a new type-binder not in scope in either term
  newB :: TyVar
newB | Bool -> Bool
not (TyVar
bL TyVar -> InScopeSet -> Bool
forall a. Var a -> InScopeSet -> Bool
`elemInScopeSet` InScopeSet
inScope) = TyVar
bL
       | Bool -> Bool
not (TyVar
bR TyVar -> InScopeSet -> Bool
forall a. Var a -> InScopeSet -> Bool
`elemInScopeSet` InScopeSet
inScope) = TyVar
bR
       | Bool
otherwise                         = InScopeSet -> TyVar -> TyVar
forall a. (Uniquable a, ClashPretty a) => InScopeSet -> a -> a
uniqAway InScopeSet
inScope TyVar
bL

{- Note [Rebinding and shadowing]
Imagine:

@
\x -> \x -> e1  `aeq` \y -> \x -> e2
@

Then inside

@
\x \y  { [x->p] [y->p]  {p} }
\x \z  { [x->q] [y->p, z->q] {p,q} }
@

i.e. if the new var is the same as the old var, the renaming is deleted by
'extendVarEnv'
-}

-- | Applies 'rnTyBndr' to several variables: the two variable lists must be of
-- equal length.
rnTyBndrs
  :: RnEnv -> [TyVar] -> [TyVar] -> RnEnv
rnTyBndrs :: RnEnv -> [TyVar] -> [TyVar] -> RnEnv
rnTyBndrs RnEnv
env [TyVar]
tvs1 [TyVar]
tvs2 =
  (RnEnv -> (TyVar, TyVar) -> RnEnv)
-> RnEnv -> [(TyVar, TyVar)] -> RnEnv
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\RnEnv
s (TyVar
l,TyVar
r) -> RnEnv -> TyVar -> TyVar -> RnEnv
rnTyBndr RnEnv
s TyVar
l TyVar
r) RnEnv
env ([TyVar] -> [TyVar] -> [(TyVar, TyVar)]
forall a b. HasCallStack => [a] -> [b] -> [(a, b)]
List.zipEqual [TyVar]
tvs1 [TyVar]
tvs2)

-- | Look up the renaming of an occurrence in the left term
rnOccLId
  :: RnEnv -> Id -> Id
rnOccLId :: RnEnv -> Id -> Id
rnOccLId RnEnv
rn Id
v = Id -> Maybe Id -> Id
forall a. a -> Maybe a -> a
fromMaybe Id
v (Id -> VarEnv Id -> Maybe Id
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
v (RnEnv -> VarEnv Id
rn_envLTm RnEnv
rn))

-- | Look up the renaming of an occurrence in the left term
rnOccRId
  :: RnEnv -> Id -> Id
rnOccRId :: RnEnv -> Id -> Id
rnOccRId RnEnv
rn Id
v = Id -> Maybe Id -> Id
forall a. a -> Maybe a -> a
fromMaybe Id
v (Id -> VarEnv Id -> Maybe Id
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
v (RnEnv -> VarEnv Id
rn_envRTm RnEnv
rn))

-- | Simultaneously go under the binder /bL/ and binder /bR/, finds a new binder
-- /newTvB/, and return an environment mapping @[bL -> newB]@ and @[bR -> newB]@
rnTmBndr
  :: RnEnv -> Id -> Id -> RnEnv
rnTmBndr :: RnEnv -> Id -> Id -> RnEnv
rnTmBndr rv :: RnEnv
rv@(RnEnv {rn_envLTm :: RnEnv -> VarEnv Id
rn_envLTm = VarEnv Id
lenv, rn_envRTm :: RnEnv -> VarEnv Id
rn_envRTm = VarEnv Id
renv, rn_inScope :: RnEnv -> InScopeSet
rn_inScope = InScopeSet
inScope}) Id
bL Id
bR =
  RnEnv
rv { rn_envLTm = extendVarEnv bL newB lenv -- See Note [Rebinding and shadowing]
     , rn_envRTm = extendVarEnv bR newB renv
     , rn_inScope = extendInScopeSet inScope newB }
 where
  -- Find a new type-binder not in scope in either term
  newB :: Id
newB | Bool -> Bool
not (Id
bL Id -> InScopeSet -> Bool
forall a. Var a -> InScopeSet -> Bool
`elemInScopeSet` InScopeSet
inScope) = Id
bL
       | Bool -> Bool
not (Id
bR Id -> InScopeSet -> Bool
forall a. Var a -> InScopeSet -> Bool
`elemInScopeSet` InScopeSet
inScope) = Id
bR
       | Bool
otherwise                         = InScopeSet -> Id -> Id
forall a. (Uniquable a, ClashPretty a) => InScopeSet -> a -> a
uniqAway InScopeSet
inScope Id
bL

-- | Applies 'rnTmBndr' to several variables: the two variable lists must be of
-- equal length.
rnTmBndrs
  :: RnEnv -> [Id] -> [Id] -> RnEnv
rnTmBndrs :: RnEnv -> [Id] -> [Id] -> RnEnv
rnTmBndrs RnEnv
env [Id]
ids1 [Id]
ids2 =
  (RnEnv -> (Id, Id) -> RnEnv) -> RnEnv -> [(Id, Id)] -> RnEnv
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\RnEnv
s (Id
l,Id
r) -> RnEnv -> Id -> Id -> RnEnv
rnTmBndr RnEnv
s Id
l Id
r) RnEnv
env ([Id] -> [Id] -> [(Id, Id)]
forall a b. HasCallStack => [a] -> [b] -> [(a, b)]
List.zipEqual [Id]
ids1 [Id]
ids2)