{-# LANGUAGE UndecidableInstances #-}

-- | Facilities for determining which names are used in some syntactic
-- construct.  The most important interface is the 'FreeIn' class and
-- its instances, but for reasons related to the Haskell type system,
-- some constructs have specialised functions.
module Futhark.IR.Prop.Names
  ( -- * Free names
    Names,
    namesIntMap,
    namesIntSet,
    nameIn,
    notNameIn,
    oneName,
    namesFromList,
    namesToList,
    namesIntersection,
    namesIntersect,
    namesSubtract,
    mapNames,

    -- * Class
    FreeIn (..),
    freeIn,

    -- * Specialised Functions
    freeInStmsAndRes,

    -- * Bound Names
    boundInBody,
    boundByStm,
    boundByStms,
    boundByLambda,

    -- * Efficient computation
    FreeDec (..),
    FV,
    fvBind,
    fvName,
    fvNames,
  )
where

import Control.Category
import Control.Monad.State.Strict
import Data.Foldable
import Data.IntMap.Strict qualified as IM
import Data.IntSet qualified as IS
import Data.Map.Strict qualified as M
import Data.Set qualified as S
import Futhark.IR.Prop.Patterns
import Futhark.IR.Prop.Scope
import Futhark.IR.Syntax
import Futhark.IR.Traversals
import Futhark.Util.Pretty
import Prelude hiding (id, (.))

-- | A set of names.  Note that the 'Ord' instance is a dummy that
-- treats everything as 'EQ' if '==', and otherwise 'LT'.
newtype Names = Names (IM.IntMap VName)
  deriving (Names -> Names -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Names -> Names -> Bool
$c/= :: Names -> Names -> Bool
== :: Names -> Names -> Bool
$c== :: Names -> Names -> Bool
Eq, Int -> Names -> ShowS
[Names] -> ShowS
Names -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Names] -> ShowS
$cshowList :: [Names] -> ShowS
show :: Names -> String
$cshow :: Names -> String
showsPrec :: Int -> Names -> ShowS
$cshowsPrec :: Int -> Names -> ShowS
Show)

-- | Retrieve the data structure underlying the names representation.
namesIntMap :: Names -> IM.IntMap VName
namesIntMap :: Names -> IntMap VName
namesIntMap (Names IntMap VName
m) = IntMap VName
m

-- | Retrieve the set of tags in the names set.
namesIntSet :: Names -> IS.IntSet
namesIntSet :: Names -> IntSet
namesIntSet (Names IntMap VName
m) = forall a. IntMap a -> IntSet
IM.keysSet IntMap VName
m

instance Ord Names where
  Names
x compare :: Names -> Names -> Ordering
`compare` Names
y = if Names
x forall a. Eq a => a -> a -> Bool
== Names
y then Ordering
EQ else Ordering
LT

instance Semigroup Names where
  Names
vs1 <> :: Names -> Names -> Names
<> Names
vs2 = IntMap VName -> Names
Names forall a b. (a -> b) -> a -> b
$ Names -> IntMap VName
namesIntMap Names
vs1 forall a. Semigroup a => a -> a -> a
<> Names -> IntMap VName
namesIntMap Names
vs2

instance Monoid Names where
  mempty :: Names
mempty = IntMap VName -> Names
Names forall a. Monoid a => a
mempty

instance Pretty Names where
  pretty :: forall ann. Names -> Doc ann
pretty = forall a ann. Pretty a => a -> Doc ann
pretty forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Names -> [VName]
namesToList

-- | Does the set of names contain this name?
nameIn :: VName -> Names -> Bool
nameIn :: VName -> Names -> Bool
nameIn VName
v (Names IntMap VName
vs) = VName -> Int
baseTag VName
v forall a. Int -> IntMap a -> Bool
`IM.member` IntMap VName
vs

