{-# 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           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
striping out values or failure.

In order to be Compactable, the following law should hold:

[/Kleisli composition/]

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

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

[/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 identity 1/]

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

[/Applicative identity 2/]

    @applyMaybe (pure Just) = id@

[/Applicative relation/]

    @compact = applyMaybe (pure id)@

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

[/Monad nameme/]

    @bindMaybe (return (Just x)) return = return x@

[/Monad identity 1/]

    @flip bindMaybe (return . Just) = id@

[/Monad identity 2/]

    @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 name me/]

    @traverseMaybe (pure . Just) = pure@

[/Traversable composition/]

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

[/Traversable name me/]

    @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 #-}

instance Compactable Maybe where
    compact = join
    {-# INLINABLE compact #-}
    fmapMaybe f (Just x) = f x
    fmapMaybe _ _        = Nothing
    {-# 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 #-}

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 ( 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 #-}

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

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

instance Compactable Vector

instance Compactable (Const r) where
    compact (Const r) = Const r
    {-# INLINABLE compact #-}

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


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


filter :: (Compactable f, Functor f) => (a -> Bool) -> f a -> f a
filter f = fmapMaybe $ \a -> if f a then Just a else Nothing


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)