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

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

-- | A set of names where we also track their type.
newtype FV = FV {FV -> Map VName StructType
unFV :: M.Map VName StructType}
  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
showList :: [FV] -> ShowS
$cshowList :: [FV] -> ShowS
show :: FV -> String
$cshow :: FV -> String
showsPrec :: Int -> FV -> ShowS
$cshowsPrec :: Int -> FV -> ShowS
Show)

instance Semigroup FV where
  FV Map VName StructType
x <> :: FV -> FV -> FV
<> FV Map VName StructType
y = Map VName StructType -> FV
FV (Map VName StructType -> FV) -> Map VName StructType -> FV
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 FV where
  mempty :: FV
mempty = Map VName StructType -> FV
FV Map VName StructType
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 Map VName StructType
x) Set VName
y = Map VName StructType -> FV
FV (Map VName StructType -> FV) -> Map VName StructType -> FV
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

ident :: IdentBase Info VName -> FV
ident :: IdentBase Info VName -> FV
ident IdentBase Info VName
v = Map VName StructType -> FV
FV (Map VName StructType -> FV) -> Map VName StructType -> FV
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 Size Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (TypeBase Size Aliasing -> StructType)
-> TypeBase Size Aliasing -> StructType
forall a b. (a -> b) -> a -> b
$ Info (TypeBase Size Aliasing) -> TypeBase Size Aliasing
forall a. Info a -> a
unInfo (IdentBase Info VName -> Info (TypeBase Size Aliasing)
forall (f :: * -> *) vn.
IdentBase f vn -> f (TypeBase Size Aliasing)
identType IdentBase Info VName
v))

size :: VName -> FV
size :: VName -> FV
size VName
v = Map VName StructType -> FV
FV (Map VName StructType -> FV) -> Map VName StructType -> FV
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 Size () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase Size () -> StructType)
-> ScalarTypeBase Size () -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size ()
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase Size ())
-> PrimType -> ScalarTypeBase Size ()
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- | A 'FV' with these names, considered to be sizes.
sizes :: S.Set VName -> FV
sizes :: Set VName -> FV
sizes = (VName -> FV) -> Set VName -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap VName -> FV
size