-- | Does the set of names not contain this name?
notNameIn :: VName -> Names -> Bool
notNameIn :: VName -> Names -> Bool
notNameIn VName
v (Names IntMap VName
vs) = VName -> Int
baseTag VName
v forall a. Int -> IntMap a -> Bool
`IM.notMember` IntMap VName
vs

-- | Construct a name set from a list.  Slow.
namesFromList :: [VName] -> Names
namesFromList :: [VName] -> Names
namesFromList [VName]
vs = IntMap VName -> Names
Names forall a b. (a -> b) -> a -> b
$ forall a. [(Int, a)] -> IntMap a
IM.fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map VName -> Int
baseTag [VName]
vs) [VName]
vs

-- | Turn a name set into a list of names.  Slow.
namesToList :: Names -> [VName]
namesToList :: Names -> [VName]
namesToList = forall a. IntMap a -> [a]
IM.elems forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Names -> IntMap VName
namesIntMap

-- | Construct a name set from a single name.
oneName :: VName -> Names
oneName :: VName -> Names
oneName VName
v = IntMap VName -> Names
Names forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> IntMap a
IM.singleton (VName -> Int
baseTag VName
v) VName
v

-- | The intersection of two name sets.
namesIntersection :: Names -> Names -> Names
namesIntersection :: Names -> Names -> Names
namesIntersection (Names IntMap VName
vs1) (Names IntMap VName
vs2) = IntMap VName -> Names
Names forall a b. (a -> b) -> a -> b
$ forall a b. IntMap a -> IntMap b -> IntMap a
IM.intersection IntMap VName
vs1 IntMap VName
vs2

-- | Do the two name sets intersect?
namesIntersect :: Names -> Names -> Bool
namesIntersect :: Names -> Names -> Bool
namesIntersect Names
vs1 Names
vs2 = Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall a b. IntMap a -> IntMap b -> Bool
IM.disjoint (Names -> IntMap VName
namesIntMap Names
vs1) (Names -> IntMap VName
namesIntMap Names
vs2)

-- | Subtract the latter name set from the former.
namesSubtract :: Names -> Names -> Names
namesSubtract :: Names -> Names -> Names
namesSubtract (Names IntMap VName
vs1) (Names IntMap VName
vs2) = IntMap VName -> Names
Names forall a b. (a -> b) -> a -> b
$ forall a b. IntMap a -> IntMap b -> IntMap a
IM.difference IntMap VName
vs1 IntMap VName
vs2

