    {-# LANGUAGE GADTs, RankNTypes, TupleSections #-}
    -----------------------------------------------------------------------------
    -- |
    -- Module     : Control.Selective.Rigid.Free
    -- Copyright  : (c) Andrey Mokhov 2018-2024
    -- License    : MIT (see the file LICENSE)
    -- Maintainer : andrey.mokhov@gmail.com
    -- Stability  : experimental
    --
    -- This is a library for /selective applicative functors/, or just
    -- /selective functors/ for short, an abstraction between applicative functors
    -- and monads, introduced in this paper: https://dl.acm.org/doi/10.1145/3341694.
    --
    -- This module defines /free rigid selective functors/. Rigid selective functors
    -- are those that satisfy the property @\<*\> = apS@.
    --
    -- Intuitively, a selective functor @f@ is "rigid" if any expression @f a@ is
    -- equivalent to a list of effects chained with @select@ operators (the normal
    -- form given by the free construction). In contrast, "non-rigid" selective
    -- functors can have non-linear, tree-like shapes, because @<*>@ nodes can't be
    -- straightened using the @\<*\> = apS@ equation.
    --
    -----------------------------------------------------------------------------
    module Control.Selective.Rigid.Free (
        -- * Free rigid selective functors
        Select (..), liftSelect,

        -- * Static analysis
        getPure, getEffects, getNecessaryEffect, runSelect, foldSelect
        ) where

    import Control.Selective.Trans.Except
    import Control.Selective
    import Data.Bifunctor
    import Data.Functor

    -- Inspired by free applicative functors by Capriotti and Kaposi.
    -- See: https://arxiv.org/pdf/1403.0749.pdf

    -- TODO: The current approach is simple but very slow: 'fmap' costs O(N), where
    -- N is the number of effects, and 'select' is even worse -- O(N^2). It is
    -- possible to improve both bounds to O(1) by using the idea developed for free
    -- applicative functors by Dave Menendez. See this blog post:
    -- https://www.eyrie.org/~zednenem/2013/05/27/freeapp
    -- An example implementation can be found here:
    -- http://hackage.haskell.org/package/free/docs/Control-Applicative-Free-Fast.html

    -- | Free rigid selective functors.
    data Select f a where
        Pure   :: a -> Select f a
        Select :: Select f (Either a b) -> f (a -> b) -> Select f b

    -- TODO: Prove that this is a lawful 'Functor'.
    instance Functor f => Functor (Select f) where
        fmap :: forall a b. (a -> b) -> Select f a -> Select f b
    fmap a -> b
    f (Pure a
    a)     = b -> Select f b
    forall a (f :: * -> *). a -> Select f a
    Pure (a -> b
    f a
    a)
        fmap a -> b
    f (Select Select f (Either a a)
    x f (a -> a)
    y) = Select f (Either a b) -> f (a -> b) -> Select f b
    forall (f :: * -> *) a b.
    Select f (Either a b) -> f (a -> b) -> Select f b
    Select ((a -> b) -> Either a a -> Either a b
    forall a b. (a -> b) -> Either a a -> Either a b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap a -> b
    f (Either a a -> Either a b)
    -> Select f (Either a a) -> Select f (Either a b)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Select f (Either a a)
    x) ((a -> b) -> (a -> a) -> a -> b
    forall a b. (a -> b) -> (a -> a) -> a -> b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap a -> b
    f ((a -> a) -> a -> b) -> f (a -> a) -> f (a -> b)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> f (a -> a)
    y)

    -- TODO: Prove that this is a lawful 'Applicative'.
    instance Functor f => Applicative (Select f) where
        pure :: forall a. a -> Select f a
    pure  = a -> Select f a
    forall a (f :: * -> *). a -> Select f a
    Pure
        <*> :: forall a b. Select f (a -> b) -> Select f a -> Select f b
    (<*>) = Select f (a -> b) -> Select f a -> Select f b
    forall (f :: * -> *) a b. Selective f => f (a -> b) -> f a -> f b
    apS -- Rigid selective functors

    -- TODO: Prove that this is a lawful 'Selective'.
    instance Functor f => Selective (Select f) where
        -- Identity law
        select :: forall a b.
    Select f (Either a b) -> Select f (a -> b) -> Select f b
    select Select f (Either a b)
    x (Pure a -> b
    y) = (a -> b) -> (b -> b) -> Either a b -> b
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
    either a -> b
    y b -> b
    forall a. a -> a
    id (Either a b -> b) -> Select f (Either a b) -> Select f b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Select f (Either a b)
    x

        -- Associativity law
        select Select f (Either a b)
    x (Select Select f (Either a (a -> b))
    y f (a -> a -> b)
    z) = Select f (Either (a, a) b) -> f ((a, a) -> b) -> Select f b
    forall (f :: * -> *) a b.
    Select f (Either a b) -> f (a -> b) -> Select f b
    Select (Select f (Either a (Either (a, a) b))
    -> Select f (a -> Either (a, a) b) -> Select f (Either (a, a) b)
    forall a b.
    Select f (Either a b) -> Select f (a -> b) -> Select f b
    forall (f :: * -> *) a b.
    Selective f =>
    f (Either a b) -> f (a -> b) -> f b
    select (Either a b -> Either a (Either (a, a) b)
    forall {a} {a}. Either a a -> Either a (Either a a)
    f (Either a b -> Either a (Either (a, a) b))
    -> Select f (Either a b) -> Select f (Either a (Either (a, a) b))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Select f (Either a b)
    x) (Either a (a -> b) -> a -> Either (a, a) b
    forall {p :: * -> * -> *} {a} {t} {d}.
    Bifunctor p =>
    p a (t -> d) -> t -> p (a, t) d
    g (Either a (a -> b) -> a -> Either (a, a) b)
    -> Select f (Either a (a -> b)) -> Select f (a -> Either (a, a) b)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Select f (Either a (a -> b))
    y)) ((a -> a -> b) -> (a, a) -> b
    forall {a} {b} {c}. (a -> b -> c) -> (a, b) -> c
    h ((a -> a -> b) -> (a, a) -> b)
    -> f (a -> a -> b) -> f ((a, a) -> b)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> f (a -> a -> b)
    z)
          where
            f :: Either a a -> Either a (Either a a)
    f     = (a -> Either a a) -> Either a a -> Either a (Either a a)
    forall a b. (a -> b) -> Either a a -> Either a b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap a -> Either a a
    forall a b. b -> Either a b
    Right
            g :: p a (t -> d) -> t -> p (a, t) d
    g p a (t -> d)
    y t
    a = (a -> (a, t)) -> ((t -> d) -> d) -> p a (t -> d) -> p (a, t) d
    forall a b c d. (a -> b) -> (c -> d) -> p a c -> p b d
    forall (p :: * -> * -> *) a b c d.
    Bifunctor p =>
    (a -> b) -> (c -> d) -> p a c -> p b d
    bimap (,t
    a) ((t -> d) -> t -> d
    forall a b. (a -> b) -> a -> b
    $t
    a) p a (t -> d)
    y
            h :: (a -> b -> c) -> (a, b) -> c
    h     = (a -> b -> c) -> (a, b) -> c
    forall {a} {b} {c}. (a -> b -> c) -> (a, b) -> c
    uncurry

    {- The following can be used in the above implementation as select = selectOpt.

    -- An optimised implementation of select for the free instance. It accumulates
    -- the calls to fmap on the @y@ parameter to avoid traversing the list on every
    -- recursive step.
    selectOpt :: Functor f => Select f (Either a b) -> Select f (a -> b) -> Select f b
    selectOpt x y = go x y id

    -- We turn @Select f (a -> b)@ to @(Select f c, c -> (a -> b))@. Hey, co-Yoneda!
    go :: Functor f => Select f (Either a b) -> Select f c -> (c -> (a -> b)) -> Select f b
    go x (Pure y)     k = either (k y) id <$> x
    go x (Select y z) k = Select (go (f <$> x) y (g . second k)) ((h . (k.)) <$> z)
      where
        f x = Right <$> x
        g y = \a -> bimap (,a) ($a) y
        h z = uncurry z
    -}

    -- | Lift a functor into a free selective computation.
    liftSelect :: Functor f => f a -> Select f a
    liftSelect :: forall (f :: * -> *) a. Functor f => f a -> Select f a
    liftSelect f a
    f = Select f (Either () a) -> f (() -> a) -> Select f a
    forall (f :: * -> *) a b.
    Select f (Either a b) -> f (a -> b) -> Select f b
    Select (Either () a -> Select f (Either () a)
    forall a (f :: * -> *). a -> Select f a
    Pure (() -> Either () a
    forall a b. a -> Either a b
    Left ())) (a -> () -> a
    forall a b. a -> b -> a
    const (a -> () -> a) -> f a -> f (() -> a)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> f a
    f)

    -- | Given a natural transformation from @f@ to @g@, this gives a canonical
    -- natural transformation from @Select f@ to @g@.
    runSelect :: Selective g => (forall x. f x -> g x) -> Select f a -> g a
    runSelect :: forall (g :: * -> *) (f :: * -> *) a.
    Selective g =>
    (forall x. f x -> g x) -> Select f a -> g a
    runSelect forall x. f x -> g x
    _ (Pure a
    a)     = a -> g a
    forall a. a -> g a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure a
    a
    runSelect forall x. f x -> g x
    t (Select Select f (Either a a)
    x f (a -> a)
    y) = g (Either a a) -> g (a -> a) -> g a
    forall a b. g (Either a b) -> g (a -> b) -> g b
    forall (f :: * -> *) a b.
    Selective f =>
    f (Either a b) -> f (a -> b) -> f b
    select ((forall x. f x -> g x) -> Select f (Either a a) -> g (Either a a)
    forall (g :: * -> *) (f :: * -> *) a.
    Selective g =>
    (forall x. f x -> g x) -> Select f a -> g a
    runSelect f x -> g x
    forall x. f x -> g x
    t Select f (Either a a)
    x) (f (a -> a) -> g (a -> a)
    forall x. f x -> g x
    t f (a -> a)
    y)

    -- | Concatenate all effects of a free selective computation.
    foldSelect :: Monoid m => (forall x. f x -> m) -> Select f a -> m
    foldSelect :: forall m (f :: * -> *) a.
    Monoid m =>
    (forall x. f x -> m) -> Select f a -> m
    foldSelect forall x. f x -> m
    f = Over m a -> m
    forall m a. Over m a -> m
    getOver (Over m a -> m) -> (Select f a -> Over m a) -> Select f a -> m
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (forall x. f x -> Over m x) -> Select f a -> Over m a
    forall (g :: * -> *) (f :: * -> *) a.
    Selective g =>
    (forall x. f x -> g x) -> Select f a -> g a
    runSelect (m -> Over m x
    forall m a. m -> Over m a
    Over (m -> Over m x) -> (f x -> m) -> f x -> Over m x
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . f x -> m
    forall x. f x -> m
    f)

    -- | Extract the resulting value if there are no necessary effects.
    getPure :: Select f a -> Maybe a
    getPure :: forall (f :: * -> *) a. Select f a -> Maybe a
    getPure = (forall x. f x -> Maybe x) -> Select f a -> Maybe a
    forall (g :: * -> *) (f :: * -> *) a.
    Selective g =>
    (forall x. f x -> g x) -> Select f a -> g a
    runSelect (Maybe x -> f x -> Maybe x
    forall a b. a -> b -> a
    const Maybe x
    forall a. Maybe a
    Nothing)

    -- | Collect all possible effects in the order they appear in a free selective
    -- computation.
    getEffects :: Functor f => Select f a -> [f ()]
    getEffects :: forall (f :: * -> *) a. Functor f => Select f a -> [f ()]
    getEffects = (forall x. f x -> [f ()]) -> Select f a -> [f ()]
    forall m (f :: * -> *) a.
    Monoid m =>
    (forall x. f x -> m) -> Select f a -> m
    foldSelect (f () -> [f ()]
    forall a. a -> [a]
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (f () -> [f ()]) -> (f x -> f ()) -> f x -> [f ()]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . f x -> f ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void)

    -- Implementation used in the paper:
    -- getEffects = getOver . runSelect (Over . pure . void)

    -- | Extract the necessary effect from a free selective computation. Note: there
    -- can be at most one effect that is statically guaranteed to be necessary.
    getNecessaryEffect :: Functor f => Select f a -> Maybe (f ())
    getNecessaryEffect :: forall (f :: * -> *) a. Functor f => Select f a -> Maybe (f ())
    getNecessaryEffect = Either (f ()) a -> Maybe (f ())
    forall a b. Either a b -> Maybe a
    leftToMaybe (Either (f ()) a -> Maybe (f ()))
    -> (Select f a -> Either (f ()) a) -> Select f a -> Maybe (f ())
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Except (f ()) a -> Either (f ()) a
    forall e a. Except e a -> Either e a
    runExcept (Except (f ()) a -> Either (f ()) a)
    -> (Select f a -> Except (f ()) a) -> Select f a -> Either (f ()) a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (forall x. f x -> ExceptT (f ()) Identity x)
    -> Select f a -> Except (f ()) a
    forall (g :: * -> *) (f :: * -> *) a.
    Selective g =>
    (forall x. f x -> g x) -> Select f a -> g a
    runSelect (f () -> ExceptT (f ()) Identity x
    forall (m :: * -> *) e a. Monad m => e -> ExceptT e m a
    throwE (f () -> ExceptT (f ()) Identity x)
    -> (f x -> f ()) -> f x -> ExceptT (f ()) Identity x
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . f x -> f ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void)

    leftToMaybe :: Either a b -> Maybe a
    leftToMaybe :: forall a b. Either a b -> Maybe a
    leftToMaybe (Left a
    a) = a -> Maybe a
    forall a. a -> Maybe a
    Just a
    a
    leftToMaybe Either a b
    _        = Maybe a
    forall a. Maybe a
    Nothing