{-# LANGUAGE CPP #-}

-- | See 'Filtrable'.
module Data.Filtrable
  ( Filtrable (..)
  , (<$?>), (<*?>)
  , nub, nubBy, nubOrd, nubOrdBy
  ) where

import Prelude hiding (filter)

import Control.Applicative
import Control.Applicative.Backwards
import Control.Monad
import qualified Control.Monad.Trans.State as M
import Data.Bool (bool)
import Data.Functor.Compose
import Data.Functor.Product
import Data.Functor.Reverse
import Data.Functor.Sum
import Data.Proxy
import Data.Traversable

#ifdef MIN_VERSION_containers
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
#endif

import qualified Data.Set.Private as Set

-- | Class of filtrable containers, i.e. containers we can map over while selectively dropping elements.
--
-- Laws:
--
-- * @'mapMaybe' 'Just' = id@
--
-- * @'mapMaybe' f = 'catMaybes' ∘ 'fmap' f@
--
-- * @'catMaybes' = 'mapMaybe' id@
--
-- * @'filter' f = 'mapMaybe' (\\ x -> 'bool' 'Nothing' ('Just' x) (f x))@
--
-- * @'mapMaybe' g . 'mapMaybe' f = 'mapMaybe' (g '<=<' f)@
--
--   Laws if @'Foldable' f@:
--
-- * @'foldMap' g . 'filter' f = 'foldMap' (\\ x -> 'bool' 'mempty' (g x) (f x))@
class Functor f => Filtrable f where
    {-# MINIMAL mapMaybe | catMaybes #-}

    -- | Map the container with the given function, dropping the elements for which it returns 'Nothing'.
    mapMaybe :: (a -> Maybe b) -> f a -> f b
    mapMaybe a -> Maybe b
f = forall (f :: * -> *) a. Filtrable f => f (Maybe a) -> f a
catMaybes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Maybe b
f

    -- | @'catMaybes' = 'mapMaybe' 'id'@
    catMaybes :: f (Maybe a) -> f a
    catMaybes = forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe forall a. a -> a
id

    -- | Drop the elements for which the given predicate is 'False'.
    filter :: (a -> Bool) -> f a -> f a
    filter a -> Bool
f = forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (forall (f :: * -> *) a b. Functor f => a -> f b -> f a
(<$) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

    -- | Traverse the container with the given function, dropping the elements for which it returns 'Nothing'.
    mapMaybeA :: (Traversable f, Applicative p) => (a -> p (Maybe b)) -> f a -> p (f b)
    mapMaybeA a -> p (Maybe b)
f f a
xs = forall (f :: * -> *) a. Filtrable f => f (Maybe a) -> f a
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> p (Maybe b)
f f a
xs

    -- | Drop the elements for which the given predicate is 'False'.
    filterA :: (Traversable f, Applicative p) => (a -> p Bool) -> f a -> p (f a)
    filterA a -> p Bool
f = forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA (\ a
x -> (a
x forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> p Bool
f a
x)

    -- | Map the container with the given function, collecting the 'Left's and the 'Right's separately.
    mapEither :: (a -> Either b c) -> f a -> (f b, f c)
    mapEither a -> Either b c
f = (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. a -> Maybe a
Just (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing) forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either b c
f)
                      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing) forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either b c
f)

    -- | Traverse the container with the given function, collecting the 'Left's and the 'Right's separately.
    mapEitherA :: (Traversable f, Applicative p) => (a -> p (Either b c)) -> f a -> p (f b, f c)
    mapEitherA a -> p (Either b c)
f = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. a -> Maybe a
Just forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
`either` forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing) forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> p (Either b c)
f)
                              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
`either` forall a. a -> Maybe a
Just) forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> p (Either b c)
f)

    -- | @'partitionEithers' = 'mapEither' 'id'@
    partitionEithers :: f (Either a b) -> (f a, f b)
    partitionEithers = forall (f :: * -> *) a b c.
Filtrable f =>
(a -> Either b c) -> f a -> (f b, f c)
mapEither forall a. a -> a
id

instance Filtrable [] where
    mapMaybe :: forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe a -> Maybe b
f = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id (:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe b
f) []

    mapMaybeA :: forall (p :: * -> *) a b.
