{-# 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 :: 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 :: 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 :: Param (FParamInfo rep) -> State FV ()
walkOnFParam = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (Param (FParamInfo rep) -> FV -> FV)
-> Param (FParamInfo 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)
-> (Param (FParamInfo rep) -> FV)
-> Param (FParamInfo 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
. Param (FParamInfo rep) -> FV
forall a. FreeIn a => a -> FV
freeIn',
      walkOnLParam :: Param (LParamInfo rep) -> State FV ()
walkOnLParam = (FV -> FV) -> State FV ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((FV -> FV) -> State FV ())
-> (Param (LParamInfo rep) -> FV -> FV)
-> Param (LParamInfo 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)
-> (Param (LParamInfo rep) -> FV)
-> Param (LParamInfo 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
. Param (LParamInfo 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 :: 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 (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 :: 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 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 BodyT rep
body) =
    Names -> FV -> FV
fvBind ([VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (Param (FParamInfo rep) -> VName)
-> [Param (FParamInfo rep)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName [Param (FParamInfo 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
<> [Param (FParamInfo rep)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [Param (FParamInfo 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 [Param (LParamInfo 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
$ (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 [Param (LParamInfo 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
<> [Param (LParamInfo rep)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [Param (LParamInfo 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 [(Param (FParamInfo rep), SubExp)]
ctxmerge [(Param (FParamInfo rep), SubExp)]
valmerge LoopForm rep
form BodyT rep
loopbody) =
    let ([Param (FParamInfo rep)]
ctxparams, Result
ctxinits) = [(Param (FParamInfo rep), SubExp)]
-> ([Param (FParamInfo rep)], Result)
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param (FParamInfo rep), SubExp)]
ctxmerge
        ([Param (FParamInfo rep)]
valparams, Result
valinits) = [(Param (FParamInfo rep), SubExp)]
-> ([Param (FParamInfo rep)], Result)
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param (FParamInfo rep), SubExp)]
valmerge
        bound_here :: Names
bound_here =
          [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$
            Scope rep -> [VName]
forall k a. Map k a -> [k]
M.keys (Scope rep -> [VName]) -> Scope rep -> [VName]
forall a b. (a -> b) -> a -> b
$
              LoopForm rep -> Scope rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf LoopForm rep
form
                Scope rep -> Scope rep -> Scope rep
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo rep)] -> Scope rep
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams ([Param (FParamInfo rep)]
ctxparams [Param (FParamInfo rep)]
-> [Param (FParamInfo rep)] -> [Param (FParamInfo rep)]
forall a. [a] -> [a] -> [a]
++ [Param (FParamInfo rep)]
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 rep -> FV
forall a. FreeIn a => a -> FV
freeIn' LoopForm rep
form
            FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo rep)] -> FV
forall a. FreeIn a => a -> FV
freeIn' ([Param (FParamInfo rep)]
ctxparams [Param (FParamInfo rep)]
-> [Param (FParamInfo rep)] -> [Param (FParamInfo rep)]
forall a. [a] -> [a] -> [a]
++ [Param (FParamInfo rep)]
valparams)
            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 [(ShapeBase SubExp, [VName], Maybe (Lambda rep, Result))]
inputs Lambda rep
lam) =
    [(ShapeBase SubExp, [VName], Maybe (Lambda rep, Result))] -> FV
forall a. FreeIn a => a -> FV
freeIn' [(ShapeBase SubExp, [VName], Maybe (Lambda rep, Result))]
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 PatternT (LetDec rep)
pat (StmAux Certificates
cs Attrs
attrs ExpDec rep
dec) Exp rep
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 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
<> PatternT (LetDec rep) -> FV
forall a. FreeIn a => a -> FV
freeIn' PatternT (LetDec 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 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 rep) => FreeIn (LoopForm rep) where
  freeIn' :: LoopForm rep -> FV
freeIn' (ForLoop VName
_ IntType
_ SubExp
bound [(Param (LParamInfo 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
<> [(Param (LParamInfo rep), VName)] -> FV
forall a. FreeIn a => a -> FV
freeIn' [(Param (LParamInfo 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 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 rep -> Names
boundInBody :: forall rep. Body rep -> Names
boundInBody = Stms rep -> Names
forall rep. Stms rep -> Names
boundByStms (Stms rep -> Names)
-> (BodyT rep -> Stms rep) -> BodyT rep -> Names
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BodyT rep -> Stms rep
forall rep. BodyT rep -> Stms rep
bodyStms

-- | The names bound by a binding.
boundByStm :: Stm rep -> Names
boundByStm :: forall rep. 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
. PatternT (LetDec rep) -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT (LetDec rep) -> [VName])
-> (Stm rep -> PatternT (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 -> PatternT (LetDec rep)
forall rep. Stm rep -> Pattern rep
stmPattern

-- | The names bound by the bindings.
boundByStms :: Stms rep -> Names
boundByStms :: forall rep. Stms rep -> Names
boundByStms = (Stm rep -> Names) -> Seq (Stm 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 :: forall rep. 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)