{-# 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 (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 (Op lore)) =>
Walker lore (State FV)
freeWalker =
  Walker lore (State FV)
forall (m :: * -> *) lore. Monad m => Walker lore m
identityWalker
    { 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'
    }

-- | 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),
    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),
 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 [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 (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 (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 (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),
 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 (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' 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 (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 (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 shape => FreeIn (TypeBase shape u) where
  freeIn' :: TypeBase shape u -> FV
freeIn' (Array PrimType
_ shape
shape u
_) = 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 PrimType
_) = FV
forall a. Monoid a => a
mempty

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)