1-- Copyright (c) Facebook, Inc. and its affiliates.
2--
3-- This source code is licensed under the MIT license found in the
4-- LICENSE file in the root directory of this source tree.
5--
6module Retrie.FreeVars
7  ( FreeVars
8  , elemFVs
9  , freeVars
10  , substFVs
11  ) where
12
13import Data.Generics hiding (Fixity)
14
15import Retrie.ExactPrint.Annotated
16import Retrie.GHC
17import Retrie.Quantifiers
18import Retrie.Substitution
19
20--------------------------------------------------------------------------------
21
22newtype FreeVars = FreeVars (UniqSet FastString)
23
24emptyFVs :: FreeVars
25emptyFVs = FreeVars emptyUniqSet
26
27instance Semigroup FreeVars where
28  (<>) = mappend
29
30instance Monoid FreeVars where
31  mempty = emptyFVs
32  mappend (FreeVars s1) (FreeVars s2) = FreeVars $ s1 <> s2
33
34instance Show FreeVars where
35  show (FreeVars m) = show (nonDetEltsUniqSet m)
36
37substFVs :: Substitution -> FreeVars
38substFVs = foldSubst (f . snd) emptyFVs
39  where
40    f (HoleExpr e) fvs = freeVars emptyQs (astA e) <> fvs
41    f (HoleRdr rdr) fvs = rdrFV rdr <> fvs
42    f _ fvs = fvs -- TODO(anfarmer) types?
43
44-- | This is an over-approximation, but that is fine for our purposes.
45freeVars :: (Data a, Typeable a) => Quantifiers -> a -> FreeVars
46freeVars qs = everything (<>) (mkQ emptyFVs fvsExpr `extQ` fvsType)
47  where
48    fvsExpr :: HsExpr GhcPs -> FreeVars
49    fvsExpr e
50      | Just (L _ rdr) <- varRdrName e
51      , not $ isQ rdr qs = rdrFV rdr
52    fvsExpr _ = emptyFVs
53
54    fvsType :: HsType GhcPs -> FreeVars
55    fvsType ty
56      | Just (L _ rdr) <- tyvarRdrName ty
57      , not $ isQ rdr qs = rdrFV rdr
58    fvsType _ = emptyFVs
59
60elemFVs :: RdrName -> FreeVars -> Bool
61elemFVs rdr (FreeVars m) = elementOfUniqSet (rdrFS rdr) m
62
63rdrFV :: RdrName -> FreeVars
64rdrFV = FreeVars . unitUniqSet . rdrFS
65