-- | Map over the names in a set.
mapNames :: (VName -> VName) -> Names -> Names
mapNames :: (VName -> VName) -> Names -> Names
mapNames VName -> VName
f Names
vs = [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> VName
f forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
vs

-- | A computation to build a free variable set.
newtype FV = FV {FV -> Names
unFV :: Names}

-- Right now the variable set is just stored explicitly, without the
-- fancy functional representation that GHC uses.  Turns out it's
-- faster this way.

instance Monoid FV where
  mempty :: FV
mempty = Names -> FV
FV forall a. Monoid a => a
mempty

instance Semigroup FV where
  FV Names
fv1 <> :: FV -> FV -> FV
<> FV Names
fv2 = Names -> FV
FV forall a b. (a -> b) -> a -> b
$ Names
fv1 forall a. Semigroup a => a -> a -> a
<> Names
fv2

-- | Consider a variable to be bound in the given 'FV' computation.
fvBind :: Names -> FV -> FV
fvBind :: Names -> FV -> FV
fvBind Names
vs (FV Names
fv) = Names -> FV
FV forall a b. (a -> b) -> a -> b
$ Names
fv Names -> Names -> Names
`namesSubtract` Names
vs

-- | Take note of a variable reference.
fvName :: VName -> FV
fvName :: VName -> FV
fvName VName
v = Names -> FV
FV forall a b. (a -> b) -> a -> b
$ VName -> Names
oneName VName
v

-- | Take note of a set of variable references.
fvNames :: Names -> FV
fvNames :: Names -> FV
fvNames = Names -> FV
FV

freeWalker ::
  ( FreeDec (ExpDec rep),
    FreeDec (BodyDec rep),
    FreeIn (FParamInfo rep),
    FreeIn (LParamInfo rep),
    FreeIn (LetDec rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    FreeIn (Op rep)
  ) =>
  Walker rep (State FV)
freeWalker :: forall rep.
(FreeDec (ExpDec rep), FreeDec (BodyDec rep),
 FreeIn (FParamInfo rep), FreeIn (LParamInfo rep),
 FreeIn (LetDec rep), FreeIn (RetType rep), FreeIn (BranchType rep),
 FreeIn (Op rep)) =>
Walker rep (State FV)
freeWalker =
  Walker
    { walkOnSubExp :: SubExp -> State FV ()
walkOnSubExp = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Semigroup a => a -> a -> a
(<>) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> FV
freeIn',
      walkOnBody :: Scope rep -> Body rep -> State FV ()
walkOnBody = \Scope rep
scope Body rep
body -> do
        forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall a. Semigroup a => a -> a -> a
(<>) forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> FV
freeIn' Body rep
body
        forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList (forall k a. Map k a -> [k]
M.keys Scope rep
scope)),
      walkOnVName :: VName -> State FV ()
walkOnVName = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Semigroup a => a -> a -> a
(<>) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VName -> FV
fvName,
      walkOnOp :: Op rep -> State FV ()
walkOnOp = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Semigroup a => a -> a -> a
(<>) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> FV
freeIn',
      walkOnFParam :: Param (FParamInfo rep) -> State FV ()
walkOnFParam = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Semigroup a => a -> a -> a
(<>) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> FV
freeIn',
      walkOnLParam :: Param (LParamInfo rep) -> State FV ()
walkOnLParam = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Semigroup a => a -> a -> a
(<>) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> FV
freeIn',
      walkOnRetType :: RetType rep -> State FV ()
walkOnRetType = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Semigroup a => a -> a -> a
(<>) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> FV
freeIn',
      walkOnBranchType :: BranchType rep -> State FV ()
walkOnBranchType = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Semigroup a => a -> a -> a
(<>) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> FV
freeIn'
    }

-- | Return the set of variable names that are free in the given
-- statements and result.  Filters away the names that are bound by
-- the statements.
freeInStmsAndRes ::
  ( FreeIn (Op rep),
    FreeIn (LetDec rep),
    FreeIn (LParamInfo rep),
    FreeIn (FParamInfo rep),
    FreeDec (BodyDec rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    FreeDec (ExpDec rep)
  ) =>
  Stms rep ->
  Result ->
  FV
freeInStmsAndRes :: forall rep.
(FreeIn (Op rep), FreeIn (LetDec rep), FreeIn (LParamInfo rep),
 FreeIn (FParamInfo rep), FreeDec (BodyDec rep),
 FreeIn (RetType rep), FreeIn (BranchType rep),
 FreeDec (ExpDec rep)) =>
Stms rep -> Result -> FV
freeInStmsAndRes Stms rep
stms Result
res =
  Names -> FV -> FV
fvBind (forall rep. Stms rep -> Names
boundByStms Stms rep
stms) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. FreeIn a => a -> FV
freeIn' Stms rep
stms forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Result
res

-- | A class indicating that we can obtain free variable information
-- from values of this type.
class FreeIn a where
  freeIn' :: a -> FV
  freeIn' = Names -> FV
fvNames forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> Names
freeIn

-- | The free variables of some syntactic construct.
freeIn :: FreeIn a => a -> Names
freeIn :: forall a. FreeIn a => a -> Names
freeIn = FV -> Names
unFV forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn FV where
  freeIn' :: FV -> FV
freeIn' = forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id

instance FreeIn () where
  freeIn' :: () -> FV
freeIn' () = forall a. Monoid a => a
mempty

instance FreeIn Int where
  freeIn' :: Int -> FV
freeIn' = forall a b. a -> b -> a
const forall a. Monoid a => a
mempty

instance (FreeIn a, FreeIn b) => FreeIn (a, b) where
  freeIn' :: (a, b) -> FV
freeIn' (a
a, b
b) = forall a. FreeIn a => a -> FV
freeIn' a
a forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' b
b

instance (FreeIn a, FreeIn b, FreeIn c) => FreeIn (a, b, c) where
  freeIn' :: (a, b, c) -> FV
freeIn' (a
a, b
b, c
c) = forall a. FreeIn a => a -> FV
freeIn' a
a forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' b
b forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' c
c

