{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PatternGuards #-}

-- |
-- Module      : Data.Binding.Hobbits.NameSet
-- Copyright   : (c) 2020 Edwin Westbrook
--
-- License     : BSD3
--
-- Maintainer  : westbrook@galois.com
-- Stability   : experimental
-- Portability : GHC
--
-- Implements sets of 'Name's using 'Data.IntSet.Strict'. Note that these
-- mappings are strict.

module Data.Binding.Hobbits.NameSet (
  NameSet(), SomeName(..)
  , empty, singleton, fromList, toList
  , insert, delete, member, null, size
  , union, unions, difference, (\\), intersection
  , map, foldr, foldl
  , liftNameSet
  ) where

import Prelude hiding (lookup, null, map, foldr, foldl)
import qualified Prelude as Prelude (map)
import Data.Maybe
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.Kind
import qualified Data.Foldable as Foldable

import Data.Binding.Hobbits.Internal.Name
import Data.Binding.Hobbits.Mb
import Data.Binding.Hobbits.NuMatching
import Data.Binding.Hobbits.QQ
import Data.Binding.Hobbits.Liftable

-- | A set of 'Name's whose types all have kind @k@
newtype NameSet k = NameSet { NameSet k -> IntSet
unNameSet :: IntSet }

-- | A 'Name' of some unknown type of kind @k@
data SomeName k = forall (a :: k). SomeName (Name a)

$(mkNuMatching [t| forall k. SomeName k |])

-- | The empty 'NameSet'
empty :: NameSet k
empty :: NameSet k
empty = IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ IntSet
IntSet.empty

-- | The singleton 'NameSet'
singleton :: Name (a::k) -> NameSet k
singleton :: Name a -> NameSet k
singleton (MkName Int
i) = IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ Int -> IntSet
IntSet.singleton (Int -> IntSet) -> Int -> IntSet
forall a b. (a -> b) -> a -> b
$ Int
i

-- | Convert a list of names to a 'NameSet'
fromList :: [SomeName k] -> NameSet k
fromList :: [SomeName k] -> NameSet k
fromList =
  IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k)
-> ([SomeName k] -> IntSet) -> [SomeName k] -> NameSet k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> IntSet
IntSet.fromList ([Int] -> IntSet)
-> ([SomeName k] -> [Int]) -> [SomeName k] -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SomeName k -> Int) -> [SomeName k] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map (\(SomeName (MkName Int
i)) -> Int
i)

-- | Convert a 'NameSet' to a list
toList :: NameSet k -> [SomeName k]
toList :: NameSet k -> [SomeName k]
toList (NameSet IntSet
s) = (Int -> SomeName k) -> [Int] -> [SomeName k]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map (Name Any -> SomeName k
forall k (a :: k). Name a -> SomeName k
SomeName (Name Any -> SomeName k) -> (Int -> Name Any) -> Int -> SomeName k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Name Any
forall k (a :: k). Int -> Name a
MkName) (IntSet -> [Int]
IntSet.toList IntSet
s)

-- | Insert a name into a 'NameSet'
insert :: Name (a::k) -> NameSet k -> NameSet k
insert :: Name a -> NameSet k -> NameSet k
insert (MkName Int
i) (NameSet IntSet
s) = IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ Int -> IntSet -> IntSet
IntSet.insert Int
i IntSet
s

-- | Delete a name from a 'NameSet'
delete :: Name (a::k) -> NameSet k -> NameSet k
delete :: Name a -> NameSet k -> NameSet k
delete (MkName Int
i) (NameSet IntSet
s) = IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ Int -> IntSet -> IntSet
IntSet.delete Int
i IntSet
s

-- | Test if a name is in a 'NameSet'
member :: Name (a::k) -> NameSet k -> Bool
member :: Name a -> NameSet k -> Bool
member (MkName Int
i) (NameSet IntSet
s) = Int -> IntSet -> Bool
IntSet.member Int
i IntSet
s

-- | Test if a 'NameSet' is empty
null :: NameSet k -> Bool
null :: NameSet k -> Bool
null (NameSet IntSet
s) = IntSet -> Bool
IntSet.null IntSet
s

-- | Compute the cardinality of a 'NameSet'
size :: NameSet k -> Int
size :: NameSet k -> Int
size (NameSet IntSet
s) = IntSet -> Int
IntSet.size IntSet
s

-- | Take the union of two 'NameSet's
union :: NameSet k -> NameSet k -> NameSet k
union :: NameSet k -> NameSet k -> NameSet k
union (NameSet IntSet
s1) (NameSet IntSet
s2) = IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ IntSet -> IntSet -> IntSet
IntSet.union IntSet
s1 IntSet
s2

-- | Take the union of a list of 'NameSet's
unions :: Foldable f => f (NameSet k) -> NameSet k
unions :: f (NameSet k) -> NameSet k
unions = (NameSet k -> NameSet k -> NameSet k)
-> NameSet k -> f (NameSet k) -> NameSet k
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' NameSet k -> NameSet k -> NameSet k
forall k (k :: k). NameSet k -> NameSet k -> NameSet k
union NameSet k
forall k (k :: k). NameSet k
empty

-- | Take the set of all elements of the first 'NameSet' not in the second
difference :: NameSet k -> NameSet k -> NameSet k
difference :: NameSet k -> NameSet k -> NameSet k
difference (NameSet IntSet
s1) (NameSet IntSet
s2) = IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ IntSet -> IntSet -> IntSet
IntSet.difference IntSet
s1 IntSet
s2

