-- | Facilities for computing free variables in an expression, which
-- we need for both lambda-lifting and defunctionalisation.
module Futhark.Internalise.FreeVars
  ( freeVars,
    without,
    member,
    ident,
    size,
    sizes,
    NameSet (..),
    patternVars,
  )
where

import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Futhark.IR.Pretty ()
import Language.Futhark

-- | A set of names where we also track uniqueness.
newtype NameSet = NameSet {NameSet -> Map VName StructType
unNameSet :: M.Map VName StructType}
  deriving (Int -> NameSet -> ShowS
[NameSet] -> ShowS
NameSet -> String
(Int -> NameSet -> ShowS)
-> (NameSet -> String) -> ([NameSet] -> ShowS) -> Show NameSet
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NameSet] -> ShowS
$cshowList :: [NameSet] -> ShowS
show :: NameSet -> String
$cshow :: NameSet -> String
showsPrec :: Int -> NameSet -> ShowS
$cshowsPrec :: Int -> NameSet -> ShowS
Show)

instance Semigroup NameSet where
  NameSet Map VName StructType
x <> :: NameSet -> NameSet -> NameSet
<> NameSet Map VName StructType
y = Map VName StructType -> NameSet
NameSet (Map VName StructType -> NameSet)
-> Map VName StructType -> NameSet
forall a b. (a -> b) -> a -> b
$ (StructType -> StructType -> StructType)
-> Map VName StructType
-> Map VName StructType
-> Map VName StructType
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith StructType -> StructType -> StructType
forall a. Ord a => a -> a -> a
max Map VName StructType
x Map VName StructType
y

instance Monoid NameSet where
  mempty :: NameSet
mempty = Map VName StructType -> NameSet
NameSet Map VName StructType
forall a. Monoid a => a
mempty

-- | Set subtraction.
without :: NameSet -> S.Set VName -> NameSet
without :: NameSet -> Set VName -> NameSet
without (NameSet Map VName StructType
x) Set VName
y = Map VName StructType -> NameSet
NameSet (Map VName StructType -> NameSet)
-> Map VName StructType -> NameSet
forall a b. (a -> b) -> a -> b
$ (VName -> StructType -> Bool)
-> Map VName StructType -> Map VName StructType
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey VName -> StructType -> Bool
forall {p}. VName -> p -> Bool
keep Map VName StructType
x
  where
    keep :: VName -> p -> Bool
keep VName
k p
_ = VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
y

