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

import Prelude hiding (filter)

import Control.Applicative
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.Sum
import Data.Proxy
import Data.Traversable

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 f :: a -> Maybe b
f = f (Maybe b) -> f b
forall (f :: * -> *) a. Filtrable f => f (Maybe a) -> f a
catMaybes (f (Maybe b) -> f b) -> (f a -> f (Maybe b)) -> f a -> f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe b) -> f a -> f (Maybe b)
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 = (Maybe a -> Maybe a) -> f (Maybe a) -> f a
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe Maybe a -> Maybe a
forall a. a -> a
id

    -- | Drop the elements for which the given predicate is 'False'.
    filter :: (a -> Bool) -> f a -> f a
    filter f :: a -> Bool
f = (a -> Maybe a) -> f a -> f a
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (a -> Maybe () -> Maybe a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
(<$) (a -> Maybe () -> Maybe a) -> (a -> Maybe ()) -> a -> Maybe a
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> (a -> Bool) -> a -> Maybe ()
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 f :: a -> p (Maybe b)
f xs :: f a
xs = f (Maybe b) -> f b
forall (f :: * -> *) a. Filtrable f => f (Maybe a) -> f a
catMaybes (f (Maybe b) -> f b) -> p (f (Maybe b)) -> p (f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> p (Maybe b)) -> f a -> p (f (Maybe 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 f :: a -> p Bool
f = (a -> p (Maybe a)) -> f a -> p (f a)
forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA (\ x :: a
x -> (a
x a -> Maybe () -> Maybe a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$) (Maybe () -> Maybe a) -> (Bool -> Maybe ()) -> Bool -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe a) -> p Bool -> p (Maybe a)
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 f :: a -> Either b c
f = (,) (f b -> f c -> (f b, f c))
-> (f a -> f b) -> f a -> f c -> (f b, f c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> Maybe b) -> f a -> f b
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe ((b -> Maybe b) -> (c -> Maybe b) -> Either b c -> Maybe b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either b -> Maybe b
forall a. a -> Maybe a
Just (Maybe b -> c -> Maybe b
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe b
forall a. Maybe a
Nothing) (Either b c -> Maybe b) -> (a -> Either b c) -> a -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either b c
f)
                      (f a -> f c -> (f b, f c)) -> (f a -> f c) -> f a -> (f b, f c)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> Maybe c) -> f a -> f c
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe ((b -> Maybe c) -> (c -> Maybe c) -> Either b c -> Maybe c
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe c -> b -> Maybe c
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe c
forall a. Maybe a
Nothing) c -> Maybe c
forall a. a -> Maybe a
Just (Either b c -> Maybe c) -> (a -> Either b c) -> a -> Maybe c
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 f :: a -> p (Either b c)
f = (f b -> f c -> (f b, f c)) -> p (f b) -> p (f c) -> p (f b, f c)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) (p (f b) -> p (f c) -> p (f b, f c))
-> (f a -> p (f b)) -> f a -> p (f c) -> p (f b, f c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> p (Maybe b)) -> f a -> p (f b)
forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA ((Either b c -> Maybe b) -> p (Either b c) -> p (Maybe b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (b -> Maybe b
forall a. a -> Maybe a
Just (b -> Maybe b) -> (c -> Maybe b) -> Either b c -> Maybe b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
`either` Maybe b -> c -> Maybe b
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe b
forall a. Maybe a
Nothing) (p (Either b c) -> p (Maybe b))
-> (a -> p (Either b c)) -> a -> p (Maybe b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> p (Either b c)
f)
                              (f a -> p (f c) -> p (f b, f c))
-> (f a -> p (f c)) -> f a -> p (f b, f c)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> p (Maybe c)) -> f a -> p (f c)
forall (f :: * -> *) (p :: * -> *) a b.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p (Maybe b)) -> f a -> p (f b)
mapMaybeA ((Either b c -> Maybe c) -> p (Either b c) -> p (Maybe c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe c -> b -> Maybe c
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe c
forall a. Maybe a
Nothing (b -> Maybe c) -> (c -> Maybe c) -> Either b c -> Maybe c
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
`either` c -> Maybe c
forall a. a -> Maybe a
Just) (p (Either b c) -> p (Maybe c))
-> (a -> p (Either b c)) -> a -> p (Maybe c)
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 = (Either a b -> Either a b) -> f (Either a b) -> (f a, f b)
forall (f :: * -> *) a b c.
Filtrable f =>
(a -> Either b c) -> f a -> (f b, f c)
mapEither Either a b -> Either a b
forall a. a -> a
id

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

    mapMaybeA :: (a -> p (Maybe b)) -> [a] -> p [b]
mapMaybeA _ [] = [b] -> p [b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    mapMaybeA f :: a -> p (Maybe b)
f (x :: a
x:xs :: [a]
xs) = ([b] -> [b]) -> (b -> [b] -> [b]) -> Maybe b -> [b] -> [b]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [b] -> [b]
forall a. a -> a
id (:) (Maybe b -> [b] -> [b]) -> p (Maybe b) -> p ([b] -> [b])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> p (Maybe b)
f a
x p ([b] -> [b]) -> p [b] -> p [b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> p (Maybe b)) -> [a] -> p [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 :: (a -> Maybe b) -> Maybe a -> Maybe b
mapMaybe = (a -> Maybe b) -> Maybe a -> Maybe b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
(=<<)
    catMaybes :: Maybe (Maybe a) -> Maybe a
catMaybes = Maybe (Maybe a) -> Maybe a
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join

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

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

instance (Filtrable f, Filtrable g) => Filtrable (Product f g) where
    mapMaybe :: (a -> Maybe b) -> Product f g a -> Product f g b
mapMaybe f :: a -> Maybe b
f (Pair as :: f a
as bs :: g a
bs) = f b -> g b -> Product f g b
forall k (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
Pair ((a -> Maybe b) -> f a -> f b
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f f a
as) ((a -> Maybe b) -> g a -> g b
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 :: (a -> Maybe b) -> Sum f g a -> Sum f g b
mapMaybe f :: a -> Maybe b
f (InL as :: f a
as) = f b -> Sum f g b
forall k (f :: k -> *) (g :: k -> *) (a :: k). f a -> Sum f g a
InL ((a -> Maybe b) -> f a -> f b
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f f a
as)
    mapMaybe f :: a -> Maybe b
f (InR bs :: g a
bs) = g b -> Sum f g b
forall k (f :: k -> *) (g :: k -> *) (a :: k). g a -> Sum f g a
InR ((a -> Maybe b) -> g a -> g b
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 :: (a -> Maybe b) -> Compose f g a -> Compose f g b
mapMaybe f :: a -> Maybe b
f (Compose as :: f (g a)
as) = f (g b) -> Compose f g b
forall k k1 (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose ((a -> Maybe b) -> g a -> g b
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe a -> Maybe b
f (g a -> g b) -> f (g a) -> f (g b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (g a)
as)

infixl 4 <$?>, <*?>

(<$?>) :: Filtrable f => (a -> Maybe b) -> f a -> f b
<$?> :: (a -> Maybe b) -> f a -> f b
(<$?>) = (a -> Maybe b) -> f a -> f b
forall (f :: * -> *) a b.
Filtrable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe

(<*?>) :: (Applicative p, Filtrable p) => p (a -> Maybe b) -> p a -> p b
f :: p (a -> Maybe b)
f <*?> :: p (a -> Maybe b) -> p a -> p b
<*?> a :: p a
a = p (Maybe b) -> p b
forall (f :: * -> *) a. Filtrable f => f (Maybe a) -> f a
catMaybes (p (a -> Maybe b)
f p (a -> Maybe b) -> p a -> p (Maybe b)
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 :: f a -> f a
nub = (a -> a -> Bool) -> f a -> f a
forall (f :: * -> *) a.
(Filtrable f, Traversable f) =>
(a -> a -> Bool) -> f a -> f a
nubBy a -> a -> Bool
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 :: (a -> a -> Bool) -> f a -> f a
nubBy eq :: a -> a -> Bool
eq = (State [a] (f a) -> f a) -> (f a -> State [a] (f a)) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((State [a] (f a) -> [a] -> f a) -> [a] -> State [a] (f a) -> f a
forall a b c. (a -> b -> c) -> b -> a -> c
flip State [a] (f a) -> [a] -> f a
forall s a. State s a -> s -> a
M.evalState []) ((f a -> State [a] (f a)) -> f a -> f a)
-> ((a -> StateT [a] Identity Bool) -> f a -> State [a] (f a))
-> (a -> StateT [a] Identity Bool)
-> f a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> StateT [a] Identity Bool) -> f a -> State [a] (f a)
forall (f :: * -> *) (p :: * -> *) a.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p Bool) -> f a -> p (f a)
filterA ((a -> StateT [a] Identity Bool) -> f a -> f a)
-> (a -> StateT [a] Identity Bool) -> f a -> f a
forall a b. (a -> b) -> a -> b
$ \ a :: a
a -> do
    [a]
as <- StateT [a] Identity [a]
forall (m :: * -> *) s. Monad m => StateT s m s
M.get
    let b :: Bool
b = (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> Bool
not (Bool -> Bool) -> (a -> Bool) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> Bool
eq a
a) [a]
as
    Bool
b Bool -> StateT [a] Identity () -> StateT [a] Identity Bool
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> StateT [a] Identity () -> StateT [a] Identity ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (([a] -> [a]) -> StateT [a] Identity ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
M.modify (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:))

-- | \(\mathcal{O}(n^2)\)
-- 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 :: f a -> f a
nubOrd = (a -> a -> Ordering) -> f a -> f a
forall (f :: * -> *) a.
(Filtrable f, Traversable f) =>
(a -> a -> Ordering) -> f a -> f a
nubOrdBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

-- | \(\mathcal{O}(n^2)\)
-- 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 :: (a -> a -> Ordering) -> f a -> f a
nubOrdBy compare :: a -> a -> Ordering
compare = (State (Set a) (f a) -> f a)
-> (f a -> State (Set a) (f a)) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((State (Set a) (f a) -> Set a -> f a)
-> Set a -> State (Set a) (f a) -> f a
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (Set a) (f a) -> Set a -> f a
forall s a. State s a -> s -> a
M.evalState Set a
forall a. Set a
Set.empty) ((f a -> State (Set a) (f a)) -> f a -> f a)
-> ((a -> StateT (Set a) Identity Bool)
    -> f a -> State (Set a) (f a))
-> (a -> StateT (Set a) Identity Bool)
-> f a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> StateT (Set a) Identity Bool) -> f a -> State (Set a) (f a)
forall (f :: * -> *) (p :: * -> *) a.
(Filtrable f, Traversable f, Applicative p) =>
(a -> p Bool) -> f a -> p (f a)
filterA ((a -> StateT (Set a) Identity Bool) -> f a -> f a)
-> (a -> StateT (Set a) Identity Bool) -> f a -> f a
forall a b. (a -> b) -> a -> b
$ \ a :: a
a -> (Set a -> (Bool, Set a)) -> StateT (Set a) Identity Bool
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
M.state ((Set a -> (Bool, Set a)) -> StateT (Set a) Identity Bool)
-> (Set a -> (Bool, Set a)) -> StateT (Set a) Identity Bool
forall a b. (a -> b) -> a -> b
$ \ as :: Set a
as ->
    case (a -> a -> Ordering) -> a -> Set a -> Maybe (Set a)
forall a. (a -> a -> Ordering) -> a -> Set a -> Maybe (Set a)
Set.insertBy' a -> a -> Ordering
compare a
a Set a
as of
        Nothing -> (Bool
False, Set a
as)
        Just as' :: Set a
as' -> (Bool
True, Set a
as')