-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
module Retrie.FreeVars
  ( FreeVars
  , elemFVs
  , freeVars
  , substFVs
  ) where

import Data.Generics hiding (Fixity)

import Retrie.ExactPrint.Annotated
import Retrie.GHC
import Retrie.Quantifiers
import Retrie.Substitution

--------------------------------------------------------------------------------

newtype FreeVars = FreeVars (UniqSet FastString)

emptyFVs :: FreeVars
emptyFVs :: FreeVars
emptyFVs = UniqSet FastString -> FreeVars
FreeVars forall a. UniqSet a
emptyUniqSet

instance Semigroup FreeVars where
  <> :: FreeVars -> FreeVars -> FreeVars
(<>) = forall a. Monoid a => a -> a -> a
mappend

instance Monoid FreeVars where
  mempty :: FreeVars
mempty = FreeVars
emptyFVs
  mappend :: FreeVars -> FreeVars -> FreeVars
mappend (FreeVars UniqSet FastString
s1) (FreeVars UniqSet FastString
s2) = UniqSet FastString -> FreeVars
FreeVars forall a b. (a -> b) -> a -> b
$ UniqSet FastString
s1 forall a. Semigroup a => a -> a -> a
<> UniqSet FastString
s2

instance Show FreeVars where
  show :: FreeVars -> String
show (FreeVars UniqSet FastString
m) = forall a. Show a => a -> String
show (forall elt. UniqSet elt -> [elt]
nonDetEltsUniqSet UniqSet FastString
m)

substFVs :: Substitution -> FreeVars
substFVs :: Substitution -> FreeVars
substFVs = forall a.
((FastString, HoleVal) -> a -> a) -> a -> Substitution -> a
foldSubst (HoleVal -> FreeVars -> FreeVars
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) FreeVars
emptyFVs
  where
    f :: HoleVal -> FreeVars -> FreeVars
f (HoleExpr AnnotatedHsExpr
e) FreeVars
fvs = forall a. (Data a, Typeable a) => Quantifiers -> a -> FreeVars
freeVars Quantifiers
emptyQs (forall ast. Annotated ast -> ast
astA AnnotatedHsExpr
e) forall a. Semigroup a => a -> a -> a
<> FreeVars
fvs
    f (HoleRdr RdrName
rdr) FreeVars
fvs = RdrName -> FreeVars
rdrFV RdrName
rdr forall a. Semigroup a => a -> a -> a
<> FreeVars
fvs
    f HoleVal
_ FreeVars
fvs = FreeVars
fvs -- TODO(anfarmer) types?

-- | This is an over-approximation, but that is fine for our purposes.
freeVars :: (Data a, Typeable a) => Quantifiers -> a -> FreeVars
freeVars :: forall a. (Data a, Typeable a) => Quantifiers -> a -> FreeVars
freeVars Quantifiers
qs = forall r. (r -> r -> r) -> GenericQ r -> GenericQ r
everything forall a. Semigroup a => a -> a -> a
(<>) (forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ FreeVars
emptyFVs HsExpr GhcPs -> FreeVars
fvsExpr forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`extQ` HsType GhcPs -> FreeVars
fvsType)
  where
    fvsExpr :: HsExpr GhcPs -> FreeVars
    fvsExpr :: HsExpr GhcPs -> FreeVars
fvsExpr HsExpr GhcPs
e
      | Just (L SrcSpanAnnN
_ RdrName
rdr) <- forall p. HsExpr p -> Maybe (LIdP p)
varRdrName HsExpr GhcPs
e
      , Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ RdrName -> Quantifiers -> Bool
isQ RdrName
rdr Quantifiers
qs = RdrName -> FreeVars
rdrFV RdrName
rdr
    fvsExpr HsExpr GhcPs
_ = FreeVars
emptyFVs

    fvsType :: HsType GhcPs -> FreeVars
    fvsType :: HsType GhcPs -> FreeVars
fvsType HsType GhcPs
ty
      | Just (L SrcSpanAnnN
_ RdrName
rdr) <- forall p. HsType p -> Maybe (LIdP p)
tyvarRdrName HsType GhcPs
ty
      , Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ RdrName -> Quantifiers -> Bool
isQ RdrName
rdr Quantifiers
qs = RdrName -> FreeVars
rdrFV RdrName
rdr
    fvsType HsType GhcPs
_ = FreeVars
emptyFVs

elemFVs :: RdrName -> FreeVars -> Bool
elemFVs :: RdrName -> FreeVars -> Bool
elemFVs RdrName
rdr (FreeVars UniqSet FastString
m) = forall a. Uniquable a => a -> UniqSet a -> Bool
elementOfUniqSet (RdrName -> FastString
rdrFS RdrName
rdr) UniqSet FastString
m

rdrFV :: RdrName -> FreeVars
rdrFV :: RdrName -> FreeVars
rdrFV = UniqSet FastString -> FreeVars
FreeVars forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Uniquable a => a -> UniqSet a
unitUniqSet forall b c a. (b -> c) -> (a -> b) -> a -> c
. RdrName -> FastString
rdrFS