{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# 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,
    nameIn,
    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 qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
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
(Names -> Names -> Bool) -> (Names -> Names -> Bool) -> Eq Names
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
(Int -> Names -> ShowS)
-> (Names -> String) -> ([Names] -> ShowS) -> Show Names
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

instance Ord Names where
  Names
x compare :: Names -> Names -> Ordering
`compare` Names
y = if Names
x Names -> Names -> Bool
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 (IntMap VName -> Names) -> IntMap VName -> Names
forall a b. (a -> b) -> a -> b
$ Names -> IntMap VName
namesIntMap Names
vs1 IntMap VName -> IntMap VName -> IntMap VName
forall a. Semigroup a => a -> a -> a
<> Names -> IntMap VName
namesIntMap Names
vs2

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

instance Pretty Names where
  ppr :: Names -> Doc
ppr = [VName] -> Doc
forall a. Pretty a => a -> Doc
ppr ([VName] -> Doc) -> (Names -> [VName]) -> Names -> Doc
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 Int -> IntMap VName -> Bool
forall a. Int -> IntMap a -> Bool
`IM.member` IntMap VName
vs

-- | Construct a name set from a list.  Slow.
namesFromList :: [VName] -> Names
namesFromList :: [VName] -> Names
namesFromList [VName]
vs = IntMap VName -> Names
Names (IntMap VName -> Names) -> IntMap VName -> Names
forall a b. (a -> b) -> a -> b
$ [(Int, VName)] -> IntMap VName
forall a. [(Int, a)] -> IntMap a
IM.fromList ([(Int, VName)] -> IntMap VName) -> [(Int, VName)] -> IntMap VName
forall a b. (a -> b) -> a -> b
$ [Int] -> [VName] -> [(Int, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((VName -> Int) -> [VName] -> [Int]
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 = IntMap VName -> [VName]
forall a. IntMap a -> [a]
IM.elems (IntMap VName -> [VName])
-> (Names -> IntMap VName) -> Names -> [VName]
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 (IntMap VName -> Names) -> IntMap VName -> Names
forall a b. (a -> b) -> a -> b
$ Int -> VName -> IntMap VName
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 (IntMap VName -> Names) -> IntMap VName -> Names
forall a b. (a -> b) -> a -> b
$ IntMap VName -> IntMap VName -> IntMap VName
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 (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ IntMap VName -> IntMap VName -> Bool
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 (IntMap VName -> Names) -> IntMap VName -> Names
forall a b. (a -> b) -> a -> b
$ IntMap VName -> IntMap VName -> IntMap VName
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 ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (VName -> VName) -> [VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> VName
f ([VName] -> [VName]) -> [VName] -> [VName]
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 Names
forall a. Monoid a => a
mempty

instance Semigroup FV where
  FV Names
fv1 <> :: FV -> FV -> FV
<> FV Names
fv2 = Names -> FV
FV (Names -> FV) -> Names -> FV
forall a b. (a -> b) -> a -> b
$ Names
fv1 Names -> Names -> Names
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 (Names -> FV) -> Names -> 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 (Names -> FV) -> Names -> 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 :: Walker rep (State FV)
freeWalker =
  Walker :: forall rep (m :: * -> *).
(SubExp -> m ())
-> (Scope rep -> Body rep -> m ())
-> (VName -> m ())
-> (RetType rep -> m ())
-> (BranchType rep -> m ())
-> (FParam rep -> m ())
-> (LParam rep -> m ())
-> (Op rep -> m ())
-> Walker rep m
Walker
    { walkOnSubExp :: SubExp -> State FV ()
walkOnSubExp = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (SubExp -> FV -> FV) -> SubExp -> State FV ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV) -> (SubExp -> FV) -> SubExp -> FV -> FV
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnBody :: Scope rep -> Body rep -> State FV ()
walkOnBody = \Scope rep
scope Body rep
body -> do
        (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ()) -> (FV -> FV) -> State FV ()
forall a b. (a -> b) -> a -> b
$ FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV) -> FV -> FV -> FV
forall a b. (a -> b) -> a -> b
$ Body rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Body rep
body
        (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ()) -> (FV -> FV) -> State FV ()
forall a b. (a -> b) -> a -> b
$ Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList (Scope rep -> [VName]
forall k a. Map k a -> [k]
M.keys Scope rep
scope)),
      walkOnVName :: VName -> State FV ()
walkOnVName = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (VName -> FV -> FV) -> VName -> State FV ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV) -> (VName -> FV) -> VName -> FV -> FV
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 = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (Op rep -> FV -> FV) -> Op rep -> State FV ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV) -> (Op rep -> FV) -> Op rep -> FV -> FV
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Op rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnFParam :: FParam rep -> State FV ()
walkOnFParam = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (FParam rep -> FV -> FV) -> FParam rep -> State FV ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV) -> (FParam rep -> FV) -> FParam rep -> FV -> FV
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FParam rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnLParam :: LParam rep -> State FV ()
walkOnLParam = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (LParam rep -> FV -> FV) -> LParam rep -> State FV ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV) -> (LParam rep -> FV) -> LParam rep -> FV -> FV
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. LParam rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnRetType :: RetType rep -> State FV ()
walkOnRetType = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (RetType rep -> FV -> FV) -> RetType rep -> State FV ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV) -> (RetType rep -> FV) -> RetType rep -> FV -> FV
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. RetType rep -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnBranchType :: BranchType rep -> State FV ()
walkOnBranchType = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (BranchType rep -> FV -> FV) -> BranchType rep -> State FV ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
(<>) (FV -> FV -> FV)
-> (BranchType rep -> FV) -> BranchType rep -> FV -> FV
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BranchType rep -> FV
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 :: Stms rep -> Result -> FV
freeInStmsAndRes Stms rep
stms Result
res =
  Names -> FV -> FV
fvBind (Stms rep -> Names
forall rep. Stms rep -> Names
boundByStms Stms rep
stms) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$ (Stm rep -> FV) -> Stms rep -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Stms rep
stms FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Result -> FV
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 (Names -> FV) -> (a -> Names) -> a -> FV
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. a -> Names
forall a. FreeIn a => a -> Names
freeIn

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

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

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

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

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

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

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 [FParam rep]
params BodyT rep
body) =
    Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (FParam rep -> VName) -> [FParam rep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> VName
forall dec. Param dec -> VName
paramName [FParam rep]
params) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$
      [RetType rep] -> FV
forall a. FreeIn a => a -> FV
freeIn' [RetType rep]
rettype FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [FParam rep] -> FV
forall a. FreeIn a => a -> FV
freeIn' [FParam rep]
params FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> BodyT rep -> FV
forall a. FreeIn a => a -> FV
freeIn' BodyT 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 [LParam rep]
params BodyT rep
body [Type]
rettype) =
    Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (LParam rep -> VName) -> [LParam rep] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map LParam rep -> VName
forall dec. Param dec -> VName
paramName [LParam rep]
params) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$
      [Type] -> FV
forall a. FreeIn a => a -> FV
freeIn' [Type]
rettype FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [LParam rep] -> FV
forall a. FreeIn a => a -> FV
freeIn' [LParam rep]
params FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> BodyT rep -> FV
forall a. FreeIn a => a -> FV
freeIn' BodyT 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) =
    BodyDec rep -> FV -> FV
forall dec. FreeDec dec => dec -> FV -> FV
precomputed BodyDec rep
dec (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$ BodyDec rep -> FV
forall a. FreeIn a => a -> FV
freeIn' BodyDec rep
dec FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Stms rep -> Result -> FV
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 [(FParam rep, SubExp)]
merge LoopForm rep
form BodyT rep
loopbody) =
    let ([FParam rep]
params, [SubExp]
args) = [(FParam rep, SubExp)] -> ([FParam rep], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam rep, SubExp)]
merge
        bound_here :: Names
bound_here =
          [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Map VName (NameInfo rep) -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName (NameInfo rep) -> [VName])
-> Map VName (NameInfo rep) -> [VName]
forall a b. (a -> b) -> a -> b
$ LoopForm rep -> Map VName (NameInfo rep)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf LoopForm rep
form Map VName (NameInfo rep)
-> Map VName (NameInfo rep) -> Map VName (NameInfo rep)
forall a. Semigroup a => a -> a -> a
<> [FParam rep] -> Map VName (NameInfo rep)
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [FParam rep]
params
     in Names -> FV -> FV
fvBind Names
bound_here (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$
          [SubExp] -> FV
forall a. FreeIn a => a -> FV
freeIn' [SubExp]
args FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> LoopForm rep -> FV
forall a. FreeIn a => a -> FV
freeIn' LoopForm rep
form FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [FParam rep] -> FV
forall a. FreeIn a => a -> FV
freeIn' [FParam rep]
params FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> BodyT rep -> FV
forall a. FreeIn a => a -> FV
freeIn' BodyT rep
loopbody
  freeIn' (WithAcc [(Shape, [VName], Maybe (Lambda rep, [SubExp]))]
inputs Lambda rep
lam) =
    [(Shape, [VName], Maybe (Lambda rep, [SubExp]))] -> FV
forall a. FreeIn a => a -> FV
freeIn' [(Shape, [VName], Maybe (Lambda rep, [SubExp]))]
inputs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Lambda rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Lambda rep
lam
  freeIn' Exp rep
e = State FV () -> FV -> FV
forall s a. State s a -> s -> s
execState (Walker rep (State FV) -> Exp rep -> State FV ()
forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Exp rep -> m ()
walkExpM Walker rep (State FV)
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) FV
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 rep
pat (StmAux Certs
cs Attrs
attrs ExpDec rep
dec) Exp rep
e) =
    Certs -> FV
forall a. FreeIn a => a -> FV
freeIn' Certs
cs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Attrs -> FV
forall a. FreeIn a => a -> FV
freeIn' Attrs
attrs
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpDec rep -> FV -> FV
forall dec. FreeDec dec => dec -> FV -> FV
precomputed ExpDec rep
dec (ExpDec rep -> FV
forall a. FreeIn a => a -> FV
freeIn' ExpDec rep
dec FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp rep
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Pat rep -> FV
forall a. FreeIn a => a -> FV
freeIn' Pat rep
pat)

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

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

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

instance FreeIn a => FreeIn (Maybe a) where
  freeIn' :: Maybe a -> FV
freeIn' = FV -> (a -> FV) -> Maybe a -> FV
forall b a. b -> (a -> b) -> Maybe a -> b
maybe FV
forall a. Monoid a => a
mempty a -> FV
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' = Type -> FV
forall a. FreeIn a => a -> FV
freeIn' (Type -> FV) -> (Ident -> Type) -> Ident -> FV
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) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v
  freeIn' Constant {} = FV
forall a. Monoid a => a
mempty

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

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

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

instance FreeIn PrimType where
  freeIn' :: PrimType -> FV
freeIn' PrimType
_ = FV
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
_) = PrimType -> FV
forall a. FreeIn a => a -> FV
freeIn' PrimType
t FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> shape -> FV
forall a. FreeIn a => a -> FV
freeIn' shape
shape
  freeIn' (Mem Space
s) = Space -> FV
forall a. FreeIn a => a -> FV
freeIn' Space
s
  freeIn' Prim {} = FV
forall a. Monoid a => a
mempty
  freeIn' (Acc VName
acc Shape
ispace [Type]
ts u
_) = (VName, Shape, [Type]) -> FV
forall a. FreeIn a => a -> FV
freeIn' (VName
acc, Shape
ispace, [Type]
ts)

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

instance FreeIn dec => FreeIn (PatElemT dec) where
  freeIn' :: PatElemT dec -> FV
freeIn' (PatElem VName
_ dec
dec) = dec -> FV
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 [(LParam rep, VName)]
loop_vars) = SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn' SubExp
bound FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [(LParam rep, VName)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [(LParam rep, VName)]
loop_vars
  freeIn' (WhileLoop VName
cond) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
cond

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

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

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

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

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

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

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

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

instance FreeIn Attrs where
  freeIn' :: Attrs -> FV
freeIn' (Attrs Set Attr
_) = FV
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) = Certs -> FV
forall a. FreeIn a => a -> FV
freeIn' Certs
cs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Attrs -> FV
forall a. FreeIn a => a -> FV
freeIn' Attrs
attrs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> dec -> FV
forall a. FreeIn a => a -> FV
freeIn' dec
dec

instance FreeIn a => FreeIn (IfDec a) where
  freeIn' :: IfDec a -> FV
freeIn' (IfDec [a]
r IfSort
_) = [a] -> FV
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
_ = FV -> FV
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
_) = a -> FV -> FV
forall dec. FreeDec dec => dec -> FV -> FV
precomputed a
a

instance FreeDec a => FreeDec [a] where
  precomputed :: [a] -> FV -> FV
precomputed [] = FV -> FV
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id
  precomputed (a
a : [a]
_) = a -> FV -> FV
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 = FV -> FV
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id
  precomputed (Just a
a) = a -> FV -> FV
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 :: Body rep -> Names
boundInBody = Stms rep -> Names
forall rep. Stms rep -> Names
boundByStms (Stms rep -> Names) -> (Body rep -> Stms rep) -> Body rep -> Names
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Body rep -> Stms rep
forall rep. BodyT rep -> Stms rep
bodyStms

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

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

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