(Traversable [], Applicative p) =>
(a -> p (Maybe b)) -> [a] -> p [b]
mapMaybeA a -> p (Maybe b)
_ [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    mapMaybeA a -> p (Maybe b)
f (a
x:[a]
xs) = forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id (:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> p (Maybe b)
f a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA a -> p (Maybe b)
f [a]
xs

instance Filtrable Maybe where
    mapMaybe :: forall a b. (a -> Maybe b) -> Maybe a -> Maybe b
mapMaybe = forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
(=<<)
    catMaybes :: forall a. Maybe (Maybe a) -> Maybe a
catMaybes = forall (m :: * -> *) a. Monad m => m (m a) -> m a
join

instance Filtrable Proxy where
    mapMaybe :: forall a b. (a -> Maybe b) -> Proxy a -> Proxy b
mapMaybe a -> Maybe b
_ Proxy a
Proxy = forall {k} (t :: k). Proxy t
Proxy

instance Filtrable (Const a) where
    mapMaybe :: forall a b. (a -> Maybe b) -> Const a a -> Const a b
mapMaybe a -> Maybe b
_ (Const a
x) = forall {k} a (b :: k). a -> Const a b
Const a
x

instance (Filtrable f, Filtrable g) => Filtrable (Product f g) where
    mapMaybe :: forall a b. (a -> Maybe b) -> Product f g a -> Product f g b
mapMaybe a -> Maybe b
f (Pair f a
as g a
bs) = forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
Pair (forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f f a
as) (forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f g a
bs)

instance (Filtrable f, Filtrable g) => Filtrable (Sum f g) where
    mapMaybe :: forall a b. (a -> Maybe b) -> Sum f g a -> Sum f g b
mapMaybe a -> Maybe b
f = \ case
        InL f a
as -> forall {k} (f :: k -> *) (g :: k -> *) (a :: k). f a -> Sum f g a
InL (forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f f a
as)
        InR g a
bs -> forall {k} (f :: k -> *) (g :: k -> *) (a :: k). g a -> Sum f g a
InR (forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f g a
bs)

instance (Functor f, Filtrable g) => Filtrable (Compose f g) where
    mapMaybe :: forall a b. (a -> Maybe b) -> Compose f g a -> Compose f g b
mapMaybe a -> Maybe b
f = forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe) a -> Maybe b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose

instance Filtrable f => Filtrable (Backwards f) where
    mapMaybe :: forall a b. (a -> Maybe b) -> Backwards f a -> Backwards f b
mapMaybe a -> Maybe b
f = forall {k} (f :: k -> *) (a :: k). f a -> Backwards f a
Backwards forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (f :: k -> *) (a :: k). Backwards f a -> f a
forwards

instance Filtrable f => Filtrable (Reverse f) where
    mapMaybe :: forall a b. (a -> Maybe b) -> Reverse f a -> Reverse f b
mapMaybe a -> Maybe b
f = forall {k} (f :: k -> *) (a :: k). f a -> Reverse f a
Reverse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (f :: k -> *) (a :: k). Reverse f a -> f a
getReverse

infixl 4 <$?>, <*?>

-- | Infix synonym of 'mapMaybe'
(<$?>) :: Filtrable f => (a -> Maybe b) -> f a -> f b
<$?> :: forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
(<$?>) = forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe

-- | @f '<*?>' a = 'catMaybes' (f '<*>' a)@
(<*?>) :: (Applicative p, Filtrable p) => p (a -> Maybe b) -> p a -> p b
p (a -> Maybe b)
f <*?> :: forall (p :: * -> *) a b.
(Applicative p, Filtrable p) =>
p (a -> Maybe b) -> p a -> p b
<*?> p a
a = forall (f :: * -> *) a. Filtrable f => f (Maybe a) -> f a
catMaybes (p (a -> Maybe b)
f forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> p a
a)

-- | \(\mathcal{O}(n^2)\)
-- Delete all but the first copy of each element, special case of 'nubBy'.
nub :: (Filtrable f, Traversable f, Eq a) => f a -> f a
nub :: forall (f :: * -> *) a.
(Filtrable f, Traversable f, Eq a) =>
f a -> f a
nub = forall (f :: * -> *) a.
(Filtrable f, Traversable f) =>
(a -> a -> Bool) -> f a -> f a
nubBy forall a. Eq a => a -> a -> Bool
(==)

-- | \(\mathcal{O}(n^2)\)
-- Delete all but the first copy of each element, with the given relation.
nubBy :: (Filtrable f, Traversable f) => (a -> a -> Bool) -> f a -> f a
nubBy :: forall (f :: * -> *) a.
(Filtrable f, Traversable f) =>
(a -> a -> Bool) -> f a -> f a
nubBy a -> a -> Bool
eq = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> a
M.evalState []) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (p :: * -> *) a.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p Bool) -> f a -> p (f a)
filterA forall a b. (a -> b) -> a -> b
$ \ a
a -> do
    [a]
