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

import Data.Map.Strict qualified as M
import Data.Set qualified as S
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
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 forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith 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 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 forall a b. (a -> b) -> a -> b
$ forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey forall {p}. VName -> p -> Bool
keep Map VName StructType
x
  where
    keep :: VName -> p -> Bool
keep VName
k p
_ = VName
k 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 forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton (forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
v) (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo (forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType IdentBase Info VName
v))

size :: VName -> FV
size :: VName -> FV
size VName
v = Map VName StructType -> FV
FV forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton VName
v forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim 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 = 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 {} -> forall a. Monoid a => a
mempty
  IntLit {} -> forall a. Monoid a => a
mempty
  FloatLit {} -> forall a. Monoid a => a
mempty
  StringLit {} -> forall a. Monoid a => a
mempty
  Hole {} -> 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
_ -> 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
_ -> 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 PatType
t SrcLoc
_) = IdentBase Info VName -> FV
ident forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
vn -> f PatType -> SrcLoc -> IdentBase f vn
Ident VName
vn Info PatType
t forall a. Monoid a => a
mempty
  ArrayLit [ExpBase Info VName]
es Info PatType
t SrcLoc
_ ->
    forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp [ExpBase Info VName]
es forall a. Semigroup a => a -> a -> a
<> Set VName -> FV
sizes (forall as. TypeBase Size as -> Set VName
freeInType forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo Info PatType
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 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ExpBase Info VName -> FV
freeInExp Maybe (ExpBase Info VName)
me forall a. Semigroup a => a -> a -> a
<> 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 PatType
t) SrcLoc
_ -> Map VName StructType -> FV
FV forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn) forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
  Ascript ExpBase Info VName
e TypeExp Info VName
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e
  AppExp (Coerce ExpBase Info VName
e TypeExp Info VName
_ SrcLoc
_) (Info AppRes
ar) ->
    ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e forall a. Semigroup a => a -> a -> a
