module Data.TrieSet (
TSet,
(\\),
null,
size,
member,
notMember,
isSubsetOf,
isProperSubsetOf,
empty,
singleton,
insert,
delete,
union,
symmetricDifference,
intersection,
difference,
filter,
partition,
split,
splitMember,
map,
mapMonotonic,
foldl,
foldr,
findMin,
findMax,
deleteMin,
deleteMax,
deleteFindMin,
deleteFindMax,
minView,
maxView,
elemAt,
deleteAt,
lookupIndex,
mapSet,
elems,
toList,
fromList,
toVector,
fromVector,
toAscList,
fromAscList,
fromDistinctAscList,
fromAscVector,
fromDistinctAscVector)
where
import Control.Monad
import Control.Monad.Ends
import Control.Monad.Lookup
import Data.TrieMap.Class
import Data.TrieMap.Class.Instances ()
import Data.TrieMap.TrieKey hiding (foldr, foldl, toList, union, diff, isect)
import qualified Data.TrieMap.TrieKey.SetOp as Set
import Data.TrieMap.Representation.Class
import Data.Vector.Build
import qualified Data.Vector.Generic as G
import Data.Vector.Fusion.Util (unId)
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import qualified Data.Vector.Fusion.Stream.Monadic as S
import Data.Maybe(fromJust)
import qualified Data.Foldable as F
import Data.Monoid (Monoid (..))
import GHC.Exts
import Prelude hiding (foldr, foldl, map, filter, null)
instance TKey a => Eq (TSet a) where
s1 == s2 = size s1 == size s2 && s1 `isSubsetOf` s2
instance (TKey a, Ord a) => Ord (TSet a) where
s1 `compare` s2 = elems s1 `compare` elems s2
instance (TKey a, Show a) => Show (TSet a) where
show s = "fromList " ++ show (elems s)
instance TKey a => Monoid (TSet a) where
mempty = empty
mappend = union
empty :: TKey a => TSet a
empty = TSet emptyM
insert :: TKey a => a -> TSet a -> TSet a
insert a (TSet s) = TSet (insertWithM (const (Elem a)) (toRep a) (Elem a) s)
delete :: TKey a => a -> TSet a -> TSet a
delete a (TSet s) = TSet (searchMC (toRep a) s clearM (const clearM))
singleton :: TKey a => a -> TSet a
singleton a = TSet (singletonM (toRep a) (Elem a))
union :: TKey a => TSet a -> TSet a -> TSet a
TSet s1 `union` TSet s2 = TSet (Set.union (const . Just) s1 s2)
symmetricDifference :: TKey a => TSet a -> TSet a -> TSet a
TSet s1 `symmetricDifference` TSet s2 = TSet (Set.union (\ _ _ -> Nothing) s1 s2)
difference :: TKey a => TSet a -> TSet a -> TSet a
TSet s1 `difference` TSet s2 = TSet (Set.diff (\ _ _ -> Nothing) s1 s2)
intersection :: TKey a => TSet a -> TSet a -> TSet a
TSet s1 `intersection` TSet s2 = TSet (Set.isect (const . Just) s1 s2)
filter :: TKey a => (a -> Bool) -> TSet a -> TSet a
filter p (TSet s) = TSet (mapMaybe (\ (Elem a) -> if p a then return (Elem a) else mzero) s)
partition :: TKey a => (a -> Bool) -> TSet a -> (TSet a, TSet a)
partition p (TSet s) = case mapEither f s of
(# s1, s2 #) -> (TSet s1, TSet s2)
where f e@(Elem a)
| p a = (# Just e, Nothing #)
| otherwise = (# Nothing, Just e #)
split :: TKey a => a -> TSet a -> (TSet a, TSet a)
split a s = case splitMember a s of
(sL, _, sR) -> (sL, sR)
splitMember :: TKey a => a -> TSet a -> (TSet a, Bool, TSet a)
splitMember a (TSet s) = searchMC (toRep a) s nomatch match where
nomatch hole = (TSet (beforeM hole), False, TSet (afterM hole))
match _ hole = (TSet (beforeM hole), True, TSet (afterM hole))
map :: (TKey a, TKey b) => (a -> b) -> TSet a -> TSet b
map f s = fromList [f x | x <- elems s]
mapMonotonic :: (TKey a, TKey b) => (a -> b) -> TSet a -> TSet b
mapMonotonic f s = fromDistinctAscList [f x | x <- toAscList s]
foldr :: TKey a => (a -> b -> b) -> b -> TSet a -> b
foldr f z (TSet s) = F.foldr (flip $ F.foldr f) z s
foldl :: TKey b => (a -> b -> a) -> a -> TSet b -> a
foldl f z (TSet s) = F.foldl (F.foldl f) z s
findMin :: TKey a => TSet a -> a
findMin = fst . deleteFindMin
findMax :: TKey a => TSet a -> a
findMax = fst . deleteFindMax
deleteMin :: TKey a => TSet a -> TSet a
deleteMin s = maybe s snd (minView s)
deleteMax :: TKey a => TSet a -> TSet a
deleteMax s = maybe s snd (maxView s)
deleteFindMin :: TKey a => TSet a -> (a, TSet a)
deleteFindMin = fromJust . minView
deleteFindMax :: TKey a => TSet a -> (a, TSet a)
deleteFindMax = fromJust . maxView
minView :: TKey a => TSet a -> Maybe (a, TSet a)
minView (TSet s) = case getFirst (extractHoleM s) of
Nothing -> Nothing
Just (Elem a, hole) -> Just (a, TSet (afterM hole))
maxView :: TKey a => TSet a -> Maybe (a, TSet a)
maxView (TSet s) = case getLast (extractHoleM s) of
Nothing -> Nothing
Just (Elem a, hole) -> Just (a, TSet (beforeM hole))
elems :: TKey a => TSet a -> [a]
elems = toAscList
toList :: TKey a => TSet a -> [a]
toList = toAscList
toAscList :: TKey a => TSet a -> [a]
toAscList s = build (\ c n -> foldr c n s)
fromFoldStream :: (Monad m, Repr a, TrieKey (Rep a)) => FromList z (Rep a) (Elem a) -> Stream m a -> m (TSet a)
fromFoldStream Foldl{..} (Stream suc s0 _) = run s0 where
run s = do
step <- suc s
case step of
Done -> return empty
Skip s' -> run s'
Yield x s' -> run' (begin (toRep x) (Elem x)) s'
run' stack s = do
step <- suc s
case step of
Done -> return (TSet (done stack))
Skip s' -> run' stack s'
Yield x s' -> run' (snoc stack (toRep x) (Elem x)) s'
fromList :: TKey a => [a] -> TSet a
fromList xs = unId (fromFoldStream (uFold const) (S.fromList xs))
fromVector :: (TKey a, G.Vector v a) => v a -> TSet a
fromVector xs = unId (fromFoldStream (uFold const) (G.stream xs))
fromAscList :: TKey a => [a] -> TSet a
fromAscList xs = unId (fromFoldStream (aFold const) (S.fromList xs))
fromAscVector :: (TKey a, G.Vector v a) => v a -> TSet a
fromAscVector xs = unId (fromFoldStream (aFold const) (G.stream xs))
fromDistinctAscList :: TKey a => [a] -> TSet a
fromDistinctAscList xs = unId (fromFoldStream daFold (S.fromList xs))
fromDistinctAscVector :: (TKey a, G.Vector v a) => v a -> TSet a
fromDistinctAscVector xs = unId (fromFoldStream daFold (G.stream xs))
toVector :: (TKey a, G.Vector v a) => TSet a -> v a
toVector (TSet s) = toVectorMapN (sizeM s) getElem s
null :: TKey a => TSet a -> Bool
null (TSet s) = isNull s
size :: TKey a => TSet a -> Int
size (TSet s) = getSize s
member :: TKey a => a -> TSet a -> Bool
member a (TSet s) = runLookup (lookupMC (toRep a) s) False (const True)
notMember :: TKey a => a -> TSet a -> Bool
notMember = not .: member
isSubsetOf :: TKey a => TSet a -> TSet a -> Bool
TSet s1 `isSubsetOf` TSet s2 = let ?le = \ _ _ -> True in s1 <=? s2
isProperSubsetOf :: TKey a => TSet a -> TSet a -> Bool
s1 `isProperSubsetOf` s2 = size s1 < size s2 && s1 `isSubsetOf` s2
(\\) :: TKey a => TSet a -> TSet a -> TSet a
(\\) = difference
mapSet :: TKey a => (a -> b) -> TSet a -> TMap a b
mapSet f (TSet s) = TMap (fmap (\ (Elem a) -> Assoc a (f a)) s)
elemAt :: TKey a => Int -> TSet a -> a
elemAt i (TSet s) = case indexM s (unbox i) of
(# _, Elem a, _ #) -> a
deleteAt :: TKey a => Int -> TSet a -> TSet a
deleteAt i (TSet s) = case indexM s (unbox i) of
(# _, _, hole #) -> TSet (clearM hole)
lookupIndex :: TKey a => a -> TSet a -> Maybe Int
lookupIndex a (TSet s) = searchMC (toRep a) s (\ _ -> Nothing) (\ _ hole -> Just $ sizeM (beforeM hole))