-- 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.
--
{-# LANGUAGE TypeFamilies #-}
module Retrie.Quantifiers
  ( Quantifiers
  , emptyQs
  , exceptQ
  , isQ
  , mkQs
  , mkFSQs
  , qList
  , qSet
  , unionQ
  ) where

import GHC.Exts (IsList(..))
import Retrie.GHC

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

-- Note [Why not RdrNames?]
-- RdrNames carry their namespace in their OccName. The unification of holes
-- already handles the namespace more finely, so the extra distinction inside the
-- RdrName is just a pain when manipulating the substitution, checking for
-- equality, etc.

-- Thus, Quantifiers are keyed on FastString (an OccName is a namespace plus
-- FastString, so this is just the OccName without its namespace).

instance IsList Quantifiers where
  type Item Quantifiers = String

  fromList :: [Item Quantifiers] -> Quantifiers
fromList = [FastString] -> Quantifiers
mkFSQs ([FastString] -> Quantifiers)
-> ([String] -> [FastString]) -> [String] -> Quantifiers
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> FastString) -> [String] -> [FastString]
forall a b. (a -> b) -> [a] -> [b]
map String -> FastString
mkFastString
  toList :: Quantifiers -> [Item Quantifiers]
toList = (FastString -> String) -> [FastString] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map FastString -> String
unpackFS ([FastString] -> [String])
-> (Quantifiers -> [FastString]) -> Quantifiers -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Quantifiers -> [FastString]
qList

-- | 'Quantifiers' is a set of variable names. If you enable the
-- "OverloadedLists" language extension, you can construct using a literal
-- list of strings.
newtype Quantifiers = Quantifiers
  { Quantifiers -> UniqSet FastString
qSet :: UniqSet FastString
    -- ^ Convert to a 'UniqSet'.
  }

-- | Construct from GHC's 'RdrName's.
mkQs :: [RdrName] -> Quantifiers
mkQs :: [RdrName] -> Quantifiers
mkQs = [FastString] -> Quantifiers
mkFSQs ([FastString] -> Quantifiers)
-> ([RdrName] -> [FastString]) -> [RdrName] -> Quantifiers
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RdrName -> FastString) -> [RdrName] -> [FastString]
forall a b. (a -> b) -> [a] -> [b]
map RdrName -> FastString
rdrFS

-- | Construct from 'FastString's.
mkFSQs :: [FastString] -> Quantifiers
mkFSQs :: [FastString] -> Quantifiers
mkFSQs = UniqSet FastString -> Quantifiers
Quantifiers (UniqSet FastString -> Quantifiers)
-> ([FastString] -> UniqSet FastString)
-> [FastString]
-> Quantifiers
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [FastString] -> UniqSet FastString
forall a. Uniquable a => [a] -> UniqSet a
mkUniqSet

-- | The empty set.
emptyQs :: Quantifiers
emptyQs :: Quantifiers
emptyQs = UniqSet FastString -> Quantifiers
Quantifiers UniqSet FastString
forall a. UniqSet a
emptyUniqSet

-- | Existence check.
isQ :: RdrName -> Quantifiers -> Bool
isQ :: RdrName -> Quantifiers -> Bool
isQ RdrName
r (Quantifiers UniqSet FastString
s) = FastString -> UniqSet FastString -> Bool
forall a. Uniquable a => a -> UniqSet a -> Bool
elementOfUniqSet (RdrName -> FastString
rdrFS RdrName
r) UniqSet FastString
s

-- | Set union.
unionQ :: Quantifiers -> Quantifiers -> Quantifiers
unionQ :: Quantifiers -> Quantifiers -> Quantifiers
unionQ (Quantifiers UniqSet FastString
s) (Quantifiers UniqSet FastString
t) = UniqSet FastString -> Quantifiers
Quantifiers (UniqSet FastString -> Quantifiers)
-> UniqSet FastString -> Quantifiers
forall a b. (a -> b) -> a -> b
$ UniqSet FastString -> UniqSet FastString -> UniqSet FastString
forall a. UniqSet a -> UniqSet a -> UniqSet a
unionUniqSets UniqSet FastString
s UniqSet FastString
t

-- | Remove a set of 'RdrName's from the set.
exceptQ :: Quantifiers -> [RdrName] -> Quantifiers
exceptQ :: Quantifiers -> [RdrName] -> Quantifiers
exceptQ (Quantifiers UniqSet FastString
s) [RdrName]
rs =
  UniqSet FastString -> Quantifiers
Quantifiers (UniqSet FastString -> Quantifiers)
-> UniqSet FastString -> Quantifiers
forall a b. (a -> b) -> a -> b
$ UniqSet FastString -> [FastString] -> UniqSet FastString
forall a. Uniquable a => UniqSet a -> [a] -> UniqSet a
delListFromUniqSet UniqSet FastString
s ((RdrName -> FastString) -> [RdrName] -> [FastString]
forall a b. (a -> b) -> [a] -> [b]
map RdrName -> FastString
rdrFS [RdrName]
rs)

-- | Convert to a list.
qList :: Quantifiers -> [FastString]
qList :: Quantifiers -> [FastString]
qList (Quantifiers UniqSet FastString
s) = UniqSet FastString -> [FastString]
forall elt. UniqSet elt -> [elt]
nonDetEltsUniqSet UniqSet FastString
s