instance (FreeIn a, FreeIn b, FreeIn c, FreeIn d) => FreeIn (a, b, c, d) where
  freeIn' :: (a, b, c, d) -> FV
freeIn' (a
a, b
b, c
c, d
d) = forall a. FreeIn a => a -> FV
freeIn' a
a forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' b
b forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' c
c forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' d
d

instance (FreeIn a, FreeIn b) => FreeIn (Either a b) where
  freeIn' :: Either a b -> FV
freeIn' = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. FreeIn a => a -> FV
freeIn' forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn a => FreeIn [a] where
  freeIn' :: [a] -> FV
freeIn' = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn a => FreeIn (S.Set a) where
  freeIn' :: Set a -> FV
freeIn' = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn (NoOp rep) where
  freeIn' :: NoOp rep -> FV
freeIn' NoOp rep
NoOp = forall a. Monoid a => a
mempty

instance
  ( FreeDec (ExpDec rep),
    FreeDec (BodyDec rep),
    FreeIn (FParamInfo rep),
    FreeIn (LParamInfo rep),
    FreeIn (LetDec rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    FreeIn (Op rep)
  ) =>
  FreeIn (FunDef rep)
  where
  freeIn' :: FunDef rep -> FV
freeIn' (FunDef Maybe EntryPoint
_ Attrs
_ Name
_ [RetType rep]
rettype [Param (FParamInfo rep)]
params Body rep
body) =
    Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName [Param (FParamInfo rep)]
params) forall a b. (a -> b) -> a -> b
$
      forall a. FreeIn a => a -> FV
freeIn' [RetType rep]
rettype forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' [Param (FParamInfo rep)]
params forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Body rep
body

instance
  ( FreeDec (ExpDec rep),
    FreeDec (BodyDec rep),
    FreeIn (FParamInfo rep),
    FreeIn (LParamInfo rep),
    FreeIn (LetDec rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    FreeIn (Op rep)
  ) =>
  FreeIn (Lambda rep)
  where
  freeIn' :: Lambda rep -> FV
freeIn' (Lambda [Param (LParamInfo rep)]
params Body rep
body [Type]
rettype) =
    Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName [Param (LParamInfo rep)]
params) forall a b. (a -> b) -> a -> b
$
      forall a. FreeIn a => a -> FV
freeIn' [Type]
rettype forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' [Param (LParamInfo rep)]
params forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Body rep
body

instance
  ( FreeDec (ExpDec rep),
    FreeDec (BodyDec rep),
    FreeIn (FParamInfo rep),
    FreeIn (LParamInfo rep),
    FreeIn (LetDec rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    FreeIn (Op rep)
  ) =>
  FreeIn (Body rep)
  where
  freeIn' :: Body rep -> FV
freeIn' (Body BodyDec rep
dec Stms rep
stms Result
res) =
    forall dec. FreeDec dec => dec -> FV -> FV
precomputed BodyDec rep
dec forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> FV
freeIn' BodyDec rep
dec forall a. Semigroup a => a -> a -> a
<> forall rep.
(FreeIn (Op rep), FreeIn (LetDec rep), FreeIn (LParamInfo rep),
 FreeIn (FParamInfo rep), FreeDec (BodyDec rep),
 FreeIn (RetType rep), FreeIn (BranchType rep),
 FreeDec (ExpDec rep)) =>
