{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeSynonymInstances #-}

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
    -- ** Working with predicates
    -- *** Searching
  , elemVarSet
  , notElemVarSet
    -- ** 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)
import           Data.Text.Prettyprint.Doc
import           GHC.Exts                  (Any)
import           GHC.Generics

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

-- | 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

-- | 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)