{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# OPTIONS_GHC -Wno-unused-imports #-}
{-# OPTIONS_GHC -Wno-unused-matches #-}

module Control.Iterate.Collect where

import qualified Control.Applicative as AP
import qualified Control.Monad as CM
import qualified Data.Map.Strict as Map

-- =========================================================================
-- Sample continuation monad to study. We don't actually use this monad, but
-- we put it here since it is the simplest continuation monad, and studying
-- it, helped me define the Collect monad.

newtype Cont ans x = Cont {forall ans x. Cont ans x -> (x -> ans) -> ans
runCont :: (x -> ans) -> ans} -- ans is the final result type of the whole computation

instance Functor (Cont ans) where
  fmap :: forall a b. (a -> b) -> Cont ans a -> Cont ans b
fmap a -> b
f (Cont (a -> ans) -> ans
k2) = forall ans x. ((x -> ans) -> ans) -> Cont ans x
Cont (\b -> ans
k1 -> (a -> ans) -> ans
k2 (b -> ans
k1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f))

instance Applicative (Cont ans) where
  pure :: forall a. a -> Cont ans a
pure a
x = forall ans x. ((x -> ans) -> ans) -> Cont ans x
Cont (\a -> ans
ret -> a -> ans
ret a
x)
  Cont ans (a -> b)
f <*> :: forall a b. Cont ans (a -> b) -> Cont ans a -> Cont ans b
<*> Cont ans a
x = do a -> b
g <- Cont ans (a -> b)
f; a
y <- Cont ans a
x; forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> b
g a
y)

instance Monad (Cont r) where
  Cont (a -> r) -> r
c >>= :: forall a b. Cont r a -> (a -> Cont r b) -> Cont r b
>>= a -> Cont r b
f = forall ans x. ((x -> ans) -> ans) -> Cont ans x
Cont forall a b. (a -> b) -> a -> b
$ \b -> r
k -> (a -> r) -> r
c (\a
a -> forall ans x. Cont ans x -> (x -> ans) -> ans
runCont (a -> Cont r b
f a
a) b -> r
k) -- i.e. c >>= f = \k -> c (\a -> f a k)

-- ========================================================================
-- Now we want to make the following, more complicated continuation a Monad
-- Here the answer type is completely abstract.

newtype Collect tuple = Collect {forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect :: forall ans. ans -> (tuple -> ans -> ans) -> ans}

instance Functor Collect where
  fmap :: forall a b. (a -> b) -> Collect a -> Collect b
fmap a -> b
f (Collect forall ans. ans -> (a -> ans -> ans) -> ans
g) = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
x b -> ans -> ans
c -> forall ans. ans -> (a -> ans -> ans) -> ans
g ans
x (\a
t ans
a -> b -> ans -> ans
c (a -> b
f a
t) ans
a))

-- Playing type tetris find this term    ^----------------------^
-- given
-- f:: t -> s
-- g:: a -> (t -> a -> a) -> a
-- x:: a
-- c:: s -> a -> a

instance Applicative Collect where
  pure :: forall a. a -> Collect a
pure a
x = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
ans a -> ans -> ans
f -> a -> ans -> ans
f a
x ans
ans)
  Collect (a -> b)
f <*> :: forall a b. Collect (a -> b) -> Collect a -> Collect b
<*> Collect a
x = do a -> b
g <- Collect (a -> b)
f; a
y <- Collect a
x; forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> b
g a
y)

instance Monad Collect where
  (Collect forall ans. ans -> (a -> ans -> ans) -> ans
g) >>= :: forall a b. Collect a -> (a -> Collect b) -> Collect b
>>= a -> Collect b
f = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
x b -> ans -> ans
c -> forall ans. ans -> (a -> ans -> ans) -> ans
g ans
x (\a
t ans
a -> forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect (a -> Collect b
f a
t) ans
a b -> ans -> ans
c))

-- Playing type tetris find this term  ^--------------------------------^
-- given
-- g:: a -> (t -> a -> a) -> a
-- f:: t -> (Collect s)
-- x:: a
-- c:: (s -> a -> a)

instance Foldable Collect where
  foldr :: forall a b. (a -> b -> b) -> b -> Collect a -> b
foldr a -> b -> b
f b
z (Collect forall ans. ans -> (a -> ans -> ans) -> ans
g) = forall ans. ans -> (a -> ans -> ans) -> ans
g b
z a -> b -> b
f

-- ===========================================================================
-- Operations on the collect Monad.

-- | A (Collect t) is completely agnostic over how 't's are beging collected.
-- We can make this abstraction concrete by using fixAction.
fixAction :: Collect tuple -> ans -> (tuple -> ans -> ans) -> ans
fixAction :: forall tuple ans.
Collect tuple -> ans -> (tuple -> ans -> ans) -> ans
fixAction = forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect

mapify :: Ord a => Collect (a, b) -> Map.Map a b
mapify :: forall a b. Ord a => Collect (a, b) -> Map a b
mapify Collect (a, b)
m = forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect (a, b)
m forall k a. Map k a
Map.empty (\(a
a, b
b) Map a b
ans -> forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
a b
b Map a b
ans)

listify :: Collect (a, b) -> [(a, b)]
listify :: forall a b. Collect (a, b) -> [(a, b)]
listify Collect (a, b)
m = forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect (a, b)
m [] (:)

count :: Collect (a, b) -> Int
count :: forall a b. Collect (a, b) -> Int
count Collect (a, b)
m = forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect (a, b)
m Int
0 (\(a, b)
t Int
n -> Int
n forall a. Num a => a -> a -> a
+ Int
1)

-- | Here are several ways to add a new t to what is being collected.

-- | The `one` and `none` interface are used when we want collections with 0 or 1 elements
one :: t -> Collect t
one :: forall a. a -> Collect a
one t
t = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
a t -> ans -> ans
f -> t -> ans -> ans
f t
t ans
a)

none :: Collect t
none :: forall t. Collect t
none = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
a t -> ans -> ans
f -> ans
a)

-- | The `front` and `rear` interface can add to either end of the sequence (both in constant time)
front :: t -> Collect t -> Collect t
front :: forall t. t -> Collect t -> Collect t
front t
t (Collect forall ans. ans -> (t -> ans -> ans) -> ans
g) = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
a t -> ans -> ans
f -> forall ans. ans -> (t -> ans -> ans) -> ans
g (t -> ans -> ans
f t
t ans
a) t -> ans -> ans
f)

rear :: Collect t -> t -> Collect t
rear :: forall t. Collect t -> t -> Collect t
rear (Collect forall ans. ans -> (t -> ans -> ans) -> ans
g) t
t = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
a t -> ans -> ans
f -> t -> ans -> ans
f t
t (forall ans. ans -> (t -> ans -> ans) -> ans
g ans
a t -> ans -> ans
f))

-- | Conditional collecting
when :: Bool -> Collect ()
when :: Bool -> Collect ()
when Bool
True = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
ans () -> ans -> ans
f -> () -> ans -> ans
f () ans
ans)
when Bool
False = forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
ans () -> ans -> ans
f -> ans
ans)

takeC :: Int -> Collect t -> [t]
takeC :: forall t. Int -> Collect t -> [t]
takeC Int
n (Collect forall ans. ans -> (t -> ans -> ans) -> ans
f) = forall a b. (a, b) -> a
fst (forall ans. ans -> (t -> ans -> ans) -> ans
f ([], Int
n) forall {b} {a}. (Eq b, Num b) => a -> ([a], b) -> ([a], b)
next)
  where
    next :: a -> ([a], b) -> ([a], b)
next a
x ([a]
xs, b
0) = ([a]
xs, b
0)
    next a
x ([a]
xs, b
m) = (a
x forall a. a -> [a] -> [a]
: [a]
xs, b
m forall a. Num a => a -> a -> a
- b
1)

isempty :: Collect t -> Bool
isempty :: forall a. Collect a -> Bool
isempty Collect t
col = forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect t
col Bool
True (\t
t Bool
a -> Bool
False)

nonempty :: Collect t -> Bool
nonempty :: forall a. Collect a -> Bool
nonempty Collect t
col = forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect t
col Bool
False (\t
t Bool
a -> Bool
True)

hasElem :: Collect t -> Maybe t
hasElem :: forall t. Collect t -> Maybe t
hasElem Collect t
col = forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect t
col forall a. Maybe a
Nothing (\t
t Maybe t
_ -> forall a. a -> Maybe a
Just t
t)