Stms rep -> Result -> FV
freeInStmsAndRes Stms rep
stms Result
res

instance
  ( FreeDec (ExpDec rep),
    FreeDec (BodyDec rep),
    FreeIn (FParamInfo rep),
    FreeIn (LParamInfo rep),
    FreeIn (LetDec rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    FreeIn (Op rep)
  ) =>
  FreeIn (Exp rep)
  where
  freeIn' :: Exp rep -> FV
freeIn' (DoLoop [(Param (FParamInfo rep), SubExp)]
merge LoopForm rep
form Body rep
loopbody) =
    let ([Param (FParamInfo rep)]
params, [SubExp]
args) = forall a b. [(a, b)] -> ([a], [b])
unzip [(Param (FParamInfo rep), SubExp)]
merge
        bound_here :: Names
bound_here =
          [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys forall a b. (a -> b) -> a -> b
$ forall rep a. Scoped rep a => a -> Scope rep
scopeOf LoopForm rep
form forall a. Semigroup a => a -> a -> a
<> forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param (FParamInfo rep)]
params
     in Names -> FV -> FV
fvBind Names
bound_here forall a b. (a -> b) -> a -> b
$
          forall a. FreeIn a => a -> FV
freeIn' [SubExp]
args forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' LoopForm rep
form forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' [Param (FParamInfo rep)]
params forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Body rep
loopbody
  freeIn' (WithAcc [WithAccInput rep]
inputs Lambda rep
lam) =
    forall a. FreeIn a => a -> FV
freeIn' [WithAccInput rep]
inputs forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Lambda rep
lam
  freeIn' Exp rep
e = forall s a. State s a -> s -> s
execState (forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Exp rep -> m ()
walkExpM forall rep.
(FreeDec (ExpDec rep), FreeDec (BodyDec rep),
 FreeIn (FParamInfo rep), FreeIn (LParamInfo rep),
 FreeIn (LetDec rep), FreeIn (RetType rep), FreeIn (BranchType rep),
 FreeIn (Op rep)) =>
Walker rep (State FV)
freeWalker Exp rep
e) forall a. Monoid a => a
mempty

instance
  ( FreeDec (ExpDec rep),
    FreeDec (BodyDec rep),
    FreeIn (FParamInfo rep),
    FreeIn (LParamInfo rep),
    FreeIn (LetDec rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    FreeIn (Op rep)
  ) =>
  FreeIn (Stm rep)
  where
  freeIn' :: Stm rep -> FV
freeIn' (Let Pat (LetDec rep)
pat (StmAux Certs
cs Attrs
attrs ExpDec rep
dec) Exp rep
e) =
    forall a. FreeIn a => a -> FV
freeIn' Certs
cs
      forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Attrs
attrs
      forall a. Semigroup a => a -> a -> a
<> forall dec. FreeDec dec => dec -> FV -> FV
precomputed ExpDec rep
dec (forall a. FreeIn a => a -> FV
freeIn' ExpDec rep
dec forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Exp rep
e forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Pat (LetDec rep)
pat)

instance FreeIn (Stm rep) => FreeIn (Stms rep) where
  freeIn' :: Stms rep -> FV
freeIn' = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn body => FreeIn (Case body) where
  freeIn' :: Case body -> FV
freeIn' = forall a. FreeIn a => a -> FV
freeIn' forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall body. Case body -> body
caseBody

instance FreeIn Names where
  freeIn' :: Names -> FV
freeIn' = Names -> FV
fvNames

instance FreeIn Bool where
  freeIn' :: Bool -> FV
freeIn' Bool
_ = forall a. Monoid a => a
mempty

instance FreeIn a => FreeIn (Maybe a) where
  freeIn' :: Maybe a -> FV
