{-# 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
    -- ** 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

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)

#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.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 :: VarEnv a
emptyVarEnv = VarEnv a
forall a. UniqMap a
emptyUniqMap

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

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

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

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

-- | Remove a variable-value pair from the environment
delVarEnv
  :: VarEnv a
  -> Var b
  -> VarEnv a
delVarEnv :: VarEnv a -> Var b -> VarEnv a
delVarEnv = VarEnv a -> Var b -> VarEnv a
forall a b. Uniquable a => UniqMap b -> a -> UniqMap b
delUniqMap

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

-- | 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 :: Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv = Var b -> a -> VarEnv a -> VarEnv a
forall a b. Uniquable a => a -> b -> UniqMap b -> UniqMap b
extendUniqMap

-- | 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 :: Var b -> a -> (a -> a -> a) -> VarEnv a -> VarEnv a
extendVarEnvWith = Var b -> a -> (a -> a -> a) -> VarEnv a -> VarEnv a
forall a b.
Uniquable a =>
a -> b -> (b -> b -> b) -> UniqMap b -> UniqMap b
extendUniqMapWith

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

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

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

-- | 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 :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
unionVarEnvWith = (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
forall a. (a -> a -> a) -> UniqMap a -> UniqMap a -> UniqMap a
unionUniqMapWith

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

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

-- | 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 :: (a -> Maybe b) -> VarEnv a -> VarEnv b
mapMaybeVarEnv = (a -> Maybe b) -> VarEnv a -> VarEnv b
forall a b. (a -> Maybe b) -> UniqMap a -> UniqMap b
mapMaybeUniqMap

-- | 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' :: (a -> Unique -> b -> a) -> a -> VarEnv b -> a
foldlWithUniqueVarEnv' = (a -> Unique -> b -> a) -> a -> VarEnv b -> a
forall a b. (a -> Unique -> b -> a) -> a -> UniqMap b -> a
foldlWithUnique'

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

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

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

-- * VarSet

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

-- | The empty set
emptyVarSet
  :: VarSet
emptyVarSet :: VarSet
emptyVarSet = VarSet
forall a. UniqSet a
emptyUniqSet

-- | The set of a single variable
unitVarSet
  :: Var a
  -> VarSet
unitVarSet :: Var a -> VarSet
unitVarSet Var a
v = Var Any -> VarSet
forall a. Uniquable a => a -> UniqSet a
unitUniqSet (Var a -> Var Any
coerce Var a
v)

-- | Add a variable to the set
extendVarSet
  :: VarSet
  -> Var a
  -> VarSet
extendVarSet :: VarSet -> Var a -> VarSet
extendVarSet VarSet
env Var a
v = VarSet -> Var Any -> VarSet
forall a. Uniquable a => UniqSet a -> a -> UniqSet a
extendUniqSet VarSet
env (Var a -> Var Any
coerce Var a
v)

-- | Union two sets
unionVarSet
  :: VarSet
  -> VarSet
  -> VarSet
unionVarSet :: VarSet -> VarSet -> VarSet
unionVarSet = VarSet -> VarSet -> VarSet
forall a. UniqSet a -> UniqSet a -> UniqSet a
unionUniqSet

-- | Take the difference of two sets
differenceVarSet
  :: VarSet
  -> VarSet
  -> VarSet
differenceVarSet :: VarSet -> VarSet -> VarSet
differenceVarSet = VarSet -> VarSet -> VarSet
forall a. UniqSet a -> UniqSet a -> UniqSet a
differenceUniqSet

-- | Is the variable an element in the set
elemVarSet
  :: Var a
  -> VarSet
  -> Bool
elemVarSet :: Var a -> VarSet -> Bool
elemVarSet Var a
v = Var Any -> VarSet -> Bool
forall a. Uniquable a => a -> UniqSet a -> Bool
elemUniqSet (Var a -> Var Any
coerce Var a
v)

-- | Is the variable not an element in the set
notElemVarSet
  :: Var a
  -> VarSet
  -> Bool
notElemVarSet :: Var a -> VarSet -> Bool
notElemVarSet Var a
v = Var Any -> VarSet -> Bool
forall a. Uniquable a => a -> UniqSet a -> Bool
notElemUniqSet (Var a -> Var Any
coerce 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 :: VarSet -> VarSet -> Bool
subsetVarSet = VarSet -> VarSet -> Bool
forall a. UniqSet a -> UniqSet a -> Bool
subsetUniqSet

-- | Are the sets of variables disjoint
disjointVarSet
  :: VarSet
  -> VarSet
  -> Bool
disjointVarSet :: VarSet -> VarSet -> Bool
disjointVarSet = VarSet -> VarSet -> Bool
forall a. UniqSet a -> UniqSet a -> Bool
disjointUniqSet

-- | Check whether a varset is empty
nullVarSet
  :: VarSet
  -> Bool
nullVarSet :: VarSet -> Bool
nullVarSet = VarSet -> Bool
forall a. UniqSet a -> Bool
nullUniqSet

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

-- | Remove a variable from the set based on its 'Unique'
delVarSetByKey
  :: Unique
  -> VarSet
  -> VarSet
delVarSetByKey :: Unique -> VarSet -> VarSet
delVarSetByKey = Unique -> VarSet -> VarSet
forall b. Unique -> UniqSet b -> UniqSet b
delUniqSetDirectly

-- | Create a set from a list of variables
mkVarSet
  :: [Var a]
  -> VarSet
mkVarSet :: [Var a] -> VarSet
mkVarSet [Var a]
xs = [Var Any] -> VarSet
forall a. Uniquable a => [a] -> UniqSet a
mkUniqSet ([Var a] -> [Var Any]
coerce [Var a]
xs)

eltsVarSet
  :: VarSet
  -> [Var Any]
eltsVarSet :: VarSet -> [Var Any]
eltsVarSet = VarSet -> [Var Any]
forall a. UniqSet a -> [a]
eltsUniqSet

-- * InScopeSet

-- | Set of variables that is in scope at some point
--
-- The 'Int' 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 #-} !Int
  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
$cto :: forall x. Rep InScopeSet x -> InScopeSet
$cfrom :: forall x. InScopeSet -> Rep InScopeSet x
Generic, InScopeSet -> ()
(InScopeSet -> ()) -> NFData InScopeSet
forall a. (a -> ()) -> NFData a
rnf :: InScopeSet -> ()
$crnf :: 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
putList :: [InScopeSet] -> Put
$cputList :: [InScopeSet] -> Put
get :: Get InScopeSet
$cget :: Get InScopeSet
put :: InScopeSet -> Put
$cput :: InScopeSet -> Put
Binary)

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

-- | The empty set
extendInScopeSet
  :: InScopeSet
  -> Var a
  -> InScopeSet
extendInScopeSet :: InScopeSet -> Var a -> InScopeSet
extendInScopeSet (InScopeSet VarSet
inScope Unique
n) Var a
v =
  VarSet -> Unique -> InScopeSet
InScopeSet (VarSet -> Var a -> VarSet
forall a. VarSet -> Var a -> VarSet
extendVarSet VarSet
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 :: InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList (InScopeSet VarSet
inScope Unique
n) [Var a]
vs =
  VarSet -> Unique -> InScopeSet
InScopeSet ((VarSet -> Var a -> VarSet) -> VarSet -> [Var a] -> VarSet
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' VarSet -> Var a -> VarSet
forall a. VarSet -> Var a -> VarSet
extendVarSet VarSet
inScope [Var a]
vs) (Unique
n Unique -> Unique -> Unique
forall a. Num a => a -> a -> a
+ [Var a] -> Unique
forall (t :: Type -> Type) a. Foldable t => t a -> Unique
length [Var a]
vs)

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

-- | Is the set of variables in scope
varSetInScope
  :: VarSet
  -> InScopeSet
  -> Bool
varSetInScope :: VarSet -> InScopeSet -> Bool
varSetInScope VarSet
vars (InScopeSet VarSet
s1 Unique
_)
  = VarSet
vars VarSet -> VarSet -> Bool
`subsetVarSet` VarSet
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 :: InScopeSet -> Var a -> Maybe (Var Any)
lookupInScope (InScopeSet VarSet
s Unique
_) Var a
v = Var a -> VarSet -> Maybe (Var Any)
forall a. Var a -> VarSet -> Maybe (Var Any)
lookupVarSet Var a
v VarSet
s

-- | Is the variable in scope
elemInScopeSet
  :: Var a
  -> InScopeSet
  -> Bool
elemInScopeSet :: Var a -> InScopeSet -> Bool
elemInScopeSet Var a
v (InScopeSet VarSet
s Unique
_) = Var a -> VarSet -> Bool
forall a. Var a -> VarSet -> Bool
elemVarSet Var a
v VarSet
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 VarSet
s Unique
_) = Unique
u Unique -> VarSet -> Bool
forall a. Unique -> UniqSet a -> Bool
`elemUniqSetDirectly` VarSet
s

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

-- | Create a set of variables in scope
mkInScopeSet
  :: VarSet
  -> InScopeSet
mkInScopeSet :: VarSet -> InScopeSet
mkInScopeSet VarSet
is = VarSet -> Unique -> InScopeSet
InScopeSet VarSet
is Unique
1

-- | The empty set
emptyInScopeSet
  :: InScopeSet
emptyInScopeSet :: InScopeSet
emptyInScopeSet = VarSet -> InScopeSet
mkInScopeSet VarSet
emptyVarSet

-- | Ensure that the 'Unique' of a variable does not occur in the 'InScopeSet'
uniqAway
  :: (Uniquable a, ClashPretty a)
  => InScopeSet
  -> a
  -> a
uniqAway :: InScopeSet -> a -> a
uniqAway (InScopeSet VarSet
set Unique
n) a
a =
  (Unique -> Bool) -> Unique -> a -> a
forall a.
(Uniquable a, ClashPretty a) =>
(Unique -> Bool) -> Unique -> a -> a
uniqAway' (Unique -> VarSet -> Bool
forall a. Unique -> UniqSet a -> Bool
`elemUniqSetDirectly` VarSet
set) Unique
n a
a

