{-# LANGUAGE
  BangPatterns,
  GADTs,
  KindSignatures,
  RankNTypes,
  ScopedTypeVariables,
  StandaloneKindSignatures,
  TypeOperators #-}

-- | Nondeterministic choice as an algebraic effect.
--
-- === Warning: Non-linear continuations
--
-- The handlers 'forAllChoices' and 'toList' call continuations zero or twice.
-- Don't use them to handle a computation that must ensure linear usage of resources.
module Bluefin.Algae.NonDeterminism
  ( -- * Operations
    Choice(..)
  , choose
  , nil
  , assume
  , pick
  , removeFrom
    -- * Handlers
  , forAllChoices
  , toList
  , foldChoice
  ) where

import Control.Monad ((>=>), join)
import Bluefin.Internal (insertFirst)
import Bluefin.Eff (Eff, type (:&), type (:>))
import Bluefin.Algae

-- | Choice effect.
data Choice :: AEffect where
  -- | Choose one of two alternatives.
  Choose :: a -> a -> Choice a
  -- | No choice.
  Nil :: Choice a

-- | Choose one of two alternatives. Call the 'Choose' operation.
choose :: z :> zz => Handler Choice z -> a -> a -> Eff zz a
choose :: forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> a -> a -> Eff zz a
choose Handler Choice z
h a
x a
y = Handler Choice z -> Choice a -> Eff zz a
forall (s :: Effects) (ss :: Effects) (f :: AEffect) a.
(s :> ss) =>
Handler f s -> f a -> Eff ss a
call Handler Choice z
h (a -> a -> Choice a
forall a. a -> a -> Choice a
Choose a
x a
y)

-- | No choice. Call the 'Nil' operation.
nil :: z :> zz => Handler Choice z -> Eff zz a
nil :: forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> Eff zz a
nil Handler Choice z
h = Handler Choice z -> Choice a -> Eff zz a
forall (s :: Effects) (ss :: Effects) (f :: AEffect) a.
(s :> ss) =>
Handler f s -> f a -> Eff ss a
call Handler Choice z
h Choice a
forall a. Choice a
Nil

-- | Do nothing if @True@. Discard (call 'nil') if @False@.
assume :: z :> zz => Handler Choice z -> Bool -> Eff zz ()
assume :: forall (z :: Effects) (zz :: Effects).
(z :> zz) =>
Handler Choice z -> Bool -> Eff zz ()
assume Handler Choice z
_ Bool
True = () -> Eff zz ()
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure ()
assume Handler Choice z
h Bool
False = Handler Choice z -> Eff zz ()
forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> Eff zz a
nil Handler Choice z
h

-- | Pick an element in a list.
pick :: z :> zz => Handler Choice z -> [a] -> Eff zz a
pick :: forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> [a] -> Eff zz a
pick Handler Choice z
h [] = Handler Choice z -> Eff zz a
forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> Eff zz a
nil Handler Choice z
h
pick Handler Choice z
h (a
x : [a]
xs) = Eff zz (Eff zz a) -> Eff zz a
forall (m :: AEffect) a. Monad m => m (m a) -> m a
join (Eff zz (Eff zz a) -> Eff zz a) -> Eff zz (Eff zz a) -> Eff zz a
forall a b. (a -> b) -> a -> b
$ Handler Choice z -> Eff zz a -> Eff zz a -> Eff zz (Eff zz a)
forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> a -> a -> Eff zz a
choose Handler Choice z
h (a -> Eff zz a
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure a
x) (Handler Choice z -> [a] -> Eff zz a
forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> [a] -> Eff zz a
pick Handler Choice z
h [a]
xs)

-- | Remove an element from a list, returning the resulting list as well.
-- The order of elements is preserved.
removeFrom :: z :> zz => Handler Choice z -> [a] -> Eff zz (a, [a])
removeFrom :: forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> [a] -> Eff zz (a, [a])
removeFrom Handler Choice z
h = [a] -> [a] -> Eff zz (a, [a])
loop []
  where
    loop :: [a] -> [a] -> Eff zz (a, [a])
loop [a]
_ [] = Handler Choice z -> Eff zz (a, [a])
forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> Eff zz a
nil Handler Choice z
h
    loop [a]
ys (a
x : [a]
xs) = Eff zz (Eff zz (a, [a])) -> Eff zz (a, [a])
forall (m :: AEffect) a. Monad m => m (m a) -> m a
join (Eff zz (Eff zz (a, [a])) -> Eff zz (a, [a]))
-> Eff zz (Eff zz (a, [a])) -> Eff zz (a, [a])
forall a b. (a -> b) -> a -> b
$ Handler Choice z
-> Eff zz (a, [a]) -> Eff zz (a, [a]) -> Eff zz (Eff zz (a, [a]))
forall (z :: Effects) (zz :: Effects) a.
(z :> zz) =>
Handler Choice z -> a -> a -> Eff zz a
choose Handler Choice z
h ((a, [a]) -> Eff zz (a, [a])
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure (a
x, [a] -> [a]
forall a. [a] -> [a]
reverse [a]
ys [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
xs)) ([a] -> [a] -> Eff zz (a, [a])
loop (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys) [a]
xs)

-- | Apply a function to every result of the nondeterministic computation.
forAllChoices :: forall a zz.
  (forall z. Handler Choice z -> Eff (z :& zz) a) ->
  (a -> Eff zz ()) ->
  Eff zz ()
forAllChoices :: forall a (zz :: Effects).
(forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a)
-> (a -> Eff zz ()) -> Eff zz ()
forAllChoices forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a
f a -> Eff zz ()
h = (a -> Eff zz ())
-> Eff zz ()
-> (Eff zz () -> Eff zz () -> Eff zz ())
-> (forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a)
-> Eff zz ()
forall a r (zz :: Effects).
(a -> Eff zz r)
-> Eff zz r
-> (Eff zz r -> Eff zz r -> Eff zz r)
-> ScopedEff Choice zz a
-> Eff zz r
foldChoice a -> Eff zz ()
h (() -> Eff zz ()
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure ()) Eff zz () -> Eff zz () -> Eff zz ()
forall a b. Eff zz a -> Eff zz b -> Eff zz b
forall (m :: AEffect) a b. Monad m => m a -> m b -> m b
(>>) Handler Choice s -> Eff (s :& zz) a
forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a
f

-- | Collect the results of a nondeterministic computation in a list.
toList :: forall a zz.
  (forall z. Handler Choice z -> Eff (z :& zz) a) ->
  Eff zz [a]
toList :: forall a (zz :: Effects).
(forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a)
-> Eff zz [a]
toList forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a
f = Eff zz ([a] -> [a]) -> Eff zz [a]
unwrap ((a -> Eff zz ([a] -> [a]))
-> Eff zz ([a] -> [a])
-> (Eff zz ([a] -> [a])
    -> Eff zz ([a] -> [a]) -> Eff zz ([a] -> [a]))
-> (forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a)
-> Eff zz ([a] -> [a])
forall a r (zz :: Effects).
(a -> Eff zz r)
-> Eff zz r
-> (Eff zz r -> Eff zz r -> Eff zz r)
-> ScopedEff Choice zz a
-> Eff zz r
foldChoice (([a] -> [a]) -> Eff zz ([a] -> [a])
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure (([a] -> [a]) -> Eff zz ([a] -> [a]))
-> (a -> [a] -> [a]) -> a -> Eff zz ([a] -> [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (:)) (([a] -> [a]) -> Eff zz ([a] -> [a])
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure [a] -> [a]
forall a. a -> a
id) ((([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a])
-> Eff zz ([a] -> [a])
-> Eff zz ([a] -> [a])
-> Eff zz ([a] -> [a])
forall a b c. (a -> b -> c) -> Eff zz a -> Eff zz b -> Eff zz c
forall (f :: AEffect) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.)) Handler Choice s -> Eff (s :& zz) a
forall (z :: Effects). Handler Choice z -> Eff (z :& zz) a
f)
  where
    unwrap :: Eff zz ([a] -> [a]) -> Eff zz [a]
    unwrap :: Eff zz ([a] -> [a]) -> Eff zz [a]
unwrap = (([a] -> [a]) -> [a]) -> Eff zz ([a] -> [a]) -> Eff zz [a]
forall a b. (a -> b) -> Eff zz a -> Eff zz b
forall (f :: AEffect) a b. Functor f => (a -> b) -> f a -> f b
fmap (([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ [])

-- | Generic 'Choice' handler parameterized by a monoid.
foldChoice :: forall a r zz.
  (a -> Eff zz r) ->           -- ^ Injection
  Eff zz r ->                            -- ^ Identity element
  (Eff zz r -> Eff zz r -> Eff zz r) ->  -- ^ Binary operation
  ScopedEff Choice zz a ->
  Eff zz r
foldChoice :: forall a r (zz :: Effects).
(a -> Eff zz r)
-> Eff zz r
-> (Eff zz r -> Eff zz r -> Eff zz r)
-> ScopedEff Choice zz a
-> Eff zz r
foldChoice a -> Eff zz r
oneE Eff zz r
nilE Eff zz r -> Eff zz r -> Eff zz r
appendE ScopedEff Choice zz a
f = HandlerBody Choice zz r -> ScopedEff Choice zz r -> Eff zz r
forall (f :: AEffect) (ss :: Effects) a.
HandlerBody f ss a -> ScopedEff f ss a -> Eff ss a
handle Choice x -> (x -> Eff zz r) -> Eff zz r
HandlerBody Choice zz r
choiceHandler (Handler Choice s -> Eff (s :& zz) a
ScopedEff Choice zz a
f (Handler Choice s -> Eff (s :& zz) a)
-> (a -> Eff (s :& zz) r) -> Handler Choice s -> Eff (s :& zz) r
forall (m :: AEffect) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Eff zz r -> Eff (s :& zz) r
forall (b :: Effects) r (c1 :: Effects). Eff b r -> Eff (c1 :& b) r
insertFirst (Eff zz r -> Eff (s :& zz) r)
-> (a -> Eff zz r) -> a -> Eff (s :& zz) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Eff zz r
oneE)
  where
    choiceHandler :: HandlerBody Choice zz r
    choiceHandler :: HandlerBody Choice zz r
choiceHandler (Choose x
x x
y) x -> Eff zz r
k = Eff zz r -> Eff zz r -> Eff zz r
appendE (x -> Eff zz r
k x
x) (x -> Eff zz r
k x
y)
    choiceHandler Choice x
Nil x -> Eff zz r
_k = Eff zz r
nilE