{-# LANGUAGE ConstrainedClassMethods #-}
{-# LANGUAGE DefaultSignatures       #-}
{-# LANGUAGE KindSignatures          #-}

module Control.Compactable  where

import           Control.Applicative
import           Control.Monad                   (join)
import           Control.Monad.Trans.Maybe
import           Data.Functor.Compose
import qualified Data.Functor.Product            as FP
import qualified Data.IntMap                     as IntMap
import qualified Data.Map                        as Map
import           Data.Maybe
import           Data.Proxy
import           Data.Semigroup
import qualified Data.Sequence                   as Seq
import qualified Data.Set                        as Set
import           Data.Vector                     (Vector)
import           GHC.Conc
import           Text.ParserCombinators.ReadP
import           Text.ParserCombinators.ReadPrec

{-|
This is a generalization of catMaybes as a new function compact. Compact
has relations with Functor, Applicative, Monad, Alternative, and Traversable.
In that we can use these class to provide the ability to operate on a data type
by throwing away intermediate Nothings. This is useful for representing
stripping out values or failure.

To be compactable alone, no laws must be satisfied other than the type signature.

If the data type is also a Functor the following should hold:

[/Kleisli composition/]

    @fmapMaybe (l <=< r) = fmapMaybe l . fmapMaybe r@

[/Functor identity 1/]

    @compact . fmap Just = id@

[/Functor identity 2/]

    @fmapMaybe Just = id@

[/Functor relation/]

    @compact = fmapMaybe id@

According to Kmett, (Compactable f, Functor f) is a functor from the
kleisli category of Maybe to the category of haskell data types.
@Kleisli Maybe -> Hask@.

If the data type is also Applicative the following should hold:

[/Applicative left identity/]

    @compact . (pure Just \<*\>) = id@

[/Applicative right identity/]

    @applyMaybe (pure Just) = id@

[/Applicative relation/]

    @compact = applyMaybe (pure id)@

If the data type is also a Monad the following should hold:

[/Monad left identity/]

    @flip bindMaybe (return . Just) = id@

[/Monad right identity/]

    @compact . (return . Just =<<) = id@

[/Monad relation/]

    @compact = flip bindMaybe return@

If the data type is also Alternative the following should hold:

[/Alternative identity/]

    @compact empty = empty@

[/Alternative annihilation/]

    @compact (const Nothing \<$\> xs) = empty@

If the data type is also Traversable the following should hold:

[/Traversable Applicative relation/]

    @traverseMaybe (pure . Just) = pure@

[/Traversable composition/]

    @Compose . fmap (traverseMaybe f) . traverseMaybe g = traverseMaybe (Compose . fmap (traverseMaybe f) . g)@

[/Traversable Functor relation/]

    @traverse f = traverseMaybe (fmap Just . f)@

[/Traversable naturality/]

    @t . traverseMaybe f = traverseMaybe (t . f)@

If you know of more useful laws, or have better names for the ones above
(especially those marked "name me"). Please let me know.
-}

class Compactable (f :: * -> *) where
    compact :: f (Maybe a) -> f a
    default compact :: (Monad f, Alternative f) => f (Maybe a) -> f a
    compact = (>>= maybe empty return)
    {-# INLINABLE compact #-}

    fmapMaybe :: Functor f => (a -> Maybe b) -> f a -> f b
    fmapMaybe f = compact . fmap f
    {-# INLINABLE fmapMaybe #-}

    applyMaybe :: Applicative f => f (a -> Maybe b) -> f a -> f b
    applyMaybe fa = compact . (fa <*>)
    {-# INLINABLE applyMaybe #-}

    bindMaybe :: Monad f => f a -> (a -> f (Maybe b)) -> f b
    bindMaybe x = compact . (x >>=)
    {-# INLINABLE bindMaybe #-}

    traverseMaybe :: (Applicative g, Traversable f)
                  => (a -> g (Maybe b)) -> f a -> g (f b)
    traverseMaybe f = fmap compact . traverse f
    {-# INLINABLE traverseMaybe #-}

    filter :: (a -> Bool) -> f a -> f a
    default filter :: Functor f => (a -> Bool) -> f a -> f a
    filter f = fmapMaybe $ \a -> if f a then Just a else Nothing
    {-# INLINABLE filter #-}

instance Compactable Maybe where
    compact = join
    {-# INLINABLE compact #-}
    fmapMaybe = (=<<)
    {-# INLINABLE fmapMaybe #-}

instance Compactable [] where
    compact = catMaybes
    {-# INLINABLE compact #-}
    fmapMaybe _ []    = []
    fmapMaybe f (h:t) = maybe (fmapMaybe f t) (: fmapMaybe f t) (f h)
    {-# INLINABLE fmapMaybe #-}
    filter = Prelude.filter
    {-# INLINABLE filter #-}

instance Compactable IO

instance Compactable STM

instance Compactable Proxy

instance Compactable Option where
    compact (Option x) = Option (join x)
    {-# INLINABLE compact #-}
    fmapMaybe f (Option (Just x)) = Option (f x)
    fmapMaybe _ _                 = Option Nothing
    {-# INLINABLE fmapMaybe #-}

instance Compactable ReadP

instance Compactable ReadPrec

instance ( Functor f, Functor g, Compactable f, Compactable g )
         => Compactable (FP.Product f g) where
    compact (FP.Pair x y) = FP.Pair (compact x) (compact y)
    {-# INLINABLE compact #-}

instance ( Functor f, Functor g, Compactable g )
         => Compactable (Compose f g) where
    compact = fmapMaybe id
    {-# INLINABLE compact #-}
    fmapMaybe f (Compose fg) = Compose $ fmap (fmapMaybe f) fg
    {-# INLINABLE fmapMaybe #-}

instance Compactable IntMap.IntMap where
    compact = IntMap.mapMaybe id
    {-# INLINABLE compact #-}
    fmapMaybe = IntMap.mapMaybe
    {-# INLINABLE fmapMaybe #-}
    filter = IntMap.filter
    {-# INLINABLE filter #-}

instance Compactable (Map.Map k) where
    compact = Map.mapMaybe id
    {-# INLINABLE compact #-}
    fmapMaybe = Map.mapMaybe
    {-# INLINABLE fmapMaybe #-}
    filter = Map.filter
    {-# INLINABLE filter #-}

instance Compactable Seq.Seq where
    compact = fmap fromJust . Seq.filter isJust
    {-# INLINABLE compact #-}
    filter = Seq.filter
    {-# INLINABLE filter #-}

instance Compactable Vector

instance Compactable (Const r) where
    compact (Const r) = Const r
    {-# INLINABLE compact #-}
    fmapMaybe _ (Const r) = Const r
    {-# INLINABLE fmapMaybe #-}
    bindMaybe (Const r) _ = Const r
    {-# INLINABLE bindMaybe #-}
    filter _ (Const r) = Const r
    {-# INLINABLE filter #-}

instance Monoid m => Compactable (Either m) where
    compact (Right (Just x)) = Right x
    compact (Right _)        = Left mempty
    compact (Left x)         = Left x
    {-# INLINABLE compact #-}
    fmapMaybe f (Right x) = maybe (Left mempty) Right (f x)
    fmapMaybe _ (Left x)  = Left x
    {-# INLINABLE fmapMaybe #-}

instance Compactable Set.Set where
    compact = Set.fromDistinctAscList . compact . Set.toAscList
    {-# INLINABLE compact #-}
    filter f = Set.fromDistinctAscList . Control.Compactable.filter f . Set.toAscList
    {-# INLINABLE filter  #-}


fforMaybe :: (Compactable f, Functor f) => f a -> (a -> Maybe b) -> f b
fforMaybe = flip fmapMaybe


fmapMaybeM :: (Compactable f, Monad f) => (a -> MaybeT f b) -> f a -> f b
fmapMaybeM f = (>>= compact . runMaybeT . f)


fforMaybeM :: (Compactable f, Monad f) => f a -> (a -> MaybeT f b) -> f b
fforMaybeM = flip fmapMaybeM


applyMaybeM :: (Compactable f, Monad f) => f (a -> MaybeT f b) -> f a -> f b
applyMaybeM fa = compact . join . fmap runMaybeT . (fa <*>)


bindMaybeM :: (Compactable f, Monad f) => f a -> (a -> f (MaybeT f b)) -> f b
bindMaybeM x = compact . join . fmap runMaybeT . (x >>=)


traverseMaybeM :: (Monad m, Compactable t, Traversable t) => (a -> MaybeT m b) -> t a -> m (t b)
traverseMaybeM f = unwrapMonad . traverseMaybe (WrapMonad . runMaybeT . f)