-- | Facilities for computing free term variables in various syntactic
-- constructs.
module Language.Futhark.FreeVars
  ( freeInExp,
    freeInPat,
    freeInType,
    freeWithout,
    FV,
    fvVars,
  )
where

import Data.Set qualified as S
import Language.Futhark.Prop
import Language.Futhark.Syntax

-- | A set of names.
newtype FV = FV {FV -> Set VName
unFV :: S.Set VName}
  deriving (Int -> FV -> ShowS
[FV] -> ShowS
FV -> String
(Int -> FV -> ShowS)
-> (FV -> String) -> ([FV] -> ShowS) -> Show FV
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> FV -> ShowS
showsPrec :: Int -> FV -> ShowS
$cshow :: FV -> String
show :: FV -> String
$cshowList :: [FV] -> ShowS
showList :: [FV] -> ShowS
Show)

-- | The set of names in an 'FV'.
fvVars :: FV -> S.Set VName
fvVars :: FV -> Set VName
fvVars = FV -> Set VName
unFV

instance Semigroup FV where
  FV Set VName
x <> :: FV -> FV -> FV
<> FV Set VName
y = Set VName -> FV
FV (Set VName -> FV) -> Set VName -> FV
forall a b. (a -> b) -> a -> b
$ Set VName
x Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
y

instance Monoid FV where
  mempty :: FV
mempty = Set VName -> FV
FV Set VName
forall a. Monoid a => a
mempty

-- | Set subtraction.  Do not consider those variables as free.
freeWithout :: FV -> S.Set VName -> FV
freeWithout :: FV -> Set VName -> FV
freeWithout (FV Set VName
x) Set VName
y = Set VName -> FV
FV (Set VName -> FV) -> Set VName -> FV
forall a b. (a -> b) -> a -> b
$ Set VName
x Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Set VName
y

