{-# 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 lore),
    FreeDec (BodyDec lore),
    FreeIn (FParamInfo lore),
    FreeIn (LParamInfo lore),
    FreeIn (LetDec lore),
    FreeIn (RetType lore),
    FreeIn (BranchType lore),
    FreeIn (Op lore)
  ) =>
  Walker lore (State FV)
freeWalker :: forall lore.
(FreeDec (ExpDec lore), FreeDec (BodyDec lore),
 FreeIn (FParamInfo lore), FreeIn (LParamInfo lore),
 FreeIn (LetDec lore), FreeIn (RetType lore),
 FreeIn (BranchType lore), FreeIn (Op lore)) =>
Walker lore (State FV)
freeWalker =
  Walker :: forall lore (m :: * -> *).
(SubExp -> m ())
-> (Scope lore -> Body lore -> m ())
-> (VName -> m ())
-> (RetType lore -> m ())
-> (BranchType lore -> m ())
-> (FParam lore -> m ())
-> (LParam lore -> m ())
-> (Op lore -> m ())
-> Walker lore 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 lore -> Body lore -> State FV ()
walkOnBody = \Scope lore
scope Body lore
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 lore -> FV
forall a. FreeIn a => a -> FV
freeIn' Body lore
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 lore -> [VName]
forall k a. Map k a -> [k]
M.keys Scope lore
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 lore -> State FV ()
walkOnOp = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (Op lore -> FV -> FV) -> Op lore -> 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 lore -> FV) -> Op lore -> 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 lore -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnFParam :: Param (FParamInfo lore) -> State FV ()
walkOnFParam = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (Param (FParamInfo lore) -> FV -> FV)
-> Param (FParamInfo lore)
-> 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)
-> (Param (FParamInfo lore) -> FV)
-> Param (FParamInfo lore)
-> FV
-> FV
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Param (FParamInfo lore) -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnLParam :: Param (LParamInfo lore) -> State FV ()
walkOnLParam = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (Param (LParamInfo lore) -> FV -> FV)
-> Param (LParamInfo lore)
-> 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)
-> (Param (LParamInfo lore) -> FV)
-> Param (LParamInfo lore)
-> FV
-> FV
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Param (LParamInfo lore) -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnRetType :: RetType lore -> State FV ()
walkOnRetType = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (RetType lore -> FV -> FV) -> RetType lore -> 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 lore -> FV) -> RetType lore -> 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 lore -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnBranchType :: BranchType lore -> State FV ()
walkOnBranchType = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (BranchType lore -> FV -> FV) -> BranchType lore -> 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 lore -> FV) -> BranchType lore -> 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 lore -> 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 lore),
    FreeIn (LetDec lore),
    FreeIn (LParamInfo lore),
    FreeIn (FParamInfo lore),
    FreeDec (BodyDec lore),
    FreeIn (RetType lore),
    FreeIn (BranchType lore),
    FreeDec (ExpDec lore)
  ) =>
  Stms lore ->
  Result ->
  FV
freeInStmsAndRes :: forall lore.
(FreeIn (Op lore), FreeIn (LetDec lore), FreeIn (LParamInfo lore),
 FreeIn (FParamInfo lore), FreeDec (BodyDec lore),
 FreeIn (RetType lore), FreeIn (BranchType lore),
 FreeDec (ExpDec lore)) =>
Stms lore -> Result -> FV
freeInStmsAndRes Stms lore
stms Result
res =
  Names -> FV -> FV
fvBind (Stms lore -> Names
forall lore. Stms lore -> Names
boundByStms Stms lore
stms) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$ (Stm lore -> FV) -> Stms lore -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm lore -> FV
forall a. FreeIn a => a -> FV
freeIn' Stms lore
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 :: forall a. FreeIn a => 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 lore),
    FreeDec (BodyDec lore),
    FreeIn (FParamInfo lore),
    FreeIn (LParamInfo lore),
    FreeIn (LetDec lore),
    FreeIn (RetType lore),
    FreeIn (BranchType lore),
    FreeIn (Op lore)
  ) =>
  FreeIn (FunDef lore)
  where
  freeIn' :: FunDef lore -> FV
freeIn' (FunDef Maybe EntryPoint
_ Attrs
_ Name
_ [RetType lore]
rettype [Param (FParamInfo lore)]
params BodyT lore
body) =
    Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (Param (FParamInfo lore) -> VName)