freeIn' = forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn VName where
  freeIn' :: VName -> FV
freeIn' = VName -> FV
fvName

instance FreeIn Ident where
  freeIn' :: Ident -> FV
freeIn' = forall a. FreeIn a => a -> FV
freeIn' forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Ident -> Type
identType

instance FreeIn SubExp where
  freeIn' :: SubExp -> FV
freeIn' (Var VName
v) = forall a. FreeIn a => a -> FV
freeIn' VName
v
  freeIn' Constant {} = forall a. Monoid a => a
mempty

instance FreeIn Space where
  freeIn' :: Space -> FV
freeIn' (ScalarSpace [SubExp]
d PrimType
_) = forall a. FreeIn a => a -> FV
freeIn' [SubExp]
d
  freeIn' Space
DefaultSpace = forall a. Monoid a => a
mempty
  freeIn' (Space String
_) = forall a. Monoid a => a
mempty

instance FreeIn d => FreeIn (ShapeBase d) where
  freeIn' :: ShapeBase d -> FV
freeIn' = forall a. FreeIn a => a -> FV
freeIn' forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall d. ShapeBase d -> [d]
shapeDims

instance FreeIn d => FreeIn (Ext d) where
  freeIn' :: Ext d -> FV
freeIn' (Free d
x) = forall a. FreeIn a => a -> FV
freeIn' d
x
  freeIn' (Ext Int
_) = forall a. Monoid a => a
mempty

instance FreeIn PrimType where
  freeIn' :: PrimType -> FV
freeIn' PrimType
_ = forall a. Monoid a => a
mempty

instance FreeIn shape => FreeIn (TypeBase shape u) where
  freeIn' :: TypeBase shape u -> FV
freeIn' (Array PrimType
t shape
shape u
_) = forall a. FreeIn a => a -> FV
freeIn' PrimType
t forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' shape
shape
  freeIn' (Mem Space
s) = forall a. FreeIn a => a -> FV
freeIn' Space
s
  freeIn' Prim {} = forall a. Monoid a => a
mempty
  freeIn' (Acc VName
acc ShapeBase SubExp
ispace [Type]
ts u
_) = forall a. FreeIn a => a -> FV
freeIn' (VName
acc, ShapeBase SubExp
ispace, [Type]
ts)

instance FreeIn dec => FreeIn (Param dec) where
  freeIn' :: Param dec -> FV
freeIn' (Param Attrs
attrs VName
_ dec
dec) = forall a. FreeIn a => a -> FV
freeIn' Attrs
attrs forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' dec
dec

instance FreeIn dec => FreeIn (PatElem dec) where
  freeIn' :: PatElem dec -> FV
freeIn' (PatElem VName
_ dec
dec) = forall a. FreeIn a => a -> FV
freeIn' dec
dec

instance FreeIn (LParamInfo rep) => FreeIn (LoopForm rep) where
  freeIn' :: LoopForm rep -> FV
freeIn' (ForLoop VName
_ IntType
_ SubExp
bound [(Param (LParamInfo rep), VName)]
loop_vars) = forall a. FreeIn a => a -> FV
freeIn' SubExp
bound forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' [(Param (LParamInfo rep), VName)]
loop_vars
  freeIn' (WhileLoop VName
cond) = forall a. FreeIn a => a -> FV
freeIn' VName
cond

instance FreeIn d => FreeIn (DimIndex d) where
  freeIn' :: DimIndex d -> FV
freeIn' = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
Data.Foldable.foldMap forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn d => FreeIn (Slice d) where
  freeIn' :: Slice d -> FV
freeIn' = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
Data.Foldable.foldMap forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn d => FreeIn (FlatDimIndex d) where
  freeIn' :: FlatDimIndex d -> FV
freeIn' = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
Data.Foldable.foldMap forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn d => FreeIn (FlatSlice d) where
  freeIn' :: FlatSlice d -> FV