-- | As 'freeWithout', but for lists.
freeWithoutL :: FV -> [VName] -> FV
freeWithoutL :: FV -> [VName] -> FV
freeWithoutL FV
fv [VName]
y = FV
fv FV -> Set VName -> FV
`freeWithout` [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
y

ident :: Ident t -> FV
ident :: forall t. Ident t -> FV
ident = Set VName -> FV
FV (Set VName -> FV) -> (Ident t -> Set VName) -> Ident t -> FV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> (Ident t -> VName) -> Ident t -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident t -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName

-- | Compute the set of free variables of an expression.
freeInExp :: ExpBase Info VName -> FV
freeInExp :: ExpBase Info VName -> FV
freeInExp ExpBase Info VName
expr = case ExpBase Info VName
expr of
  Literal {} -> FV
forall a. Monoid a => a
mempty
  IntLit {} -> FV
forall a. Monoid a => a
mempty
  FloatLit {} -> FV
forall a. Monoid a => a
mempty
  StringLit {} -> FV
forall a. Monoid a => a
mempty
  Hole {} -> FV
forall a. Monoid a => a
mempty
  Parens ExpBase Info VName
e SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  QualParens (QualName VName, SrcLoc)
_ ExpBase Info VName
e SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  TupLit [ExpBase Info VName]
es SrcLoc
_ -> (ExpBase Info VName -> FV) -> [ExpBase Info VName] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp [ExpBase Info VName]
es
  RecordLit [FieldBase Info VName]
fs SrcLoc
_ -> (FieldBase Info VName -> FV) -> [FieldBase Info VName] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap FieldBase Info VName -> FV
freeInExpField [FieldBase Info VName]
fs
    where
      freeInExpField :: FieldBase Info VName -> FV
freeInExpField (RecordFieldExplicit Name
_ ExpBase Info VName
e SrcLoc
_) = ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
      freeInExpField (RecordFieldImplicit VName
vn Info StructType
t SrcLoc
_) = Ident StructType -> FV
forall t. Ident t -> FV
ident (Ident StructType -> FV) -> Ident StructType -> FV
forall a b. (a -> b) -> a -> b
$ VName -> Info StructType -> SrcLoc -> Ident StructType
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
Ident VName
vn Info StructType
t SrcLoc
forall a. Monoid a => a
mempty
  ArrayLit [ExpBase Info VName]
es Info StructType
t SrcLoc
_ ->
    (ExpBase Info VName -> FV) -> [ExpBase Info VName] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp [ExpBase Info VName]
es FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> StructType -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType (Info StructType -> StructType
forall a. Info a -> a
unInfo Info StructType
t)
  AppExp (Range ExpBase Info VName
e Maybe (ExpBase Info VName)
me Inclusiveness (ExpBase Info VName)
incl SrcLoc
_) Info AppRes
_ ->
    ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (ExpBase Info VName -> FV) -> Maybe (ExpBase Info VName) -> FV
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp Maybe (ExpBase Info VName)
me FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (ExpBase Info VName -> FV)
-> Inclusiveness (ExpBase Info VName) -> FV
forall m a. Monoid m => (a -> m) -> Inclusiveness a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp Inclusiveness (ExpBase Info VName)
incl
  Var QualName VName
qn Info StructType
_ SrcLoc
_ -> Set VName -> FV
FV (Set VName -> FV) -> Set VName -> FV
forall a b. (a -> b) -> a -> b
$ VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
  Ascript ExpBase Info VName
e TypeExp (ExpBase Info VName) VName
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  Coerce ExpBase Info VName
e TypeExp (ExpBase Info VName) VName
_ (Info StructType
t) SrcLoc
_ ->
    ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> StructType -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType StructType
t
  AppExp (LetPat [SizeBinder VName]
let_sizes PatBase Info VName StructType
pat ExpBase Info VName
e1 ExpBase Info VName
e2 SrcLoc
_) Info AppRes
_ ->
    ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ( (PatBase Info VName StructType -> FV
forall u. Pat (TypeBase (ExpBase Info VName) u) -> FV
freeInPat PatBase Info VName StructType
pat FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2)
             FV -> [VName] -> FV
`freeWithoutL` (PatBase Info VName StructType -> [VName]
forall t. Pat t -> [VName]
patNames PatBase Info VName StructType
pat [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> (SizeBinder VName -> VName) -> [SizeBinder VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map SizeBinder VName -> VName
forall vn. SizeBinder vn -> vn
sizeName [SizeBinder VName]
let_sizes)
         )
  AppExp (LetFun VName
vn ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
pats, Maybe (TypeExp (ExpBase Info VName) VName)
_, Info ResRetType
_, ExpBase Info VName
e1) ExpBase Info VName
e2 SrcLoc
_) Info AppRes
_ ->
    ( (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (PatBase Info VName ParamType -> FV)
-> [PatBase Info VName ParamType] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName ParamType -> FV
forall u. Pat (TypeBase (ExpBase Info VName) u) -> FV
freeInPat [PatBase Info VName ParamType]
pats)
        FV -> [VName] -> FV
`freeWithoutL` ((PatBase Info VName ParamType -> [VName])
-> [PatBase Info VName ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [PatBase Info VName ParamType]
pats [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> (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)
    )
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2 FV -> Set VName -> FV
`freeWithout` VName -> Set VName
forall a. a -> Set a
S.singleton VName
vn)
  AppExp (If ExpBase Info VName
e1 ExpBase Info VName
e2 ExpBase Info VName
e3 SrcLoc
_) Info AppRes
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e3
  AppExp (Apply ExpBase Info VName
f NonEmpty (Info (Maybe VName), ExpBase Info VName)
args SrcLoc
_) Info AppRes
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
f FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ((Info (Maybe VName), ExpBase Info VName) -> FV)
-> NonEmpty (Info (Maybe VName), ExpBase Info VName) -> FV
forall m a. Monoid m => (a -> m) -> NonEmpty a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (ExpBase Info VName -> FV
freeInExp (ExpBase Info VName -> FV)
-> ((Info (Maybe VName), ExpBase Info VName) -> ExpBase Info VName)
-> (Info (Maybe VName), ExpBase Info VName)
-> FV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Info (Maybe VName), ExpBase Info VName) -> ExpBase Info VName
forall a b. (a, b) -> b
snd) NonEmpty (Info (Maybe VName), ExpBase Info VName)
args
  Negate ExpBase Info VName
e SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  Not ExpBase Info VName
e SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  Lambda [PatBase Info VName ParamType]
pats ExpBase Info VName
e0 Maybe (TypeExp (ExpBase Info VName) VName)
_ (Info (RetType [VName]
dims TypeBase (ExpBase Info VName) Uniqueness
t)) SrcLoc
_ ->
    ((PatBase Info VName ParamType -> FV)
-> [PatBase Info VName ParamType] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName ParamType -> FV
forall u. Pat (TypeBase (ExpBase Info VName) u) -> FV
freeInPat [PatBase Info VName ParamType]
pats FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e0 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> TypeBase (ExpBase Info VName) Uniqueness -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType TypeBase (ExpBase Info VName) Uniqueness
t)
      FV -> [VName] -> FV
`freeWithoutL` ((PatBase Info VName ParamType -> [VName])
-> [PatBase Info VName ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [PatBase Info VName ParamType]
pats [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dims)
  OpSection {} -> FV
forall a. Monoid a => a
mempty
  OpSectionLeft QualName VName
_ Info StructType
_ ExpBase Info VName
e (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
_ (Info ResRetType, Info [VName])
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  OpSectionRight QualName VName
_ Info StructType
_ ExpBase Info VName
e (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
_ Info ResRetType
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  ProjectSection {} -> FV
forall a. Monoid a => a
mempty
  IndexSection SliceBase Info VName
idxs Info StructType
_ SrcLoc
_ -> (DimIndexBase Info VName -> FV) -> SliceBase Info VName -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> FV
freeInDimIndex SliceBase Info VName
idxs
  AppExp (Loop [VName]
sparams PatBase Info VName ParamType
pat ExpBase Info VName
e1 LoopFormBase Info VName
form ExpBase Info VName
e3 SrcLoc
_) Info AppRes
_ ->
    let (FV
e2fv, [VName]
e2ident) = LoopFormBase Info VName -> (FV, [VName])
formVars LoopFormBase Info VName
form
     in ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1
          FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ( (FV
e2fv FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e3)
                 FV -> [VName] -> FV
`freeWithoutL` ([VName]
sparams [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> PatBase Info VName ParamType -> [VName]
forall t. Pat t -> [VName]
patNames PatBase Info VName ParamType
pat [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
e2ident)
             )
    where
      formVars :: LoopFormBase Info VName -> (FV, [VName])
formVars (For Ident StructType
v ExpBase Info VName
e2) = (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2, [Ident StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName Ident StructType
v])
      formVars (ForIn PatBase Info VName StructType
p ExpBase Info VName
e2) = (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2, PatBase Info VName StructType -> [VName]
forall t. Pat t -> [VName]
patNames PatBase Info VName StructType
p)
      formVars (While ExpBase Info VName
e2) = (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2, [VName]
forall a. Monoid a => a
mempty)
  AppExp (BinOp (QualName VName
qn, SrcLoc
_) Info StructType
_ (ExpBase Info VName
e1, Info (Maybe VName)
_) (ExpBase Info VName
e2, Info (Maybe VName)
_) SrcLoc
_) Info AppRes
_ ->
    Set VName -> FV
FV (VName -> Set VName
forall a. a -> Set a
S.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn))
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2
  Project Name
_ ExpBase Info VName
e Info StructType
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  AppExp (LetWith Ident StructType
id1 Ident StructType
id2 SliceBase Info VName
idxs ExpBase Info VName
e1 ExpBase Info VName
e2 SrcLoc
_) Info AppRes
_ ->
    Ident StructType -> FV
forall t. Ident t -> FV
ident Ident StructType
id2
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (DimIndexBase Info VName -> FV) -> SliceBase Info VName -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> FV
freeInDimIndex SliceBase Info VName
idxs
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2 FV -> Set VName -> FV
`freeWithout` VName -> Set VName
forall a. a -> Set a
S.singleton (Ident StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName Ident StructType
id1))
  AppExp (Index ExpBase Info VName
e SliceBase Info VName
idxs SrcLoc
_) Info AppRes
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (DimIndexBase Info VName -> FV) -> SliceBase Info VName -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> FV
freeInDimIndex SliceBase Info VName
idxs
  Update ExpBase Info VName
e1 SliceBase Info VName
idxs ExpBase Info VName
e2 SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (DimIndexBase Info VName -> FV) -> SliceBase Info VName -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> FV
freeInDimIndex SliceBase Info VName
idxs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2
  RecordUpdate ExpBase Info VName
e1 [Name]
_ ExpBase Info VName
e2 Info StructType
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2
  Assert ExpBase Info VName
e1 ExpBase Info VName
e2 Info Text
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2
  Constr Name
_ [ExpBase Info VName]
es Info StructType
_ SrcLoc
_ -> (ExpBase Info VName -> FV) -> [ExpBase Info VName] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp [ExpBase Info VName]
es
  Attr AttrInfo VName
_ ExpBase Info VName
e SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  AppExp (Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs SrcLoc
_) Info AppRes
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (CaseBase Info VName -> FV) -> NonEmpty (CaseBase Info VName) -> FV
forall m a. Monoid m => (a -> m) -> NonEmpty a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap CaseBase Info VName -> FV
caseFV NonEmpty (CaseBase Info VName)
cs
    where
      caseFV :: CaseBase Info VName -> FV
caseFV (CasePat PatBase Info VName StructType
p ExpBase Info VName
eCase SrcLoc
_) =
        (PatBase Info VName StructType -> FV
forall u. Pat (TypeBase (ExpBase Info VName) u) -> FV
freeInPat PatBase Info VName StructType
p FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
eCase)
          FV -> [VName] -> FV
`freeWithoutL` PatBase Info VName StructType -> [VName]
forall t. Pat t -> [VName]
patNames PatBase Info VName StructType
p

freeInDimIndex :: DimIndexBase Info VName -> FV
freeInDimIndex :: DimIndexBase Info VName -> FV
freeInDimIndex (DimFix ExpBase Info VName
e) = ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
freeInDimIndex (DimSlice Maybe (ExpBase Info VName)
me1 Maybe (ExpBase Info VName)
me2 Maybe (ExpBase Info VName)
me3) =
  (Maybe (ExpBase Info VName) -> FV)
-> [Maybe (ExpBase Info VName)] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((ExpBase Info VName -> FV) -> Maybe (ExpBase Info VName) -> FV
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp) [Maybe (ExpBase Info VName)
me1, Maybe (ExpBase Info VName)
me2, Maybe (ExpBase Info VName)
me3]

-- | Free variables in pattern (including types of the bound identifiers).
freeInPat :: Pat (TypeBase Size u) -> FV
freeInPat :: forall u. Pat (TypeBase (ExpBase Info VName) u) -> FV
freeInPat = (TypeBase (ExpBase Info VName) u -> FV)
-> PatBase Info VName (TypeBase (ExpBase Info VName) u) -> FV
forall m a. Monoid m => (a -> m) -> PatBase Info VName a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase (ExpBase Info VName) u -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType

-- | Free variables in the type (meaning those that are used in size expression).
freeInType :: TypeBase Size u -> FV
freeInType :: forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType TypeBase (ExpBase Info VName) u
t =
  case TypeBase (ExpBase Info VName) u
t of
    Array u
_ Shape (ExpBase Info VName)
s ScalarTypeBase (ExpBase Info VName) NoUniqueness
a ->
      StructType -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType (ScalarTypeBase (ExpBase Info VName) NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar ScalarTypeBase (ExpBase Info VName) NoUniqueness
a) FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (ExpBase Info VName -> FV) -> [ExpBase Info VName] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp (Shape (ExpBase Info VName) -> [ExpBase Info VName]
forall dim. Shape dim -> [dim]
shapeDims Shape (ExpBase Info VName)
s)
    Scalar (Record Map Name (TypeBase (ExpBase Info VName) u)
fs) ->
      (TypeBase (ExpBase Info VName) u -> FV)
-> Map Name (TypeBase (ExpBase Info VName) u) -> FV
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase (ExpBase Info VName) u -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType Map Name (TypeBase (ExpBase Info VName) u)
fs
    Scalar Prim {} ->
      FV
forall a. Monoid a => a
mempty
    Scalar (Sum Map Name [TypeBase (ExpBase Info VName) u]
cs) ->
      ([TypeBase (ExpBase Info VName) u] -> FV)
-> Map Name [TypeBase (ExpBase Info VName) u] -> FV
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((TypeBase (ExpBase Info VName) u -> FV)
-> [TypeBase (ExpBase Info VName) u] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeBase (ExpBase Info VName) u -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType) Map Name [TypeBase (ExpBase Info VName) u]
cs
    Scalar (Arrow u
_ PName
v Diet
_ StructType
t1 (RetType [VName]
dims TypeBase (ExpBase Info VName) Uniqueness
t2)) ->
      Set VName -> FV
FV (Set VName -> FV) -> (Set VName -> Set VName) -> Set VName -> FV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Bool) -> Set VName -> Set VName
forall a. (a -> Bool) -> Set a -> Set a
S.filter (\VName
k -> PName -> VName -> Bool
notV PName
v VName
k Bool -> Bool -> Bool
&& VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
notElem VName
k [VName]
dims) (Set VName -> FV) -> Set VName -> FV
forall a b. (a -> b) -> a -> b
$
        FV -> Set VName
unFV (StructType -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType StructType
t1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> TypeBase (ExpBase Info VName) Uniqueness -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType TypeBase (ExpBase Info VName) Uniqueness
t2)
    Scalar (TypeVar u
_ QualName VName
_ [TypeArg (ExpBase Info VName)]
targs) ->
      (TypeArg (ExpBase Info VName) -> FV)
-> [TypeArg (ExpBase Info VName)] -> FV
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeArg (ExpBase Info VName) -> FV
typeArgDims [TypeArg (ExpBase Info VName)]
targs
  where
    typeArgDims :: TypeArg (ExpBase Info VName) -> FV
typeArgDims (TypeArgDim ExpBase Info VName
d) = ExpBase Info VName -> FV
freeInExp ExpBase Info VName
d
    typeArgDims (TypeArgType StructType
at) = StructType -> FV
forall u. TypeBase (ExpBase Info VName) u -> FV
freeInType StructType
at

    notV :: PName -> VName -> Bool
notV PName
Unnamed = Bool -> VName -> Bool
forall a b. a -> b -> a
const Bool
True
    notV (Named VName
v) = (VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
v)