-> [Param (FParamInfo lore)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo lore) -> VName
forall dec. Param dec -> VName
paramName [Param (FParamInfo lore)]
params) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$
      [RetType lore] -> FV
forall a. FreeIn a => a -> FV
freeIn' [RetType lore]
rettype FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo lore)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [Param (FParamInfo lore)]
params FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> BodyT lore -> FV
forall a. FreeIn a => a -> FV
freeIn' BodyT lore
body

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

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

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

instance
  ( FreeDec (ExpDec lore),
    FreeDec (BodyDec lore),
    FreeIn (FParamInfo lore),
    FreeIn (LParamInfo lore),
    FreeIn (LetDec lore),
    FreeIn (RetType lore),
    FreeIn (BranchType lore),
    FreeIn (Op lore)
  ) =>
  FreeIn (Stm lore)
  where
  freeIn' :: Stm lore -> FV
freeIn' (Let PatternT (LetDec lore)
pat (StmAux Certificates
cs Attrs
attrs ExpDec lore
dec) Exp lore
e) =
    Certificates -> FV
forall a. FreeIn a => a -> FV
freeIn' Certificates
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 lore -> FV -> FV
forall dec. FreeDec dec => dec -> FV -> FV
precomputed ExpDec lore
dec (ExpDec lore -> FV
forall a. FreeIn a => a -> FV
freeIn' ExpDec lore
dec FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp lore -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp lore
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> PatternT (LetDec lore) -> FV
forall a. FreeIn a => a -> FV
freeIn' PatternT (LetDec lore)
pat)

instance FreeIn (Stm lore) => FreeIn (Stms lore) where
  freeIn' :: Stms lore -> FV
freeIn' = (Stm lore -> FV) -> Stms lore -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm lore -> 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 Result
d PrimType
_) = Result -> FV
forall a. FreeIn a => a -> FV
freeIn' Result
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 ShapeBase SubExp
ispace [Type]
ts u
_) = (VName, ShapeBase SubExp, [Type]) -> FV
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 VName
_ dec
dec) = 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 lore) => FreeIn (LoopForm lore) where
  freeIn' :: LoopForm lore -> FV
freeIn' (ForLoop VName
_ IntType
_ SubExp
bound [(Param (LParamInfo lore), VName)]
loop_vars) = SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn' SubExp
bound FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [(Param (LParamInfo lore), VName)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [(Param (LParamInfo lore), 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 dec => FreeIn (PatternT dec) where
  freeIn' :: PatternT dec -> FV
freeIn' (Pattern [PatElemT dec]
context [PatElemT dec]
values) =
    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] -> FV) -> [PatElemT dec] -> FV
forall a b. (a -> b) -> a -> b
$ [PatElemT dec]
context [PatElemT dec] -> [PatElemT dec] -> [PatElemT dec]
forall a. [a] -> [a] -> [a]
++ [PatElemT dec]
values
    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] -> [VName]) -> [PatElemT dec] -> [VName]
forall a b. (a -> b) -> a -> b
$ [PatElemT dec]
context [PatElemT dec] -> [PatElemT dec] -> [PatElemT dec]
forall a. [a] -> [a] -> [a]
++ [PatElemT dec]
values

instance FreeIn Certificates where
  freeIn' :: Certificates -> FV
freeIn' (Certificates [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 Certificates
cs Attrs
attrs dec
dec) = Certificates -> FV
forall a. FreeIn a => a -> FV
freeIn' Certificates
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 lore -> Names
boundInBody :: forall lore. Body lore -> Names
boundInBody = Stms lore -> Names
forall lore. Stms lore -> Names
boundByStms (Stms lore -> Names)
-> (BodyT lore -> Stms lore) -> BodyT lore -> Names
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BodyT lore -> Stms lore
forall lore. BodyT lore -> Stms lore
bodyStms

-- | The names bound by a binding.
boundByStm :: Stm lore -> Names
boundByStm :: forall lore. Stm lore -> Names
boundByStm = [VName] -> Names
namesFromList ([VName] -> Names) -> (Stm lore -> [VName]) -> Stm lore -> Names
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PatternT (LetDec lore) -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT (LetDec lore) -> [VName])
-> (Stm lore -> PatternT (LetDec lore)) -> Stm lore -> [VName]
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Stm lore -> PatternT (LetDec lore)
forall lore. Stm lore -> Pattern lore
stmPattern

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

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