{- Language/Haskell/TH/Desugar/FV.hs

(c) Ryan Scott 2018

Compute free variables of programs.
-}

{-# LANGUAGE CPP #-}
module Language.Haskell.TH.Desugar.FV
  ( fvDType
  , extractBoundNamesDPat
  ) where

#if __GLASGOW_HASKELL__ < 710
import Data.Foldable (foldMap)
#endif
#if __GLASGOW_HASKELL__ < 804
import Data.Monoid ((<>))
#endif
import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Desugar.AST
import qualified Language.Haskell.TH.Desugar.OSet as OS
import Language.Haskell.TH.Desugar.OSet (OSet)

-- | Compute the free variables of a 'DType'.
fvDType :: DType -> OSet Name
fvDType :: DType -> OSet Name
fvDType = DType -> OSet Name
go
  where
    go :: DType -> OSet Name
    go :: DType -> OSet Name
go (DForallT DForallTelescope
tele DType
ty)      = DForallTelescope -> OSet Name -> OSet Name
fv_dtele DForallTelescope
tele (DType -> OSet Name
go DType
ty)
    go (DConstrainedT DCxt
ctxt DType
ty) = (DType -> OSet Name) -> DCxt -> OSet Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DType -> OSet Name
fvDType DCxt
ctxt OSet Name -> OSet Name -> OSet Name
forall a. Semigroup a => a -> a -> a
<> DType -> OSet Name
go DType
ty
    go (DAppT DType
t1 DType
t2)           = DType -> OSet Name
go DType
t1 OSet Name -> OSet Name -> OSet Name
forall a. Semigroup a => a -> a -> a
<> DType -> OSet Name
go DType
t2
    go (DAppKindT DType
t DType
k)         = DType -> OSet Name
go DType
t OSet Name -> OSet Name -> OSet Name
forall a. Semigroup a => a -> a -> a
<> DType -> OSet Name
go DType
k
    go (DSigT DType
ty DType
ki)           = DType -> OSet Name
go DType
ty OSet Name -> OSet Name -> OSet Name
forall a. Semigroup a => a -> a -> a
<> DType -> OSet Name
go DType
ki
    go (DVarT Name
n)               = Name -> OSet Name
forall a. a -> OSet a
OS.singleton Name
n
    go (DConT {})              = OSet Name
forall a. OSet a
OS.empty
    go DType
DArrowT                 = OSet Name
forall a. OSet a
OS.empty
    go (DLitT {})              = OSet Name
forall a. OSet a
OS.empty
    go DType
DWildCardT              = OSet Name
forall a. OSet a
OS.empty

-----
-- Extracting bound term names
-----

-- | Extract the term variables bound by a 'DPat'.
--
-- This does /not/ extract any type variables bound by pattern signatures.
extractBoundNamesDPat :: DPat -> OSet Name
extractBoundNamesDPat :: DPat -> OSet Name
extractBoundNamesDPat = DPat -> OSet Name
go
  where
    go :: DPat -> OSet Name
    go :: DPat -> OSet Name
go (DLitP Lit
_)          = OSet Name
forall a. OSet a
OS.empty
    go (DVarP Name
n)          = Name -> OSet Name
forall a. a -> OSet a
OS.singleton Name
n
    go (DConP Name
_ DCxt
tys [DPat]
pats) = (DType -> OSet Name) -> DCxt -> OSet Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DType -> OSet Name
fvDType DCxt
tys OSet Name -> OSet Name -> OSet Name
forall a. Semigroup a => a -> a -> a
<> (DPat -> OSet Name) -> [DPat] -> OSet Name
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DPat -> OSet Name
go [DPat]
pats
    go (DTildeP DPat
p)        = DPat -> OSet Name
go DPat
p
    go (DBangP DPat
p)         = DPat -> OSet Name
go DPat
p
    go (DSigP DPat
p DType
_)        = DPat -> OSet Name
go DPat
p
    go DPat
DWildP             = OSet Name
forall a. OSet a
OS.empty

-----
-- Binding forms
-----

-- | Adjust the free variables of something following a 'DForallTelescope'.
fv_dtele :: DForallTelescope -> OSet Name -> OSet Name
fv_dtele :: DForallTelescope -> OSet Name -> OSet Name
fv_dtele (DForallVis   [DTyVarBndrUnit]
tvbs) = [DTyVarBndrUnit] -> OSet Name -> OSet Name
forall flag. [DTyVarBndr flag] -> OSet Name -> OSet Name
fv_dtvbs [DTyVarBndrUnit]
tvbs
fv_dtele (DForallInvis [DTyVarBndrSpec]
tvbs) = [DTyVarBndrSpec] -> OSet Name -> OSet Name
forall flag. [DTyVarBndr flag] -> OSet Name -> OSet Name
fv_dtvbs [DTyVarBndrSpec]
tvbs

-- | Adjust the free variables of something following 'DTyVarBndr's.
fv_dtvbs :: [DTyVarBndr flag] -> OSet Name -> OSet Name
fv_dtvbs :: [DTyVarBndr flag] -> OSet Name -> OSet Name
fv_dtvbs [DTyVarBndr flag]
tvbs OSet Name
fvs = (DTyVarBndr flag -> OSet Name -> OSet Name)
-> OSet Name -> [DTyVarBndr flag] -> OSet Name
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr DTyVarBndr flag -> OSet Name -> OSet Name
forall flag. DTyVarBndr flag -> OSet Name -> OSet Name
fv_dtvb OSet Name
fvs [DTyVarBndr flag]
tvbs

-- | Adjust the free variables of something following a 'DTyVarBndr'.
fv_dtvb :: DTyVarBndr flag -> OSet Name -> OSet Name
fv_dtvb :: DTyVarBndr flag -> OSet Name -> OSet Name
fv_dtvb (DPlainTV Name
n flag
_)    OSet Name
fvs = Name -> OSet Name -> OSet Name
forall a. Ord a => a -> OSet a -> OSet a
OS.delete Name
n OSet Name
fvs
fv_dtvb (DKindedTV Name
n flag
_ DType
k) OSet Name
fvs = Name -> OSet Name -> OSet Name
forall a. Ord a => a -> OSet a -> OSet a
OS.delete Name
n OSet Name
fvs OSet Name -> OSet Name -> OSet Name
forall a. Semigroup a => a -> a -> a
<> DType -> OSet Name
fvDType DType
k