freeIn' = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
Data.Foldable.foldMap forall a. FreeIn a => a -> FV
freeIn'

instance FreeIn SubExpRes where
  freeIn' :: SubExpRes -> FV
freeIn' (SubExpRes Certs
cs SubExp
se) = forall a. FreeIn a => a -> FV
freeIn' Certs
cs forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' SubExp
se

instance FreeIn dec => FreeIn (Pat dec) where
  freeIn' :: Pat dec -> FV
freeIn' (Pat [PatElem dec]
xs) =
    Names -> FV -> FV
fvBind Names
bound_here forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> FV
freeIn' [PatElem dec]
xs
    where
      bound_here :: Names
bound_here = [VName] -> Names
namesFromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall dec. PatElem dec -> VName
patElemName [PatElem dec]
xs

instance FreeIn Certs where
  freeIn' :: Certs -> FV
freeIn' (Certs [VName]
cs) = forall a. FreeIn a => a -> FV
freeIn' [VName]
cs

instance FreeIn Attrs where
  freeIn' :: Attrs -> FV
freeIn' (Attrs Set Attr
_) = forall a. Monoid a => a
mempty

instance FreeIn dec => FreeIn (StmAux dec) where
  freeIn' :: StmAux dec -> FV
freeIn' (StmAux Certs
cs Attrs
attrs dec
dec) = forall a. FreeIn a => a -> FV
freeIn' Certs
cs forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' Attrs
attrs forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' dec
dec

instance FreeIn a => FreeIn (MatchDec a) where
  freeIn' :: MatchDec a -> FV
freeIn' (MatchDec [a]
r MatchSort
_) = forall a. FreeIn a => a -> FV
freeIn' [a]
r

-- | Either return precomputed free names stored in the attribute, or
-- the freshly computed names.  Relies on lazy evaluation to avoid the
-- work.
class FreeIn dec => FreeDec dec where
  precomputed :: dec -> FV -> FV
  precomputed dec
_ = forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id

instance FreeDec ()

instance (FreeDec a, FreeIn b) => FreeDec (a, b) where
  precomputed :: (a, b) -> FV -> FV
precomputed (a
a, b
_) = forall dec. FreeDec dec => dec -> FV -> FV
precomputed a
a

instance FreeDec a => FreeDec [a] where
  precomputed :: [a] -> FV -> FV
precomputed [] = forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id
  precomputed (a
a : [a]
_) = forall dec. FreeDec dec => dec -> FV -> FV
precomputed a
a

instance FreeDec a => FreeDec (Maybe a) where
  precomputed :: Maybe a -> FV -> FV
precomputed Maybe a
Nothing = forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id
  precomputed (Just a
a) = forall dec. FreeDec dec => dec -> FV -> FV
precomputed a
a

instance FreeDec Names where
  precomputed :: Names -> FV -> FV
precomputed Names
_ FV
fv = FV
fv

-- | The names bound by the bindings immediately in a t'Body'.
boundInBody :: Body rep -> Names
boundInBody :: forall rep. Body rep -> Names
boundInBody = forall rep. Stms rep -> Names
boundByStms forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall rep. Body rep -> Stms rep
bodyStms

-- | The names bound by a binding.
boundByStm :: Stm rep -> Names
boundByStm :: forall rep. Stm rep -> Names
boundByStm = [VName] -> Names
namesFromList forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall dec. Pat dec -> [VName]
patNames forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall rep. Stm rep -> Pat (LetDec rep)
stmPat

-- | The names bound by the bindings.
boundByStms :: Stms rep -> Names
boundByStms :: forall rep. Stms rep -> Names
boundByStms = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall rep. Stm rep -> Names
boundByStm

-- | The names of the lambda parameters plus the index parameter.
boundByLambda :: Lambda rep -> [VName]
boundByLambda :: forall rep. Lambda rep -> [VName]
boundByLambda Lambda rep
lam = forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName (forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam)