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

module Control.Compactable
  ( Compactable (..)
  , fforMaybe
  , fforEither
  , fmapMaybeM
  , fmapEitherM
  , fforMaybeM
  , fforEitherM
  , applyMaybeM
  , bindMaybeM
  , traverseMaybeM
  , altDefaultCompact
  , altDefaultSeparate ) where

import           Control.Applicative
import           Control.Monad                   (join)
import           Control.Monad.Trans.Except
import           Control.Monad.Trans.Maybe
import           Data.Bifunctor                  (bimap)
import           Data.Either                     (partitionEithers)
import           Data.Foldable                   (foldl')
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 qualified Data.Vector                     as V
import           GHC.Conc
import           Text.ParserCombinators.ReadP
import           Text.ParserCombinators.ReadPrec

{-|
Class 'Compactable' provides two methods which can be writen in terms of each other, compact and separate.

=Compact
is generalization of catMaybes as a new function. 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)@

= Separate and filter
have recently elevated roles in this typeclass, and is not as well explored as compact. Here are the laws known today:

[/Functor identity 3/]

    @fst . separate . fmap Right = id@

[/Functor identity 4/]

    @snd . separate . fmap Left = id@

[/Applicative left identity 2/]

    @snd . separate . (pure Right \<*\>) = id@

[/Applicative right identity 2/]

    @fst . separate . (pure Left \<*\>) = id@

[/Alternative annihilation left/]

    @snd . separate . fmap (const Left) = empty@

[/Alternative annihilation right/]

    @fst , separate . fmap (const Right) = empty@

Docs for relationships between these functions and, a cleanup of laws will happen at some point.

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 :: Functor f => f (Maybe a) -> f a
    compact = snd . separate . fmap (\case Just x -> Right x; _ -> Left ())
    {-# INLINABLE compact #-}

    separate :: f (Either l r) -> (f l, f r)
    default separate :: Functor f => f (Either l r) -> (f l, f r)
    separate xs = (compact $ hush . flipEither <$> xs, compact $ hush <$> xs)
    {-# INLINABLE separate #-}

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

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

    fmapEither :: Functor f => (a -> Either l r) -> f a -> (f l, f r)
    fmapEither f = separate . fmap f
    {-# INLINABLE fmapEither #-}

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

    applyEither :: Applicative f => f (a -> Either l r) -> f a -> (f l, f r)
    applyEither fa = separate . (fa <*>)
    {-# INLINABLE applyEither #-}

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

    bindEither :: Monad f => f a -> (a -> f (Either l r)) -> (f l, f r)
    bindEither x = separate . (x >>=)
    {-# INLINABLE bindEither #-}

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

    traverseEither :: (Applicative g, Traversable f)
                   => (a -> g (Either l r)) -> f a -> g (f l, f r)
    traverseEither f = fmap separate . traverse f
    {-# INLINABLE traverseEither #-}


instance Compactable Maybe where
    compact = join
    {-# INLINABLE compact #-}
    fmapMaybe = (=<<)
    {-# INLINABLE fmapMaybe #-}
    separate = \case
      Just x -> case x of
         Left l  -> (Just l, Nothing)
         Right r -> (Nothing, Just r)
      _ -> (Nothing, Nothing)
    {-# INLINABLE separate #-}

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 #-}
    separate = \case
      Right (Left l) -> (Right l, Left mempty)
      Right (Right r) -> (Left mempty, Right r)
      Left x -> (Left x, Left x)
    {-# INLINABLE separate #-}

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 #-}
    separate = partitionEithers
    {-# INLINABLE separate #-}
    fmapEither f = foldl' (deal f) ([],[])
      where deal g ~(bs, cs) a = case g a of
             Left b  -> (b:bs, cs)
             Right c -> (bs, c:cs)
    {-# INLINABLE fmapEither #-}
    traverseMaybe f = go where
      go (x:xs) = maybe id (:) <$> f x <*> go xs
      go []     = pure []
    {-# INLINE traverseMaybe #-}

instance Compactable IO where
    compact = altDefaultCompact
    {-# INLINABLE compact #-}

instance Compactable STM where
    compact = altDefaultCompact
    {-# INLINABLE compact #-}

instance Compactable Proxy where
    compact _ = Proxy
    {-# INLINABLE compact #-}
    separate _ = (Proxy, Proxy)
    {-# INLINABLE separate #-}
    filter _ _ = Proxy
    {-# INLINABLE filter #-}
    fmapMaybe _ _ = Proxy
    {-# INLINABLE fmapMaybe #-}
    applyMaybe _ _ = Proxy
    {-# INLINABLE applyMaybe #-}
    bindMaybe _ _ = Proxy
    {-# INLINABLE bindMaybe #-}

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 #-}
    separate = altDefaultSeparate
    {-# INLINABLE separate #-}

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 $ 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 #-}
    separate = IntMap.mapEither id
    {-# INLINABLE separate #-}
    fmapEither = IntMap.mapEither
    {-# INLINABLE fmapEither #-}

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

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

instance Compactable V.Vector where
    compact = altDefaultCompact
    {-# INLINABLE compact #-}
    separate = altDefaultSeparate
    {-# INLINABLE separate #-}
    filter = V.filter
    {-# INLINABLE filter #-}

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 Compactable Set.Set where
    compact = Set.fromDistinctAscList . compact . Set.toAscList
    {-# INLINABLE compact #-}
    separate = bimap Set.fromDistinctAscList Set.fromDistinctAscList . separate . Set.toAscList
    {-# INLINABLE separate #-}
    filter = Set.filter
    {-# INLINABLE filter  #-}


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


fforEither :: (Compactable f, Functor f) => f a -> (a -> Either l r) -> (f l, f r)
fforEither = flip fmapEither


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


fmapEitherM :: (Compactable f, Monad f) => (a -> ExceptT l f r) -> f a -> (f l, f r)
fmapEitherM f x = separate $ runExceptT . f =<< x


fforEitherM :: (Compactable f, Monad f) => f a -> (a -> ExceptT l f r) -> (f l, f r)
fforEitherM = flip fmapEitherM


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)

-- | While more constrained, when available, this default is going to be faster than the one provided in the typeclass
altDefaultCompact :: (Alternative f, Monad f) => f (Maybe a) -> f a
altDefaultCompact = (>>= maybe empty return)
{-# INLINABLE altDefaultCompact #-}

-- | While more constrained, when available, this default is going to be faster than the one provided in the typeclass
altDefaultSeparate :: (Alternative f, Foldable f) => f (Either l r) -> (f l, f r)
altDefaultSeparate = foldl' (\(l', r') -> \case
  Left l  -> (l' <|> pure l ,r')
  Right r -> (l', r' <|> pure r)) (empty, empty)
{-# INLINABLE altDefaultSeparate #-}


hush :: Either l r -> Maybe r
hush = \case (Right x) -> Just x; _ -> Nothing


flipEither :: Either a b -> Either b a
flipEither = \case (Right x) -> Left x; (Left x) -> Right x