-- | 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 (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 (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 (TypeBase Size Aliasing)
t SrcLoc
_) = IdentBase Info VName -> FV
ident (IdentBase Info VName -> FV) -> IdentBase Info VName -> FV
forall a b. (a -> b) -> a -> b
$ VName
-> Info (TypeBase Size Aliasing) -> SrcLoc -> IdentBase Info VName
forall (f :: * -> *) vn.
vn -> f (TypeBase Size Aliasing) -> SrcLoc -> IdentBase f vn
Ident VName
vn Info (TypeBase Size Aliasing)
t SrcLoc
forall a. Monoid a => a
mempty
  ArrayLit [ExpBase Info VName]
es Info (TypeBase Size Aliasing)
t SrcLoc
_ ->
    (ExpBase Info VName -> FV) -> [ExpBase Info VName] -> FV
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
<> Set VName -> FV
sizes (TypeBase Size Aliasing -> Set VName
forall as. TypeBase Size as -> Set VName
freeInType (TypeBase Size Aliasing -> Set VName)
-> TypeBase Size Aliasing -> Set VName
forall a b. (a -> b) -> a -> b
$ Info (TypeBase Size Aliasing) -> TypeBase Size Aliasing
forall a. Info a -> a
unInfo Info (TypeBase Size Aliasing)
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 (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 (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 TypeBase Size Aliasing
t) SrcLoc
_ -> Map VName StructType -> FV
FV (Map VName StructType -> FV) -> Map VName StructType -> FV
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 Size Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase Size Aliasing
t
  Ascript ExpBase Info VName
e TypeExp VName
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  AppExp (Coerce ExpBase Info VName
e TypeExp VName
_ SrcLoc
_) (Info AppRes
ar) ->
    ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Set VName -> FV
sizes (TypeBase Size Aliasing -> Set VName
forall as. TypeBase Size as -> Set VName
freeInType (AppRes -> TypeBase Size Aliasing
appResType AppRes
ar))
  AppExp (LetPat [SizeBinder VName]
let_sizes PatBase Info VName
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
<> ( (Set VName -> FV
sizes (PatBase Info VName -> Set VName
freeInPat PatBase Info VName
pat) FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2)
             FV -> Set VName -> FV
`freeWithout` (PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
pat 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 ((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]
pats, Maybe (TypeExp VName)
_, Info StructRetType
_, 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
<> Set VName -> FV
sizes ((PatBase Info VName -> Set VName)
-> [PatBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
freeInPat [PatBase Info VName]
pats))
        FV -> Set VName -> FV
`freeWithout` ( (PatBase Info VName -> Set VName)
-> [PatBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase 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)
                      )
    )
      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
e1 ExpBase Info VName
e2 Info (Diet, Maybe VName)
_ 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
  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]
pats ExpBase Info VName
e0 Maybe (TypeExp VName)
_ (Info (Aliasing
_, RetType [VName]
dims StructType
t)) SrcLoc
_ ->
    (Set VName -> FV
sizes ((PatBase Info VName -> Set VName)
-> [PatBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
freeInPat [PatBase Info VName]
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
<> Set VName -> FV
sizes (StructType -> Set VName
forall as. TypeBase Size as -> Set VName
freeInType StructType
t))
      FV -> Set VName -> FV
`freeWithout` ((PatBase Info VName -> Set VName)
-> [PatBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase 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 [VName]
dims)
  OpSection {} -> FV
forall a. Monoid a => a
mempty
  OpSectionLeft QualName VName
_ Info (TypeBase Size Aliasing)
_ ExpBase Info VName
e (Info (PName, StructType, Maybe VName), Info (PName, StructType))
_ (Info PatRetType, Info [VName])
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  OpSectionRight QualName VName
_ Info (TypeBase Size Aliasing)
_ ExpBase Info VName
e (Info (PName, StructType), Info (PName, StructType, Maybe VName))
_ Info PatRetType
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  ProjectSection {} -> FV
forall a. Monoid a => a
mempty
  IndexSection SliceBase Info VName
idxs Info (TypeBase Size Aliasing)
_ SrcLoc
_ -> (DimIndexBase Info VName -> FV) -> SliceBase Info VName -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> FV
freeInDimIndex SliceBase Info VName
idxs
  AppExp (DoLoop [VName]
sparams PatBase Info VName
pat ExpBase Info VName
e1 LoopFormBase Info VName
form ExpBase Info VName
e3 SrcLoc
_) Info AppRes
_ ->
    let (FV
e2fv, Set VName
e2ident) = LoopFormBase Info VName -> (FV, Set 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 -> Set VName -> FV
`freeWithout` ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
sparams Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
pat Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
e2ident)
             )
    where
      formVars :: LoopFormBase Info VName -> (FV, Set VName)
formVars (For IdentBase Info VName
v ExpBase Info VName
e2) = (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2, VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
v)
      formVars (ForIn PatBase Info VName
p ExpBase Info VName
e2) = (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2, PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
p)
      formVars (While ExpBase Info VName
e2) = (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2, Set VName
forall a. Monoid a => a
mempty)
  AppExp (BinOp (QualName VName
qn, SrcLoc
_) (Info TypeBase Size Aliasing
qn_t) (ExpBase Info VName
e1, Info (StructType, Maybe VName)
_) (ExpBase Info VName
e2, Info (StructType, Maybe VName)
_) SrcLoc
_) Info AppRes
_ ->
    Map VName StructType -> FV
FV (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 Size Aliasing -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct TypeBase Size Aliasing
qn_t)
      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 (TypeBase Size Aliasing)
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  AppExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 SliceBase Info VName
idxs ExpBase Info VName
e1 ExpBase Info VName
e2 SrcLoc
_) Info AppRes
_ ->
    IdentBase Info VName -> FV
ident IdentBase Info VName
id2
      FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (DimIndexBase Info VName -> FV) -> SliceBase Info VName -> FV
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 (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
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 (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 (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 (TypeBase Size Aliasing)
_ 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 String
_ 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 (TypeBase Size Aliasing)
_ SrcLoc
_ -> (ExpBase Info VName -> FV) -> [ExpBase Info VName] -> FV
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 (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
p ExpBase Info VName
eCase SrcLoc
_) =
        (Set VName -> FV
sizes (PatBase Info VName -> Set VName
freeInPat PatBase Info VName
p) FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
eCase)
          FV -> Set VName -> FV
`freeWithout` PatBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
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 (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((ExpBase Info VName -> FV) -> Maybe (ExpBase Info VName) -> FV
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 :: PatBase Info VName -> S.Set VName
freeInPat :: PatBase Info VName -> Set VName
freeInPat (TuplePat [PatBase Info VName]
ps SrcLoc
_) = (PatBase Info VName -> Set VName)
-> [PatBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
freeInPat [PatBase Info VName]
ps
freeInPat (RecordPat [(Name, PatBase Info VName)]
fs SrcLoc
_) = ((Name, PatBase Info VName) -> Set VName)
-> [(Name, PatBase Info VName)] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatBase Info VName -> Set VName
freeInPat (PatBase Info VName -> Set VName)
-> ((Name, PatBase Info VName) -> PatBase Info VName)
-> (Name, PatBase Info VName)
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatBase Info VName) -> PatBase Info VName
forall a b. (a, b) -> b
snd) [(Name, PatBase Info VName)]
fs
freeInPat (PatParens PatBase Info VName
p SrcLoc
_) = PatBase Info VName -> Set VName
freeInPat PatBase Info VName
p
freeInPat (Id VName
_ (Info TypeBase Size Aliasing
tp) SrcLoc
_) = TypeBase Size Aliasing -> Set VName
forall as. TypeBase Size as -> Set VName
freeInType TypeBase Size Aliasing
tp
freeInPat (Wildcard (Info TypeBase Size Aliasing
tp) SrcLoc
_) = TypeBase Size Aliasing -> Set VName
forall as. TypeBase Size as -> Set VName
freeInType TypeBase Size Aliasing
tp
freeInPat (PatAscription PatBase Info VName
p TypeExp VName
_ SrcLoc
_) = PatBase Info VName -> Set VName
freeInPat PatBase Info VName
p
freeInPat (PatLit PatLit
_ (Info TypeBase Size Aliasing
tp) SrcLoc
_) = TypeBase Size Aliasing -> Set VName
forall as. TypeBase Size as -> Set VName
freeInType TypeBase Size Aliasing
tp
freeInPat (PatConstr Name
_ Info (TypeBase Size Aliasing)
_ [PatBase Info VName]
ps SrcLoc
_) = (PatBase Info VName -> Set VName)
-> [PatBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
freeInPat [PatBase Info VName]
ps
freeInPat (PatAttr AttrInfo VName
_ PatBase Info VName
p SrcLoc
_) = PatBase Info VName -> Set VName
freeInPat PatBase Info VName
p

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

    onSize :: Size -> Set VName
onSize (NamedSize QualName VName
qn) = 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
    onSize Size
_ = Set VName
forall a. Monoid a => a
mempty