-- 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 UniqSet FastString
forall a. UniqSet a
emptyUniqSet

instance Semigroup FreeVars where
  <> :: FreeVars -> FreeVars -> FreeVars
(<>) = 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 (UniqSet FastString -> FreeVars) -> UniqSet FastString -> FreeVars
forall a b. (a -> b) -> a -> b
$ UniqSet FastString
s1 UniqSet FastString -> UniqSet FastString -> UniqSet FastString
forall a. Semigroup a => a -> a -> a
<> UniqSet FastString
s2

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

substFVs :: Substitution -> FreeVars
substFVs :: Substitution -> FreeVars
substFVs = ((FastString, HoleVal) -> FreeVars -> FreeVars)
-> FreeVars -> Substitution -> FreeVars
forall a.
((FastString, HoleVal) -> a -> a) -> a -> Substitution -> a
foldSubst (HoleVal -> FreeVars -> FreeVars
f (HoleVal -> FreeVars -> FreeVars)
-> ((FastString, HoleVal) -> HoleVal)
-> (FastString, HoleVal)
-> FreeVars
-> FreeVars
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FastString, HoleVal) -> HoleVal
forall a b. (a, b) -> b
snd) FreeVars
emptyFVs
  where
    f :: HoleVal -> FreeVars -> FreeVars
f (HoleExpr AnnotatedHsExpr
e) FreeVars
fvs = Quantifiers -> GenLocated SrcSpanAnnA (HsExpr GhcPs) -> FreeVars
forall a. (Data a, Typeable a) => Quantifiers -> a -> FreeVars
freeVars Quantifiers
emptyQs (Annotated (GenLocated SrcSpanAnnA (HsExpr GhcPs))
-> GenLocated SrcSpanAnnA (HsExpr GhcPs)
forall ast. Annotated ast -> ast
astA AnnotatedHsExpr
Annotated (GenLocated SrcSpanAnnA (HsExpr GhcPs))
e) FreeVars -> FreeVars -> FreeVars
forall a. Semigroup a => a -> a -> a
<> FreeVars
fvs
    f (HoleRdr RdrName
rdr) FreeVars
fvs = RdrName -> FreeVars
rdrFV RdrName
rdr FreeVars -> FreeVars -> FreeVars
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 = (FreeVars -> FreeVars -> FreeVars)
-> GenericQ FreeVars -> GenericQ FreeVars
forall r. (r -> r -> r) -> GenericQ r -> GenericQ r
everything FreeVars -> FreeVars -> FreeVars
forall a. Semigroup a => a -> a -> a
(<>) (FreeVars -> (HsExpr GhcPs -> FreeVars) -> a -> FreeVars
forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ FreeVars
emptyFVs HsExpr GhcPs -> FreeVars
fvsExpr (a -> FreeVars) -> (HsType GhcPs -> FreeVars) -> a -> FreeVars
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` HsType GhcPs -> FreeVars
fvsType)
  where
    fvsExpr :: HsExpr GhcPs -> FreeVars
    fvsExpr :: HsExpr GhcPs -> FreeVars
fvsExpr HsExpr GhcPs
e
      | Just (L SrcSpanAnnN
_ RdrName
rdr) <- HsExpr GhcPs -> Maybe (LIdP GhcPs)
forall p. HsExpr p -> Maybe (LIdP p)
varRdrName HsExpr GhcPs
e
      , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
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) <- HsType GhcPs -> Maybe (LIdP GhcPs)
forall p. HsType p -> Maybe (LIdP p)
tyvarRdrName HsType GhcPs
ty
      , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
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) = FastString -> UniqSet FastString -> Bool
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 (UniqSet FastString -> FreeVars)
-> (RdrName -> UniqSet FastString) -> RdrName -> FreeVars
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FastString -> UniqSet FastString
forall a. Uniquable a => a -> UniqSet a
unitUniqSet (FastString -> UniqSet FastString)
-> (RdrName -> FastString) -> RdrName -> UniqSet FastString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RdrName -> FastString
rdrFS