as <- forall (m :: * -> *) s. Monad m => StateT s m s
M.get
    let b :: Bool
b = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> Bool
eq a
a) [a]
as
    Bool
b forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
M.modify (a
aforall a. a -> [a] -> [a]
:))

-- | \(\mathcal{O}(n\;\mathrm{log}\;n)\)
-- Delete all but the first copy of each element, special case of 'nubOrdBy'.
nubOrd :: (Filtrable f, Traversable f, Ord a) => f a -> f a
nubOrd :: forall (f :: * -> *) a.
(Filtrable f, Traversable f, Ord a) =>
f a -> f a
nubOrd = forall (f :: * -> *) a.
(Filtrable f, Traversable f) =>
(a -> a -> Ordering) -> f a -> f a
nubOrdBy forall a. Ord a => a -> a -> Ordering
compare

-- | \(\mathcal{O}(n\;\mathrm{log}\;n)\)
-- Delete all but the first copy of each element, with the given relation.
nubOrdBy :: (Filtrable f, Traversable f) => (a -> a -> Ordering) -> f a -> f a
nubOrdBy :: forall (f :: * -> *) a.
(Filtrable f, Traversable f) =>
(a -> a -> Ordering) -> f a -> f a
nubOrdBy a -> a -> Ordering
compare = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> a
M.evalState forall a. Set a
Set.empty) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (p :: * -> *) a.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p Bool) -> f a -> p (f a)
filterA forall a b. (a -> b) -> a -> b
$ \ a
a -> forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
M.state forall a b. (a -> b) -> a -> b
$ \ Set a
as ->
    case forall a. (a -> a -> Ordering) -> a -> Set a -> Maybe (Set a)
Set.insertBy' a -> a -> Ordering
compare a
a Set a
as of
        Maybe (Set a)
Nothing -> (Bool
False, Set a
as)
        Just Set a
as' -> (Bool
True, Set a
as')

#ifdef MIN_VERSION_containers
instance Filtrable IntMap where
    mapMaybe :: forall a b. (a -> Maybe b) -> IntMap a -> IntMap b
mapMaybe = forall a b. (a -> Maybe b) -> IntMap a -> IntMap b
IntMap.mapMaybe
    mapEither :: forall a b c. (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
mapEither = forall a b c. (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c)
IntMap.mapEither
    filter :: forall a. (a -> Bool) -> IntMap a -> IntMap a
filter = forall a. (a -> Bool) -> IntMap a -> IntMap a
IntMap.filter

instance Ord k => Filtrable (Map k) where
    mapMaybe :: forall a b. (a -> Maybe b) -> Map k a -> Map k b
mapMaybe = forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe
    mapEither :: forall a b c. (a -> Either b c) -> Map k a -> (Map k b, Map k c)
mapEither = forall a b c k. (a -> Either b c) -> Map k a -> (Map k b, Map k c)
Map.mapEither
    filter :: forall a. (a -> Bool) -> Map k a -> Map k a
filter = forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter

instance Filtrable Seq where
    mapMaybe :: forall a b. (a -> Maybe b) -> Seq a -> Seq b
mapMaybe a -> Maybe b
f = Seq a -> Seq b
go
      where
        go :: Seq a -> Seq b
go = \ case
            Seq a
Seq.Empty -> forall a. Seq a
Seq.Empty
            a
a Seq.:<| Seq a
as -> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id forall a. a -> Seq a -> Seq a
(Seq.:<|) (a -> Maybe b
f a
a) (Seq a -> Seq b
go Seq a
as)
#endif