-- | Even though a (Collect t) is a function, if we can (Show t), we can pick an action
-- that collects all the shown t, and turn them into a big multi-line string.
instance Show t => Show (Collect t) where
  show :: Collect t -> String
show Collect t
c2 = [String] -> String
unlines (forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect t
c2 [] (\t
t [String]
ans -> forall a. Show a => a -> String
show t
t forall a. a -> [a] -> [a]
: [String]
ans))

-- =======================================================
-- Collection with mplus

newtype ColPlus tuple = ColPlus
  { forall tuple.
ColPlus tuple
-> forall ans.
   ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans
runColPlus :: forall ans. ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans
  }

instance Functor ColPlus where
  fmap :: forall a b. (a -> b) -> ColPlus a -> ColPlus b
fmap a -> b
f (ColPlus forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g) = forall tuple.
(forall ans.
 ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus tuple
ColPlus (\ans
x b -> ans -> ans
c ans -> ans -> ans
m -> forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g ans
x (\a
t ans
a -> b -> ans -> ans
c (a -> b
f a
t) ans
a) ans -> ans -> ans
m)

instance Applicative ColPlus where
  pure :: forall a. a -> ColPlus a
pure a
x = forall tuple.
(forall ans.
 ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus tuple
ColPlus (\ans
ans a -> ans -> ans
f ans -> ans -> ans
m -> a -> ans -> ans
f a
x ans
ans)
  ColPlus (a -> b)
f <*> :: forall a b. ColPlus (a -> b) -> ColPlus a -> ColPlus b
<*> ColPlus a
x = do a -> b
g <- ColPlus (a -> b)
f; a
y <- ColPlus a
x; forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> b
g a
y)

instance Monad ColPlus where
  (ColPlus forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g) >>= :: forall a b. ColPlus a -> (a -> ColPlus b) -> ColPlus b
>>= a -> ColPlus b
f = forall tuple.
(forall ans.
 ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus tuple
ColPlus (\ans
x b -> ans -> ans
c ans -> ans -> ans
m -> forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g ans
x (\a
t ans
a -> forall tuple.
ColPlus tuple
-> forall ans.
   ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans
runColPlus (a -> ColPlus b
f a
t) ans
a b -> ans -> ans
c ans -> ans -> ans
m) ans -> ans -> ans
m)

runPlus :: Monoid a => ColPlus t -> a -> (t -> a -> a) -> a
runPlus :: forall a t. Monoid a => ColPlus t -> a -> (t -> a -> a) -> a
runPlus (ColPlus forall ans. ans -> (t -> ans -> ans) -> (ans -> ans -> ans) -> ans
g) a
a t -> a -> a
f = forall ans. ans -> (t -> ans -> ans) -> (ans -> ans -> ans) -> ans
g a
a t -> a -> a
f forall a. Monoid a => a -> a -> a
mappend

instance AP.Alternative ColPlus where
  empty :: forall a. ColPlus a
empty = forall tuple.
(forall ans.
 ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus tuple
ColPlus (\ans
a a -> ans -> ans
h ans -> ans -> ans
m -> ans
a)
  <|> :: forall a. ColPlus a -> ColPlus a -> ColPlus a
(<|>) (ColPlus forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
f) (ColPlus forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g) = forall tuple.
(forall ans.
 ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus tuple
ColPlus (\ans
a a -> ans -> ans
h ans -> ans -> ans
m -> ans -> ans -> ans
m (forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
f ans
a a -> ans -> ans
h ans -> ans -> ans
m) (forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g ans
a a -> ans -> ans
h ans -> ans -> ans
m))

instance CM.MonadPlus ColPlus where
  mzero :: forall a. ColPlus a
mzero = forall tuple.
(forall ans.
 ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus tuple
ColPlus (\ans
a a -> ans -> ans
h ans -> ans -> ans
m -> ans
a)
  mplus :: forall a. ColPlus a -> ColPlus a -> ColPlus a
mplus (ColPlus forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
f) (ColPlus forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g) = forall tuple.
(forall ans.
 ans -> (tuple -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus tuple
ColPlus (\ans
a a -> ans -> ans
h ans -> ans -> ans
m -> ans -> ans -> ans
m (forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
f ans
a a -> ans -> ans
h ans -> ans -> ans
m) (forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g ans
a a -> ans -> ans
h ans -> ans -> ans
m))