<> Set VName -> FV
sizes (forall as. TypeBase Size as -> Set VName
freeInType (AppRes -> PatType
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
      forall a. Semigroup a => a -> a -> a
<> ( (Set VName -> FV
sizes (PatBase Info VName -> Set VName
freeInPat PatBase Info VName
pat) forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2)
             FV -> Set VName -> FV
`freeWithout` (forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
pat forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList (forall a b. (a -> b) -> [a] -> [b]
map forall vn. SizeBinder vn -> vn
sizeName [SizeBinder VName]
let_sizes))
         )
  AppExp (LetFun VName
vn ([TypeParamBase VName]
tparams, [PatBase Info VName]
pats, Maybe (TypeExp Info VName)
_, Info StructRetType
_, ExpBase Info VName
e1) ExpBase Info VName
e2 SrcLoc
_) Info AppRes
_ ->
    ( (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1 forall a. Semigroup a => a -> a -> a
<> Set VName -> FV
sizes (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` ( forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase Info VName]
pats
                          forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams)
                      )
    )
      forall a. Semigroup a => a -> a -> a
<> (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2 FV -> Set VName -> FV
`freeWithout` 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 forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2 forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e3
  AppExp (Apply ExpBase Info VName
f NonEmpty (Info (Diet, Maybe VName), ExpBase Info VName)
args SrcLoc
_) Info AppRes
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
f forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (ExpBase Info VName -> FV
freeInExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) NonEmpty (Info (Diet, 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]
pats ExpBase Info VName
e0 Maybe (TypeExp Info VName)
_ (Info (Aliasing
_, RetType [VName]
dims StructType
t)) SrcLoc
_ ->
    (Set VName -> FV
sizes (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
freeInPat [PatBase Info VName]
pats) forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e0 forall a. Semigroup a => a -> a -> a
<> Set VName -> FV
sizes (forall as. TypeBase Size as -> Set VName
freeInType StructType
t))
      FV -> Set VName -> FV
`freeWithout` (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase Info VName]
pats forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims)
  OpSection {} -> forall a. Monoid a => a
mempty
  OpSectionLeft QualName VName
_ Info PatType
_ 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 PatType
_ 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 {} -> forall a. Monoid a => a
mempty
  IndexSection SliceBase Info VName
idxs Info PatType
_ SrcLoc
_ -> 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
          forall a. Semigroup a => a -> a -> a
<> ( (FV
e2fv forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e3)
                 FV -> Set VName -> FV
`freeWithout` (forall a. Ord a => [a] -> Set a
S.fromList [VName]
sparams forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames PatBase Info VName
pat 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, forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ 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, 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, forall a. Monoid a => a
mempty)
  AppExp (BinOp (QualName VName
qn, SrcLoc
_) (Info PatType
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 (forall k a. k -> a -> Map k a
M.singleton (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn) forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
qn_t)
      forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1
      forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2
  Project Name
_ ExpBase Info VName
e Info PatType
_ 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
      forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> FV
freeInDimIndex SliceBase Info VName
idxs
      forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1
      forall a. Semigroup a => a -> a -> a
<> (ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2 FV -> Set VName -> FV
`freeWithout` forall a. a -> Set a
S.singleton (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 forall a. Semigroup a => a -> a -> a
<> 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 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimIndexBase Info VName -> FV
freeInDimIndex SliceBase Info VName
idxs 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 PatType
_ SrcLoc
_ -> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e1 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 forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
e2
  Constr Name
_ [ExpBase Info VName]
es Info PatType
_ SrcLoc
_ -> 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 forall a. Semigroup a => a -> a -> a
<> 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) forall a. Semigroup a => a -> a -> a
<> ExpBase Info VName -> FV
freeInExp ExpBase Info VName
eCase)
          FV -> Set VName -> FV
`freeWithout` 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) =
  forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (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
_) = 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
_) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatBase Info VName -> Set VName
freeInPat forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 PatType
tp) SrcLoc
_) = forall as. TypeBase Size as -> Set VName
freeInType PatType
tp
freeInPat (Wildcard (Info PatType
tp) SrcLoc
_) = forall as. TypeBase Size as -> Set VName
freeInType PatType
tp
freeInPat (PatAscription PatBase Info VName
p TypeExp Info VName
_ SrcLoc
_) = PatBase Info VName -> Set VName
freeInPat PatBase Info VName
p
freeInPat (PatLit PatLit
_ (Info PatType
tp) SrcLoc
_) = forall as. TypeBase Size as -> Set VName
freeInType PatType
tp
freeInPat (PatConstr Name
_ Info PatType
_ [PatBase Info VName]
ps SrcLoc
_) = 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 :: forall as. 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 ->
      forall as. TypeBase Size as -> Set VName
freeInType (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase Size ()
a) forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Size -> Set VName
onSize (forall dim. Shape dim -> [dim]
shapeDims Shape Size
s)
    Scalar (Record Map Name (TypeBase Size as)
fs) ->
      forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall as. TypeBase Size as -> Set VName
freeInType Map Name (TypeBase Size as)
fs
    Scalar Prim {} ->
      forall a. Monoid a => a
mempty
    Scalar (Sum Map Name [TypeBase Size as]
cs) ->
      forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall as. TypeBase Size as -> Set VName
freeInType) Map Name [TypeBase Size as]
cs
    Scalar (Arrow as
_ PName
v Diet
_ StructType
t1 (RetType [VName]
dims TypeBase Size as
t2)) ->
      forall a. (a -> Bool) -> Set a -> Set a
S.filter (PName -> VName -> Bool
notV PName
v) forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> Set a -> Set a
S.filter (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
dims) forall a b. (a -> b) -> a -> b
$ forall as. TypeBase Size as -> Set VName
freeInType StructType
t1 forall a. Semigroup a => a -> a -> a
<> forall as. TypeBase Size as -> Set VName
freeInType TypeBase Size as
t2
    Scalar (TypeVar as
_ Uniqueness
_ QualName VName
_ [TypeArg Size]
targs) ->
      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
_) = forall as. TypeBase Size as -> Set VName
freeInType StructType
at

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

    onSize :: Size -> Set VName
onSize (NamedSize QualName VName
qn) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
    onSize Size
_ = forall a. Monoid a => a
mempty