uniqAway'
  :: (Uniquable a, ClashPretty a)
  => (Unique -> Bool)
  -- ^ Unique in scope test
  -> Int
  -- ^ Seed
  -> a
  -> a
uniqAway' :: (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
  -> Int
  -> 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 :: VarEnv TyVar
-> VarEnv Id -> VarEnv TyVar -> VarEnv Id -> InScopeSet -> RnEnv
RnEnv
  { rn_envLTy :: VarEnv TyVar
rn_envLTy  = VarEnv TyVar
forall a. UniqMap a
emptyVarEnv
  , rn_envLTm :: VarEnv Id
rn_envLTm  = VarEnv Id
forall a. UniqMap a
emptyVarEnv
  , rn_envRTy :: VarEnv TyVar
rn_envRTy  = VarEnv TyVar
forall a. UniqMap a
emptyVarEnv
  , rn_envRTm :: VarEnv Id
rn_envRTm  = VarEnv Id
forall a. UniqMap 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 :: VarEnv TyVar
rn_envLTy = TyVar -> TyVar -> VarEnv TyVar -> VarEnv TyVar
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
bL TyVar
newB VarEnv TyVar
lenv -- See Note [Rebinding and shadowing]
     , rn_envRTy :: VarEnv TyVar
rn_envRTy = TyVar -> TyVar -> VarEnv TyVar -> VarEnv TyVar
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
bR TyVar
newB VarEnv TyVar
renv
     , rn_inScope :: InScopeSet
rn_inScope = InScopeSet -> TyVar -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
inScope TyVar
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 (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 :: VarEnv Id
rn_envLTm = Id -> Id -> VarEnv Id -> VarEnv Id
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Id
bL Id
newB VarEnv Id
lenv -- See Note [Rebinding and shadowing]
     , rn_envRTm :: VarEnv Id
rn_envRTm = Id -> Id -> VarEnv Id -> VarEnv Id
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Id
bR Id
newB VarEnv Id
renv
     , rn_inScope :: InScopeSet
rn_inScope = InScopeSet -> Id -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
inScope Id
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 (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)