{-# 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 F
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 forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
Pulsable pulse =>
pulse m ans -> Bool
done pulse m ans
p
      then forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
Pulsable pulse =>
pulse m ans -> ans
current pulse m ans
p)
      else do pulse m ans
p' <- forall (pulse :: (* -> *) -> * -> *) (m :: * -> *) ans.
(Pulsable pulse, Monad m) =>
pulse m ans -> m (pulse m ans)
pulseM pulse m ans
p; 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 " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n forall a. [a] -> [a] -> [a]
++ forall a. [a] -> String
status [a]
t forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show ans
a forall a. [a] -> [a] -> [a]
++ String
")"

status :: [a] -> String
status :: forall a. [a] -> String
status [a]
x = if 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 " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n forall a. [a] -> [a] -> [a]
++ (if forall k a. Map k a -> Bool
Map.null Map k v
t then String
" Done " else String
" More ") forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show ans
a 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 :: forall t1 t2.
Int -> (t1 -> t2 -> t1) -> [t2] -> t1 -> PulseListM Identity t1
pulseList Int
n t1 -> t2 -> t1
accum [t2]
xs t1
zero =
  forall ans k (m :: * -> *).
Int -> (ans -> k -> m ans) -> [k] -> ans -> PulseListM m ans
PulseList Int
n (\t1
ans t2
x -> 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 :: forall a k v.
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 = 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 -> 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 :: forall (p :: (* -> *) -> * -> *) ans.
Pulsable p =>
p Identity ans -> p Identity ans
pulse p Identity ans
p = forall a. Identity a -> a
runIdentity (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 :: forall (p :: (* -> *) -> * -> *) ans.
Pulsable p =>
p Identity ans -> ans
complete p Identity ans
p = forall a. Identity a -> a
runIdentity (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 :: forall (m :: * -> *) ans. PulseListM m ans -> Bool
done (PulseList Int
_ ans -> a -> m ans
_ [a]
zs ans
_) = forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
zs
  current :: forall (m :: * -> *) ans. PulseListM m ans -> ans
current (PulseList Int
_ ans -> a -> m ans
_ [a]
_ ans
ans) = ans
ans
  pulseM :: forall (m :: * -> *) ans.
Monad m =>
PulseListM m ans -> m (PulseListM m ans)
pulseM (ll :: PulseListM m ans
ll@(PulseList Int
_ ans -> a -> m ans
_ [a]
balance ans
_)) | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
balance = 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') = forall a. Int -> [a] -> ([a], [a])
List.splitAt Int
n [a]
balance
    ans
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
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ans k (m :: * -> *).
Int -> (ans -> k -> m ans) -> [k] -> ans -> PulseListM m ans
PulseList Int
n ans -> a -> m ans
accum [a]
balance' ans
ans')
  completeM :: forall (m :: * -> *) ans. Monad m => PulseListM m ans -> m ans
completeM (PulseList Int
_ ans -> a -> m ans
accum [a]
balance ans
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 :: forall (m :: * -> *) ans. PulseMapM m ans -> Bool
done (PulseMap Int
_ ans -> k -> v -> m ans
_ Map k v
m ans
_) = forall k a. Map k a -> Bool
Map.null Map k v
m
  current :: forall (m :: * -> *) ans. PulseMapM m ans -> ans
current (PulseMap Int
_ ans -> k -> v -> m ans
_ Map k v
_ ans
ans) = ans
ans
  pulseM :: forall (m :: * -> *) ans.
Monad m =>
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
_)) | forall k a. Map k a -> Bool
Map.null Map k v
balance = 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') = forall k a. Int -> Map k a -> (Map k a, Map k a)
Map.splitAt Int
n Map k v
balance
    ans
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
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (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' :: 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
ans t k
acc = case forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList t k
acc of
  [] -> 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; 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' :: forall (m :: * -> *) a k b.
Monad m =>
(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 = 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.