{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

module Data.Pulse
  ( -- * The class that defines operations on pulsers.
    Pulsable (..),

    -- * Two reusable types that have Pulsable instances
    PulseMapM (..),
    PulseListM (..),

    -- * Virtual versions of PulseMapM and PulseListM specialized to be non-monadic.
    PulseMap,
    PulseList,
    pulseList,
    pulseMap,
    pulse,
    complete,

    -- * Monadic folds designed to be used inside pulsers.
    foldlM',
    foldlWithKeyM',
  )
where

import Control.Monad.Identity (Identity (..))
import qualified Data.Foldable as Foldable
import Data.Kind
import qualified Data.List as List
import Data.Map (Map)
import Data.Map.Internal (Map (..))
import qualified Data.Map.Strict as Map

-- ====================================================

-- | let T be a Pulse structure. A Pulse struture
--   is abstracted over a monad: m, and an answer type: t,
--   so the concrete type of a pulse structure is written: (T m a).
--   The Pulsable class supplies operations on the structure
--   that allow its computation to be split into many discrete
--   steps. One does this by running: "pulse p" or "pulseM p",
--   depending upon whether the computation is monadic or not,
--   to run a discrete step.  The scheduling infrastructure needs
--   to know nothing about what is going on inside the pulse structure.
class Pulsable (pulse :: (Type -> Type) -> Type -> Type) where
  done :: pulse m ans -> Bool
  current :: pulse m ans -> ans
  pulseM :: Monad m => pulse m ans -> m (pulse m ans)
  completeM :: Monad m => pulse m ans -> m ans
  completeM pulse m ans
p =
    if pulse m ans -> Bool
forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
Pulsable pulse =>
pulse m ans -> Bool
done pulse m ans
p
      then ans -> m ans
forall (f :: * -> *) a. Applicative f => a -> f a
pure (pulse m ans -> ans
forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
Pulsable pulse =>
pulse m ans -> ans
current pulse m ans
p)
      else do pulse m ans
p' <- pulse m ans -> m (pulse m ans)
forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
(Pulsable pulse, Monad m) =>
pulse m ans -> m (pulse m ans)
pulseM pulse m ans
p; pulse m ans -> m ans
forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
(Pulsable pulse, Monad m) =>
pulse m ans -> m ans
completeM pulse m ans
p'

-- =================================
-- Pulse structure for List in an arbitray monad

-- | A List based pulser
data PulseListM m ans where
  PulseList :: !Int -> !(ans -> a -> m ans) -> ![a] -> !ans -> PulseListM m ans

instance Show ans => Show (PulseListM m ans) where
  show :: PulseListM m ans -> String
show (PulseList Int
n ans -> a -> m ans
_ [a]
t ans
a) = String
"(Pulse " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ [a] -> String
forall a. [a] -> String
status [a]
t String -> ShowS
forall a. [a] -> [a] -> [a]
++ ans -> String
forall a. Show a => a -> String
show ans
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

status :: [a] -> String
status :: [a] -> String
status [a]
x = if [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
x then String
" Done " else String
" More "

-- =================================
-- Pulse structure for Map in an arbitray monad

-- | A Map based pulser.
data PulseMapM m ans where
  PulseMap :: !Int -> !(ans -> k -> v -> m ans) -> !(Map k v) -> !ans -> PulseMapM m ans

instance Show ans => Show (PulseMapM m ans) where
  show :: PulseMapM m ans -> String
show (PulseMap Int
n ans -> k -> v -> m ans
_ Map k v
t ans
a) =
    String
"(Pulse " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ (if Map k v -> Bool
forall k a. Map k a -> Bool
Map.null Map k v
t then String
" Done " else String
" More ") String -> ShowS
forall a. [a] -> [a] -> [a]
++ ans -> String
forall a. Show a => a -> String
show ans
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

-- ===============================================================
-- Pulse structures can be Specialize to the Identity Monad

-- | Type of a List based pulser in the Identity monad.
type PulseList ans = PulseListM Identity ans

-- | Type of a Map based pulser in the Identity monad.
type PulseMap ans = PulseListM Identity ans

-- Use these 'pseudo constructors' to construct Pulse structures in
-- the identity monad. They automatically lift the accumulating function

-- | Create List pulser structure in the Identity monad, a pure accumulator is lifted to a monadic one.
pulseList :: Int -> (t1 -> t2 -> t1) -> [t2] -> t1 -> PulseListM Identity t1
pulseList :: Int -> (t1 -> t2 -> t1) -> [t2] -> t1 -> PulseListM Identity t1
pulseList Int
n t1 -> t2 -> t1
accum [t2]
xs t1
zero =
  Int
-> (t1 -> t2 -> Identity t1)
-> [t2]
-> t1
-> PulseListM Identity t1
forall ans a (m :: * -> *).
Int -> (ans -> a -> m ans) -> [a] -> ans -> PulseListM m ans
PulseList Int
n (\t1
ans t2
x -> t1 -> Identity t1
forall a. a -> Identity a
Identity (t1 -> t2 -> t1
accum t1
ans t2
x)) [t2]
xs t1
zero

-- | Create Map pulser structure in the Identity monad, a pure accumulator is lifted to a monadic one.
pulseMap :: Int -> (a -> k -> v -> a) -> Map k v -> a -> PulseMapM Identity a
pulseMap :: Int -> (a -> k -> v -> a) -> Map k v -> a -> PulseMapM Identity a
pulseMap Int
n a -> k -> v -> a
accum Map k v
ts a
zero = Int
-> (a -> k -> v -> Identity a)
-> Map k v
-> a
-> PulseMapM Identity a
forall ans k v (m :: * -> *).
Int
-> (ans -> k -> v -> m ans) -> Map k v -> ans -> PulseMapM m ans
PulseMap Int
n (\a
ans k
k v
v -> a -> Identity a
forall a. a -> Identity a
Identity (a -> k -> v -> a
accum a
ans k
k v
v)) Map k v
ts a
zero

-- run Pulse structures in the Identity monad.

-- | Pulse a structure in the Identity monad
pulse :: Pulsable p => p Identity ans -> p Identity ans
pulse :: p Identity ans -> p Identity ans
pulse p Identity ans
p = Identity (p Identity ans) -> p Identity ans
forall a. Identity a -> a
runIdentity (p Identity ans -> Identity (p Identity ans)
forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
(Pulsable pulse, Monad m) =>
pulse m ans -> m (pulse m ans)
pulseM p Identity ans
p)

-- | Complete a structure in the Identity monad
complete :: Pulsable p => p Identity ans -> ans
complete :: p Identity ans -> ans
complete p Identity ans
p = Identity ans -> ans
forall a. Identity a -> a
runIdentity (p Identity ans -> Identity ans
forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
(Pulsable pulse, Monad m) =>
pulse m ans -> m ans
completeM p Identity ans
p)

-- =================================================
-- Some instances

instance Pulsable PulseListM where
  done :: PulseListM m ans -> Bool
done (PulseList Int
_ ans -> a -> m ans
_ [a]
zs ans
_) = [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
zs
  current :: PulseListM m ans -> ans
current (PulseList Int
_ ans -> a -> m ans
_ [a]
_ ans
ans) = ans
ans
  pulseM :: PulseListM m ans -> m (PulseListM m ans)
pulseM (ll :: PulseListM m ans
ll@(PulseList Int
_ ans -> a -> m ans
_ [a]
balance ans
_)) | [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
balance = PulseListM m ans -> m (PulseListM m ans)
forall (f :: * -> *) a. Applicative f => a -> f a
pure PulseListM m ans
ll
  pulseM (PulseList Int
n ans -> a -> m ans
accum [a]
balance ans
ans) = do
    let ([a]
steps, [a]
balance') = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
List.splitAt Int
n [a]
balance
    ans
ans' <- (ans -> a -> m ans) -> ans -> [a] -> m ans
forall (t :: * -> *) (m :: * -> *) ans k.
(Foldable t, Monad m) =>
(ans -> k -> m ans) -> ans -> t k -> m ans
foldlM' ans -> a -> m ans
accum ans
ans [a]
steps
    PulseListM m ans -> m (PulseListM m ans)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> (ans -> a -> m ans) -> [a] -> ans -> PulseListM m ans
forall ans a (m :: * -> *).
Int -> (ans -> a -> m ans) -> [a] -> ans -> PulseListM m ans
PulseList Int
n ans -> a -> m ans
accum [a]
balance' ans
ans')
  completeM :: PulseListM m ans -> m ans
completeM (PulseList Int
_ ans -> a -> m ans
accum [a]
balance ans
ans) = (ans -> a -> m ans) -> ans -> [a] -> m ans
forall (t :: * -> *) (m :: * -> *) ans k.
(Foldable t, Monad m) =>
(ans -> k -> m ans) -> ans -> t k -> m ans
foldlM' ans -> a -> m ans
accum ans
ans [a]
balance

instance Pulsable PulseMapM where
  done :: PulseMapM m ans -> Bool
done (PulseMap Int
_ ans -> k -> v -> m ans
_ Map k v
m ans
_) = Map k v -> Bool
forall k a. Map k a -> Bool
Map.null Map k v
m
  current :: PulseMapM m ans -> ans
current (PulseMap Int
_ ans -> k -> v -> m ans
_ Map k v
_ ans
ans) = ans
ans
  pulseM :: PulseMapM m ans -> m (PulseMapM m ans)
pulseM (ll :: PulseMapM m ans
ll@(PulseMap Int
_ ans -> k -> v -> m ans
_ Map k v
balance ans
_)) | Map k v -> Bool
forall k a. Map k a -> Bool
Map.null Map k v
balance = PulseMapM m ans -> m (PulseMapM m ans)
forall (f :: * -> *) a. Applicative f => a -> f a
pure PulseMapM m ans
ll
  pulseM (PulseMap Int
n ans -> k -> v -> m ans
accum Map k v
balance ans
ans) = do
    let (Map k v
steps, Map k v
balance') = Int -> Map k v -> (Map k v, Map k v)
forall k a. Int -> Map k a -> (Map k a, Map k a)
Map.splitAt Int
n Map k v
balance
    ans
ans' <- (ans -> k -> v -> m ans) -> ans -> Map k v -> m ans
forall (m :: * -> *) a k b.
Monad m =>
(a -> k -> b -> m a) -> a -> Map k b -> m a
foldlWithKeyM' ans -> k -> v -> m ans
accum ans
ans Map k v
steps
    PulseMapM m ans -> m (PulseMapM m ans)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
-> (ans -> k -> v -> m ans) -> Map k v -> ans -> PulseMapM m ans
forall ans k v (m :: * -> *).
Int
-> (ans -> k -> v -> m ans) -> Map k v -> ans -> PulseMapM m ans
PulseMap Int
n ans -> k -> v -> m ans
accum Map k v
balance' ans
ans')

-- ================================================================
-- Special monadic folds for use with PulseListM and PulseMapM
-- They are strict, monadic, and their arguments are in the right order.
-- These functions should appear somewhere in Data.List or Data.List or
-- Data.Foldable or Data.Traversable, or Control.Monad, but they don't.

-- | A strict, monadic, version of 'foldl'. It  associates to the left.
foldlM' :: (Foldable t, Monad m) => (ans -> k -> m ans) -> ans -> t k -> m ans
foldlM' :: (ans -> k -> m ans) -> ans -> t k -> m ans
foldlM' ans -> k -> m ans
accum !ans
ans t k
acc = case t k -> [k]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList t k
acc of
  [] -> ans -> m ans
forall (f :: * -> *) a. Applicative f => a -> f a
pure ans
ans
  (k
k : [k]
more) -> do ans
ans1 <- ans -> k -> m ans
accum ans
ans k
k; (ans -> k -> m ans) -> ans -> [k] -> m ans
forall (t :: * -> *) (m :: * -> *) ans k.
(Foldable t, Monad m) =>
(ans -> k -> m ans) -> ans -> t k -> m ans
foldlM' ans -> k -> m ans
accum ans
ans1 [k]
more

-- | /O(n)/. A strict, monadic, version of 'foldlWithKey'. Each application of the
--   operator is evaluated before using the result in the next application. This
--   function is strict in the starting value. Associates to the left.
foldlWithKeyM' :: Monad m => (a -> k -> b -> m a) -> a -> Map k b -> m a
foldlWithKeyM' :: (a -> k -> b -> m a) -> a -> Map k b -> m a
foldlWithKeyM' a -> k -> b -> m a
f a
z = a -> Map k b -> m a
go a
z
  where
    go :: a -> Map k b -> m a
go !a
z' Map k b
Tip = a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
z'
    go a
z' (Bin Int
_ k
kx b
x Map k b
l Map k b
r) =
      do
        !a
ans1 <- (a -> Map k b -> m a
go a
z' Map k b
l)
        !a
ans2 <- (a -> k -> b -> m a
f a
ans1 k
kx b
x)
        a -> Map k b -> m a
go a
ans2 Map k b
r

-- ===================================
-- We could probably generalise this to PulseFoldableM over any
-- foldable structure. We would have to devise a way to break a Foldable
-- structure into small pieces. Lets leave this to another day.