{-# 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 {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 :: (a -> b) -> Cont ans a -> Cont ans b
fmap a -> b
f (Cont (a -> ans) -> ans
k2) = ((b -> ans) -> ans) -> Cont ans b
forall ans x. ((x -> ans) -> ans) -> Cont ans x
Cont (\b -> ans
k1 -> (a -> ans) -> ans
k2 (b -> ans
k1 (b -> ans) -> (a -> b) -> a -> ans
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f))

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

instance Monad (Cont r) where
  return :: a -> Cont r a
return a
a = ((a -> r) -> r) -> Cont r a
forall ans x. ((x -> ans) -> ans) -> Cont ans x
Cont (((a -> r) -> r) -> Cont r a) -> ((a -> r) -> r) -> Cont r a
forall a b. (a -> b) -> a -> b
$ \a -> r
k -> a -> r
k a
a -- i.e. return a = \k -> k a
  (Cont (a -> r) -> r
c) >>= :: Cont r a -> (a -> Cont r b) -> Cont r b
>>= a -> Cont r b
f = ((b -> r) -> r) -> Cont r b
forall ans x. ((x -> ans) -> ans) -> Cont ans x
Cont (((b -> r) -> r) -> Cont r b) -> ((b -> r) -> r) -> Cont r b
forall a b. (a -> b) -> a -> b
$ \b -> r
k -> (a -> r) -> r
c (\a
a -> Cont r b -> (b -> r) -> r
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 {Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect :: forall ans. ans -> (tuple -> ans -> ans) -> ans}

instance Functor Collect where
  fmap :: (a -> b) -> Collect a -> Collect b
fmap a -> b
f (Collect forall ans. ans -> (a -> ans -> ans) -> ans
g) = (forall ans. ans -> (b -> ans -> ans) -> ans) -> Collect b
forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
x b -> ans -> ans
c -> ans -> (a -> ans -> ans) -> ans
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 :: a -> Collect a
pure a
x = (forall ans. ans -> (a -> ans -> ans) -> ans) -> Collect a
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 <*> :: Collect (a -> b) -> Collect a -> Collect b
<*> Collect a
x = do a -> b
g <- Collect (a -> b)
f; a
y <- Collect a
x; b -> Collect b
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) >>= :: Collect a -> (a -> Collect b) -> Collect b
>>= a -> Collect b
f = (forall ans. ans -> (b -> ans -> ans) -> ans) -> Collect b
forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
x b -> ans -> ans
c -> ans -> (a -> ans -> ans) -> ans
forall ans. ans -> (a -> ans -> ans) -> ans
g ans
x (\a
t ans
a -> Collect b -> ans -> (b -> ans -> ans) -> ans
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 :: (a -> b -> b) -> b -> Collect a -> b
foldr a -> b -> b
f b
z (Collect forall ans. ans -> (a -> ans -> ans) -> ans
g) = b -> (a -> b -> b) -> b
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 :: Collect tuple -> ans -> (tuple -> ans -> ans) -> ans
fixAction = Collect tuple -> ans -> (tuple -> ans -> ans) -> ans
forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect

mapify :: Ord a => Collect (a, b) -> Map.Map a b
mapify :: Collect (a, b) -> Map a b
mapify Collect (a, b)
m = Collect (a, b)
-> Map a b -> ((a, b) -> Map a b -> Map a b) -> Map a b
forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect (a, b)
m Map a b
forall k a. Map k a
Map.empty (\(a
a, b
b) Map a b
ans -> a -> b -> Map a b -> Map a b
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 :: Collect (a, b) -> [(a, b)]
listify Collect (a, b)
m = Collect (a, b)
-> [(a, b)] -> ((a, b) -> [(a, b)] -> [(a, b)]) -> [(a, b)]
forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect (a, b)
m [] (:)

count :: Collect (a, b) -> Int
count :: Collect (a, b) -> Int
count Collect (a, b)
m = Collect (a, b) -> Int -> ((a, b) -> Int -> Int) -> Int
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 Int -> Int -> Int
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 :: t -> Collect t
one t
t = (forall ans. ans -> (t -> ans -> ans) -> ans) -> Collect 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 :: Collect t
none = (forall ans. ans -> (t -> ans -> ans) -> ans) -> Collect t
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 :: t -> Collect t -> Collect t
front t
t (Collect forall ans. ans -> (t -> ans -> ans) -> ans
g) = (forall ans. ans -> (t -> ans -> ans) -> ans) -> Collect t
forall tuple.
(forall ans. ans -> (tuple -> ans -> ans) -> ans) -> Collect tuple
Collect (\ans
a t -> ans -> ans
f -> ans -> (t -> ans -> ans) -> ans
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 :: Collect t -> t -> Collect t
rear (Collect forall ans. ans -> (t -> ans -> ans) -> ans
g) t
t = (forall ans. ans -> (t -> ans -> ans) -> ans) -> Collect 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 -> (t -> ans -> ans) -> ans
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 ans. ans -> (() -> ans -> ans) -> ans) -> Collect ()
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 ans. ans -> (() -> ans -> ans) -> ans) -> Collect ()
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 :: Int -> Collect t -> [t]
takeC Int
n (Collect forall ans. ans -> (t -> ans -> ans) -> ans
f) = ([t], Int) -> [t]
forall a b. (a, b) -> a
fst (([t], Int) -> (t -> ([t], Int) -> ([t], Int)) -> ([t], Int)
forall ans. ans -> (t -> ans -> ans) -> ans
f ([], Int
n) t -> ([t], Int) -> ([t], Int)
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 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs, b
m b -> b -> b
forall a. Num a => a -> a -> a
- b
1)

isempty :: Collect t -> Bool
isempty :: Collect t -> Bool
isempty Collect t
col = Collect t -> Bool -> (t -> Bool -> Bool) -> Bool
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 :: Collect t -> Bool
nonempty Collect t
col = Collect t -> Bool -> (t -> Bool -> Bool) -> Bool
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 :: Collect t -> Maybe t
hasElem Collect t
col = Collect t -> Maybe t -> (t -> Maybe t -> Maybe t) -> Maybe t
forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect t
col Maybe t
forall a. Maybe a
Nothing (\t
t Maybe 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 (Collect t -> [String] -> (t -> [String] -> [String]) -> [String]
forall tuple.
Collect tuple -> forall ans. ans -> (tuple -> ans -> ans) -> ans
runCollect Collect t
c2 [] (\t
t [String]
ans -> t -> String
forall a. Show a => a -> String
show t
t String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
ans))

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

newtype ColPlus tuple = ColPlus
  { 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 :: (a -> b) -> ColPlus a -> ColPlus b
fmap a -> b
f (ColPlus forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g) = (forall ans.
 ans -> (b -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus b
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 -> ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
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 :: a -> ColPlus a
pure a
x = (forall ans.
 ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus a
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 <*> :: ColPlus (a -> b) -> ColPlus a -> ColPlus b
<*> ColPlus a
x = do a -> b
g <- ColPlus (a -> b)
f; a
y <- ColPlus a
x; b -> ColPlus b
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) >>= :: ColPlus a -> (a -> ColPlus b) -> ColPlus b
>>= a -> ColPlus b
f = (forall ans.
 ans -> (b -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus b
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 -> ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g ans
x (\a
t ans
a -> ColPlus b -> ans -> (b -> ans -> ans) -> (ans -> ans -> ans) -> ans
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 :: 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 = a -> (t -> a -> a) -> (a -> a -> a) -> a
forall ans. ans -> (t -> ans -> ans) -> (ans -> ans -> ans) -> ans
g a
a t -> a -> a
f a -> a -> a
forall a. Monoid a => a -> a -> a
mappend

instance AP.Alternative ColPlus where
  empty :: ColPlus a
empty = (forall ans.
 ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus a
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)
  <|> :: 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 ans.
 ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus a
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 (ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
f ans
a a -> ans -> ans
h ans -> ans -> ans
m) (ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
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 :: ColPlus a
mzero = (forall ans.
 ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus a
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 :: 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 ans.
 ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans)
-> ColPlus a
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 (ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
f ans
a a -> ans -> ans
h ans -> ans -> ans
m) (ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
forall ans. ans -> (a -> ans -> ans) -> (ans -> ans -> ans) -> ans
g ans
a a -> ans -> ans
h ans -> ans -> ans
m))