-- | Another name for 'difference'
(\\) :: NameSet k -> NameSet k -> NameSet k
\\ :: NameSet k -> NameSet k -> NameSet k
(\\) = NameSet k -> NameSet k -> NameSet k
forall k (k :: k). NameSet k -> NameSet k -> NameSet k
difference

-- | Take the intersection of two 'NameSet's
intersection :: NameSet k -> NameSet k -> NameSet k
intersection :: NameSet k -> NameSet k -> NameSet k
intersection (NameSet IntSet
s1) (NameSet IntSet
s2) = IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ IntSet -> IntSet -> IntSet
IntSet.intersection IntSet
s1 IntSet
s2

-- | Map a function across all elements of a 'NameSet'
map :: (forall (a::k). Name a -> Name a) -> NameSet k -> NameSet k
map :: (forall (a :: k). Name a -> Name a) -> NameSet k -> NameSet k
map forall (a :: k). Name a -> Name a
f (NameSet IntSet
s) =
  IntSet -> NameSet k
forall k (k :: k). IntSet -> NameSet k
NameSet (IntSet -> NameSet k) -> IntSet -> NameSet k
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> IntSet -> IntSet
IntSet.map (\Int
i -> let (MkName Int
j) = Name Any -> Name Any
forall (a :: k). Name a -> Name a
f (Int -> Name Any
forall k (a :: k). Int -> Name a
MkName Int
i) in Int
j) IntSet
s

-- | Perform a right fold of a function across all elements of a 'NameSet'
foldr :: (forall (a::k). Name a -> r -> r) -> r -> NameSet k -> r
foldr :: (forall (a :: k). Name a -> r -> r) -> r -> NameSet k -> r
foldr forall (a :: k). Name a -> r -> r
f r
r (NameSet IntSet
s) = (Int -> r -> r) -> r -> IntSet -> r
forall b. (Int -> b -> b) -> b -> IntSet -> b
IntSet.foldr (Name Any -> r -> r
forall (a :: k). Name a -> r -> r
f (Name Any -> r -> r) -> (Int -> Name Any) -> Int -> r -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Name Any
forall k (a :: k). Int -> Name a
MkName) r
r IntSet
s

-- | Perform a left fold of a function across all elements of a 'NameSet'
foldl :: (forall (a::k). r -> Name a -> r) -> r -> NameSet k -> r
foldl :: (forall (a :: k). r -> Name a -> r) -> r -> NameSet k -> r
foldl forall (a :: k). r -> Name a -> r
f r
r (NameSet IntSet
s) = (r -> Int -> r) -> r -> IntSet -> r
forall a. (a -> Int -> a) -> a -> IntSet -> a
IntSet.foldl (\r
r -> r -> Name Any -> r
forall (a :: k). r -> Name a -> r
f r
r (Name Any -> r) -> (Int -> Name Any) -> Int -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Name Any
forall k (a :: k). Int -> Name a
MkName) r
r IntSet
s

-- | Lift a 'NameSet' out of a name-binding by lifting all names not bound by
-- the name-binding and then forming a 'NameSet' from those lifted names
liftNameSet :: Mb ctx (NameSet (k :: Type)) -> NameSet k
liftNameSet :: Mb ctx (NameSet k) -> NameSet k
liftNameSet Mb ctx (NameSet k)
mb_s = [SomeName k] -> NameSet k
forall k. [SomeName k] -> NameSet k
fromList ([SomeName k] -> NameSet k) -> [SomeName k] -> NameSet k
forall a b. (a -> b) -> a -> b
$ (Mb ctx (SomeName k) -> Maybe (SomeName k))
-> [Mb ctx (SomeName k)] -> [SomeName k]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Mb ctx (SomeName k) -> Maybe (SomeName k)
forall k (ctx' :: RList k) k'.
Mb ctx' (SomeName k') -> Maybe (SomeName k')
helper ([Mb ctx (SomeName k)] -> [SomeName k])
-> [Mb ctx (SomeName k)] -> [SomeName k]
forall a b. (a -> b) -> a -> b
$ Mb ctx [SomeName k] -> [Mb ctx (SomeName k)]
forall k a (c :: RList k). NuMatching a => Mb c [a] -> [Mb c a]
mbList (Mb ctx [SomeName k] -> [Mb ctx (SomeName k)])
-> Mb ctx [SomeName k] -> [Mb ctx (SomeName k)]
forall a b. (a -> b) -> a -> b
$ (NameSet k -> [SomeName k])
-> Mb ctx (NameSet k) -> Mb ctx [SomeName k]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NameSet k -> [SomeName k]
forall k. NameSet k -> [SomeName k]
toList Mb ctx (NameSet k)
mb_s
  where
    helper :: forall ctx' k'. Mb ctx' (SomeName k') -> Maybe (SomeName k')
    helper :: Mb ctx' (SomeName k') -> Maybe (SomeName k')
helper Mb ctx' (SomeName k')
[nuP| SomeName mb_n |]
      | Right Name a
n <- Mb ctx' (Name a) -> Either (Member ctx' a) (Name a)
forall k1 k2 (a :: k1) (ctx :: RList k2).
Mb ctx (Name a) -> Either (Member ctx a) (Name a)
mbNameBoundP Mb ctx' (Name a)
mb_n = SomeName k' -> Maybe (SomeName k')
forall a. a -> Maybe a
Just (Name a -> SomeName k'
forall k (a :: k). Name a -> SomeName k
SomeName Name a
n)
    helper Mb ctx' (SomeName k')
_ = Maybe (SomeName k')
forall a. Maybe a
Nothing