{-# 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
  , extendVarSet
  , 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'
  )
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
#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