Selective.hs

{-# LANGUAGE TypeOperators, PatternSynonyms, ExplicitNamespaces #-}
{-# LANGUAGE LambdaCase, BlockArguments #-}

module Select
  ( type (-?)(Fun, Const, Lazy, unLazy), ($?)
  , Selective(..)
  , select, branch, whenS, ifS, whileS, fromMaybeS
  , (<||>), (<&&>), anyS, allS
  , Monad(..)
  ) where

import Prelude hiding (Monad(..), id, (.))
import Control.Category
import Control.Applicative (liftA2)


newtype a -? b = Lazy { unLazy :: Either (a -> b) b }
infixr 0 -?

{-# COMPLETE Const, Fun #-}
pattern Const ::       b  -> (a -? b)
pattern Fun   :: (a -> b) -> (a -? b)
pattern Const b = Lazy (Right b)
pattern Fun   f = Lazy (Left f)

instance Functor ((-?) a) where
  fmap = (.) . Fun

instance Applicative ((-?) a) where
  pure = Const
  Const f <*> Const b = Const (f b)
  g       <*> x       = Fun \a -> (g $? a) (x $? a)

instance Selective ((-?) a) where
  Const (Const b) <*? _ = Const b
  g               <*? x = ($?) <$> g <*> x

instance Monad ((-?) a) where
  (>>=) = \case
    Const a -> \f -> f a
    Fun   g -> \f -> Fun \e -> f (g e) $? e

instance Semigroup b => Semigroup (a -? b) where
  (<>) = liftA2 (<>)

instance Monoid b => Monoid (a -? b) where
  mempty = pure mempty

instance Category (-?) where
  id = Fun id
  Const c . _       = Const c
  Fun f   . Const b = Const (f b)
  Fun f   . Fun g   = Fun (f . g)

($?) :: (a -? b) -> (a -> b)
Const b $? _ = b
Fun f   $? a = f a
infixr 0 $?

-- | Steal laws straight from Applicative?
--
--   Identity:
--     pure id <*? x = x
--
--   Composition:
--     pure (.) <*? u <*? v <*? w = u <*? (v <*? w)
--
--   Homomorphism:
--     pure f <*? pure x = pure (f $? x)
--
--   Interchange:
--     u <*? pure y = pure ($? y) <*? u
--
class Applicative f => Selective f where
  {-# MINIMAL (<*?) | liftS2 #-}

  (<*?) :: f (a -? b) -> f a -> f b
  (<*?) = liftS2 id
  infixl 4 <*?

  liftS2 :: (a -? b -? c) -> f a -> f b -> f c
  liftS2 g fa fb = pure g <*? fa <*? fb

select :: Selective f => f (Either a b) -> f (a -> b) -> f b
select feab fatob = lazyMakeB <$> feab <*? fatob
  where
    lazyMakeB :: Either a b -> (a -> b) -? b
    lazyMakeB = \case
      Left  a -> Fun \f -> f a
      Right b -> Const b

branch :: Selective f => f (Either a b) -> f (a -> c) -> f (b -> c) -> f c
branch feab fatoc fbtoc = lazyMakeC <$> feab <*? fatoc <*? fbtoc
  where
    lazyMakeC :: Either a b -> (a -> c) -? (b -> c) -? c
    lazyMakeC = \case
      Left  a -> Fun \f -> Const (f a)
      Right b -> Const (Fun \g -> g b)

whenS :: Selective f => f Bool -> f () -> f ()
whenS fb fu = lazyIfFalse <$> fb <*? fu
  where
    lazyIfFalse :: Bool -> () -? ()
    lazyIfFalse = \case
      False -> Const ()
      True  -> id

ifS :: Selective f => f Bool -> f a -> f a -> f a
ifS fb fa1 fa2 = lazyIf <$> fb <*? fa1 <*? fa2
  where
    lazyIf :: Bool -> a -? a -? a
    lazyIf = \case
      True  -> Fun Const
      False -> Const id

whileS :: Selective f => f Bool -> f ()
whileS fb = whenS fb (whileS fb)

-- NB: if both effects run, those of the second argument run first.
fromMaybeS :: Selective f => f a -> f (Maybe a) -> f a
fromMaybeS fa fma = lazyFromMaybe <$> fma <*? fa
  where
    lazyFromMaybe :: Maybe a -> a -? a
    lazyFromMaybe = \case
      Nothing -> id
      Just a  -> Const a

(<||>) :: Selective f => f Bool -> f Bool -> f Bool
b1 <||> b2 = ifS b1 (pure True) b2

(<&&>) :: Selective f => f Bool -> f Bool -> f Bool
b1 <&&> b2 = ifS b1 b2 (pure False)

anyS :: Selective f => (a -> f Bool) -> [a] -> f Bool
anyS p = foldr (\a b -> p a <||> b) (pure False)

allS :: Selective f => (a -> f Bool) -> [a] -> f Bool
allS p = foldr (\a b -> p a <&&> b) (pure True)

class Selective f => Monad f where
  (>>=) :: f a -> (a -> f b) -> f b