withoutM :: NameSet -> NameSet -> NameSet
withoutM :: NameSet -> NameSet -> NameSet
withoutM (NameSet Map VName StructType
x) (NameSet Map VName StructType
y) = Map VName StructType -> NameSet
NameSet (Map VName StructType -> NameSet)
-> Map VName StructType -> NameSet
forall a b. (a -> b) -> a -> b
$ Map VName StructType
x Map VName StructType
-> Map VName StructType -> Map VName StructType
forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` Map VName StructType
y

-- | Is this name in the 'NameSet'?
member :: VName -> NameSet -> Bool
member :: VName -> NameSet -> Bool
member VName
v (NameSet Map VName StructType
m) = VName
v VName -> Map VName StructType -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName StructType
m

-- | A 'NameSet' with a single 'Nonunique' name.
ident :: Ident -> NameSet
ident :: IdentBase Info VName -> NameSet
ident IdentBase Info VName
v = Map VName StructType -> NameSet
NameSet (Map VName StructType -> NameSet)
-> Map VName StructType -> NameSet
forall a b. (a -> b) -> a -> b
$ VName -> StructType -> Map VName StructType
forall k a. k -> a -> Map k a
M.singleton (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
v) (TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (TypeBase (DimDecl VName) Aliasing -> StructType)
-> TypeBase (DimDecl VName) Aliasing -> StructType
forall a b. (a -> b) -> a -> b
$ Info (TypeBase (DimDecl VName) Aliasing)
-> TypeBase (DimDecl VName) Aliasing
forall a. Info a -> a
unInfo (Info (TypeBase (DimDecl VName) Aliasing)
 -> TypeBase (DimDecl VName) Aliasing)
-> Info (TypeBase (DimDecl VName) Aliasing)
-> TypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName -> Info (TypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) vn.
IdentBase f vn -> f (TypeBase (DimDecl VName) Aliasing)
identType IdentBase Info VName
v)

size :: VName -> NameSet
size :: VName -> NameSet
size VName
v = Map VName StructType -> NameSet
NameSet (Map VName StructType -> NameSet)
-> Map VName StructType -> NameSet
forall a b. (a -> b) -> a -> b
$ VName -> StructType -> Map VName StructType
forall k a. k -> a -> Map k a
M.singleton VName
v (StructType -> Map VName StructType)
-> StructType -> Map VName StructType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) () -> StructType)
-> ScalarTypeBase (DimDecl VName) () -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) ())
-> PrimType -> ScalarTypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

sizes :: S.Set VName -> NameSet
sizes :: Set VName -> NameSet
sizes = (VName -> NameSet) -> Set VName -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap VName -> NameSet
size

-- | Compute the set of free variables of an expression.
freeVars :: Exp -> NameSet
freeVars :: Exp -> NameSet
freeVars Exp
expr = case Exp
expr of
  Literal {} -> NameSet
forall a. Monoid a => a
mempty
  IntLit {} -> NameSet
forall a. Monoid a => a
mempty
  FloatLit {} -> NameSet
forall a. Monoid a => a
mempty
  StringLit {} -> NameSet
forall a. Monoid a => a
mempty
  Parens Exp
e SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e
  QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e
  TupLit [Exp]
es SrcLoc
_ -> (Exp -> NameSet) -> [Exp] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> NameSet
freeVars [Exp]
es
  RecordLit [FieldBase Info VName]
fs SrcLoc
_ -> (FieldBase Info VName -> NameSet)
-> [FieldBase Info VName] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap FieldBase Info VName -> NameSet
freeVarsField [FieldBase Info VName]
fs
    where
      freeVarsField :: FieldBase Info VName -> NameSet
freeVarsField (RecordFieldExplicit Name
_ Exp
e SrcLoc
_) = Exp -> NameSet
freeVars Exp
e
      freeVarsField (RecordFieldImplicit VName
vn Info (TypeBase (DimDecl VName) Aliasing)
t SrcLoc
_) = IdentBase Info VName -> NameSet
ident (IdentBase Info VName -> NameSet)
-> IdentBase Info VName -> NameSet
forall a b. (a -> b) -> a -> b
$ VName
-> Info (TypeBase (DimDecl VName) Aliasing)
-> SrcLoc
-> IdentBase Info VName
forall (f :: * -> *) vn.
vn
-> f (TypeBase (DimDecl VName) Aliasing)
-> SrcLoc
-> IdentBase f vn
Ident VName
vn Info (TypeBase (DimDecl VName) Aliasing)
t SrcLoc
forall a. Monoid a => a
mempty
  ArrayLit [Exp]
es Info (TypeBase (DimDecl VName) Aliasing)
t SrcLoc
_ ->
    (Exp -> NameSet) -> [Exp] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> NameSet
freeVars [Exp]
es NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Set VName -> NameSet
sizes (TypeBase (DimDecl VName) Aliasing -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames (TypeBase (DimDecl VName) Aliasing -> Set VName)
-> TypeBase (DimDecl VName) Aliasing -> Set VName
forall a b. (a -> b) -> a -> b
$ Info (TypeBase (DimDecl VName) Aliasing)
-> TypeBase (DimDecl VName) Aliasing
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) Aliasing)
t)
  Range Exp
e Maybe Exp
me Inclusiveness Exp
incl (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ ->
    Exp -> NameSet
freeVars Exp
e NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (Exp -> NameSet) -> Maybe Exp -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> NameSet
freeVars Maybe Exp
me NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (Exp -> NameSet) -> Inclusiveness Exp -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> NameSet
freeVars Inclusiveness Exp
incl
  Var QualName VName
qn (Info TypeBase (DimDecl VName) Aliasing
t) SrcLoc
_ -> Map VName StructType -> NameSet
NameSet (Map VName StructType -> NameSet)
-> Map VName StructType -> NameSet
forall a b. (a -> b) -> a -> b
$ VName -> StructType -> Map VName StructType
forall k a. k -> a -> Map k a
M.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn) (StructType -> Map VName StructType)
-> StructType -> Map VName StructType
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
t
  Ascript Exp
e TypeDeclBase Info VName
t SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Set VName -> NameSet
sizes (StructType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames (StructType -> Set VName) -> StructType -> Set VName
forall a b. (a -> b) -> a -> b
$ Info StructType -> StructType
forall a. Info a -> a
unInfo (Info StructType -> StructType) -> Info StructType -> StructType
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info StructType
forall (f :: * -> *) vn. TypeDeclBase f vn -> f StructType
expandedType TypeDeclBase Info VName
t)
  Coerce Exp
e TypeDeclBase Info VName
t (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Set VName -> NameSet
sizes (StructType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames (StructType -> Set VName) -> StructType -> Set VName
forall a b. (a -> b) -> a -> b
$ Info StructType -> StructType
forall a. Info a -> a
unInfo (Info StructType -> StructType) -> Info StructType -> StructType
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info StructType
forall (f :: * -> *) vn. TypeDeclBase f vn -> f StructType
expandedType TypeDeclBase Info VName
t)
  LetPat PatternBase Info VName
pat Exp
e1 Exp
e2 (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ ->
    Exp -> NameSet
freeVars Exp
e1
      NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> ( (Set VName -> NameSet
sizes (PatternBase Info VName -> Set VName
patternDimNames PatternBase Info VName
pat) NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e2)
             NameSet -> NameSet -> NameSet
`withoutM` PatternBase Info VName -> NameSet
patternVars PatternBase Info VName
pat
         )
  LetFun VName
vn ([TypeParamBase VName]
tparams, [PatternBase Info VName]
pats, Maybe (TypeExp VName)
_, Info StructType
_, Exp
e1) Exp
e2 Info (TypeBase (DimDecl VName) Aliasing)
_ SrcLoc
_ ->
    ( (Exp -> NameSet
freeVars Exp
e1 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Set VName -> NameSet
sizes ((PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
patternDimNames [PatternBase Info VName]
pats))
        NameSet -> Set VName -> NameSet
`without` ( (IdentBase Info VName -> VName)
-> Set (IdentBase Info VName) -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName ((PatternBase Info VName -> Set (IdentBase Info VName))
-> [PatternBase Info VName] -> Set (IdentBase Info VName)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase Info VName]
pats)
                      Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((TypeParamBase VName -> VName) -> [TypeParamBase VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams)
                  )
    )
      NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (Exp -> NameSet
freeVars Exp
e2 NameSet -> Set VName -> NameSet
`without` VName -> Set VName
forall a. a -> Set a
S.singleton VName
vn)
  If Exp
e1 Exp
e2 Exp
e3 (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e1 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e2 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e3
  Apply Exp
e1 Exp
e2 Info (Diet, Maybe VName)
_ (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e1 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e2
  Negate Exp
e SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e
  Lambda [PatternBase Info VName]
pats Exp
e0 Maybe (TypeExp VName)
_ (Info (Aliasing
_, StructType
t)) SrcLoc
_ ->
    (Set VName -> NameSet
sizes ((PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
patternDimNames [PatternBase Info VName]
pats) NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e0 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Set VName -> NameSet
sizes (StructType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames StructType
t))
      NameSet -> NameSet -> NameSet
`withoutM` (PatternBase Info VName -> NameSet)
-> [PatternBase Info VName] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> NameSet
patternVars [PatternBase Info VName]
pats
  OpSection {} -> NameSet
forall a. Monoid a => a
mempty
  OpSectionLeft QualName VName
_ Info (TypeBase (DimDecl VName) Aliasing)
_ Exp
e (Info (PName, StructType, Maybe VName), Info (PName, StructType))
_ (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e
  OpSectionRight QualName VName
_ Info (TypeBase (DimDecl VName) Aliasing)
_ Exp
e (Info (PName, StructType), Info (PName, StructType, Maybe VName))
_ Info (TypeBase (DimDecl VName) Aliasing)
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e
  ProjectSection {} -> NameSet
forall a. Monoid a => a
mempty
  IndexSection [DimIndexBase Info VName]
idxs Info (TypeBase (DimDecl VName) Aliasing)
_ SrcLoc
_ -> (DimIndexBase Info VName -> NameSet)
-> [DimIndexBase Info VName] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> NameSet
freeDimIndex [DimIndexBase Info VName]
idxs
  DoLoop [VName]
sparams PatternBase Info VName
pat Exp
e1 LoopFormBase Info VName
form Exp
e3 Info (TypeBase (DimDecl VName) Aliasing, [VName])
_ SrcLoc
_ ->
    let (NameSet
e2fv, NameSet
e2ident) = LoopFormBase Info VName -> (NameSet, NameSet)
formVars LoopFormBase Info VName
form
     in Exp -> NameSet
freeVars Exp
e1
          NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> ( (NameSet
e2fv NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e3)
                 NameSet -> NameSet -> NameSet
`withoutM` (Set VName -> NameSet
sizes ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
sparams) NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> PatternBase Info VName -> NameSet
patternVars PatternBase Info VName
pat NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> NameSet
e2ident)
             )
    where
      formVars :: LoopFormBase Info VName -> (NameSet, NameSet)
formVars (For IdentBase Info VName
v Exp
e2) = (Exp -> NameSet
freeVars Exp
e2, IdentBase Info VName -> NameSet
ident IdentBase Info VName
v)
      formVars (ForIn PatternBase Info VName
p Exp
e2) = (Exp -> NameSet
freeVars Exp
e2, PatternBase Info VName -> NameSet
patternVars PatternBase Info VName
p)
      formVars (While Exp
e2) = (Exp -> NameSet
freeVars Exp
e2, NameSet
forall a. Monoid a => a
mempty)
  BinOp (QualName VName
qn, SrcLoc
_) (Info TypeBase (DimDecl VName) Aliasing
qn_t) (Exp
e1, Info (StructType, Maybe VName)
_) (Exp
e2, Info (StructType, Maybe VName)
_) Info (TypeBase (DimDecl VName) Aliasing)
_ Info [VName]
_ SrcLoc
_ ->
    Map VName StructType -> NameSet
NameSet (VName -> StructType -> Map VName StructType
forall k a. k -> a -> Map k a
M.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn) (StructType -> Map VName StructType)
-> StructType -> Map VName StructType
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase (DimDecl VName) Aliasing
qn_t)
      NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e1
      NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e2
  Project Name
_ Exp
e Info (TypeBase (DimDecl VName) Aliasing)
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e
  LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs Exp
e1 Exp
e2 Info (TypeBase (DimDecl VName) Aliasing)
_ SrcLoc
_ ->
    IdentBase Info VName -> NameSet
ident IdentBase Info VName
id2 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (DimIndexBase Info VName -> NameSet)
-> [DimIndexBase Info VName] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> NameSet
freeDimIndex [DimIndexBase Info VName]
idxs NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e1
      NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (Exp -> NameSet
freeVars Exp
e2 NameSet -> Set VName -> NameSet
`without` VName -> Set VName
forall a. a -> Set a
S.singleton (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
id1))
  Index Exp
e [DimIndexBase Info VName]
idxs (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (DimIndexBase Info VName -> NameSet)
-> [DimIndexBase Info VName] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> NameSet
freeDimIndex [DimIndexBase Info VName]
idxs
  Update Exp
e1 [DimIndexBase Info VName]
idxs Exp
e2 SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e1 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (DimIndexBase Info VName -> NameSet)
-> [DimIndexBase Info VName] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> NameSet
freeDimIndex [DimIndexBase Info VName]
idxs NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e2
  RecordUpdate Exp
e1 [Name]
_ Exp
e2 Info (TypeBase (DimDecl VName) Aliasing)
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e1 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e2
  Assert Exp
e1 Exp
e2 Info String
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e1 NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
e2
  Constr Name
_ [Exp]
es Info (TypeBase (DimDecl VName) Aliasing)
_ SrcLoc
_ -> (Exp -> NameSet) -> [Exp] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> NameSet
freeVars [Exp]
es
  Attr AttrInfo
_ Exp
e SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e
  Match Exp
e NonEmpty (CaseBase Info VName)
cs (Info (TypeBase (DimDecl VName) Aliasing), Info [VName])
_ SrcLoc
_ -> Exp -> NameSet
freeVars Exp
e NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (CaseBase Info VName -> NameSet)
-> NonEmpty (CaseBase Info VName) -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap CaseBase Info VName -> NameSet
caseFV NonEmpty (CaseBase Info VName)
cs
    where
      caseFV :: CaseBase Info VName -> NameSet
caseFV (CasePat PatternBase Info VName
p Exp
eCase SrcLoc
_) =
        (Set VName -> NameSet
sizes (PatternBase Info VName -> Set VName
patternDimNames PatternBase Info VName
p) NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> Exp -> NameSet
freeVars Exp
eCase)
          NameSet -> NameSet -> NameSet
`withoutM` PatternBase Info VName -> NameSet
patternVars PatternBase Info VName
p

freeDimIndex :: DimIndexBase Info VName -> NameSet
freeDimIndex :: DimIndexBase Info VName -> NameSet
freeDimIndex (DimFix Exp
e) = Exp -> NameSet
freeVars Exp
e
freeDimIndex (DimSlice Maybe Exp
me1 Maybe Exp
me2 Maybe Exp
me3) =
  (Maybe Exp -> NameSet) -> [Maybe Exp] -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Exp -> NameSet) -> Maybe Exp -> NameSet
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> NameSet
freeVars) [Maybe Exp
me1, Maybe Exp
me2, Maybe Exp
me3]

-- | Extract all the variable names bound in a pattern.
patternVars :: Pattern -> NameSet
patternVars :: PatternBase Info VName -> NameSet
patternVars = [NameSet] -> NameSet
forall a. Monoid a => [a] -> a
mconcat ([NameSet] -> NameSet)
-> (PatternBase Info VName -> [NameSet])
-> PatternBase Info VName
-> NameSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IdentBase Info VName -> NameSet)
-> [IdentBase Info VName] -> [NameSet]
forall a b. (a -> b) -> [a] -> [b]
map IdentBase Info VName -> NameSet
ident ([IdentBase Info VName] -> [NameSet])
-> (PatternBase Info VName -> [IdentBase Info VName])
-> PatternBase Info VName
-> [NameSet]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set (IdentBase Info VName) -> [IdentBase Info VName]
forall a. Set a -> [a]
S.toList (Set (IdentBase Info VName) -> [IdentBase Info VName])
-> (PatternBase Info VName -> Set (IdentBase Info VName))
-> PatternBase Info VName
-> [IdentBase Info VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents