{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ScopedTypeVariables #-}

module DSV.Fold
  ( Fold (Fold), FoldM (FoldM)
  , foldDrop, foldDropM
  , foldProducer, foldProducerM
  , foldVectorM
  ) where

import DSV.Numbers
import DSV.Pipes
import DSV.Prelude

-- foldl
import qualified Control.Foldl as L
import Control.Foldl (Fold (Fold), FoldM (FoldM))

-- pipes
import qualified Pipes.Prelude as P

foldDrop ::
    forall a b .
    Natural -> Fold a b -> Fold a b

foldDrop :: Natural -> Fold a b -> Fold a b
foldDrop Natural
n (Fold x -> a -> x
step x
begin x -> b
done) =
    ((Natural, x) -> a -> (Natural, x))
-> (Natural, x) -> ((Natural, x) -> b) -> Fold a b
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold (Natural, x) -> a -> (Natural, x)
forall a. (Eq a, Num a) => (a, x) -> a -> (a, x)
step' (Natural, x)
begin' (Natural, x) -> b
forall a. (a, x) -> b
done'
  where
    begin' :: (Natural, x)
begin'          = (Natural
n, x
begin)
    step' :: (a, x) -> a -> (a, x)
step' (a
0,  x
s) a
x = (a
0, x -> a -> x
step x
s a
x)
    step' (a
n', x
s) a
_ = (a
n' a -> a -> a
forall a. Num a => a -> a -> a
- a
1, x
s)
    done' :: (a, x) -> b
done' (a
_,  x
s)   = x -> b
done x
s

foldDropM ::
    forall m a b .
    Monad m
    => Natural -> FoldM m a b -> FoldM m a b

foldDropM :: Natural -> FoldM m a b -> FoldM m a b
foldDropM Natural
n (FoldM x -> a -> m x
step m x
begin x -> m b
done) =
    ((Natural, x) -> a -> m (Natural, x))
-> m (Natural, x) -> ((Natural, x) -> m b) -> FoldM m a b
forall (m :: * -> *) a b x.
(x -> a -> m x) -> m x -> (x -> m b) -> FoldM m a b
FoldM (Natural, x) -> a -> m (Natural, x)
forall a. (Eq a, Num a) => (a, x) -> a -> m (a, x)
step' m (Natural, x)
begin' (Natural, x) -> m b
forall a. (a, x) -> m b
done'
  where
    begin' :: m (Natural, x)
begin'          = (x -> (Natural, x)) -> m x -> m (Natural, x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\x
s  -> (Natural
n, x
s))  m x
begin
    step' :: (a, x) -> a -> m (a, x)
step' (a
0,  x
s) a
x = (x -> (a, x)) -> m x -> m (a, x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\x
s' -> (a
0, x
s')) (x -> a -> m x
step x
s a
x)
    step' (a
n', x
s) a
_ = (a, x) -> m (a, x)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
n' a -> a -> a
forall a. Num a => a -> a -> a
- a
1, x
s)
    done' :: (a, x) -> m b
done' (a
_,  x
s)   = x -> m b
done x
s

foldProducer ::
    forall a b m r.
    Monad m
    => Fold a b -> Producer a m r -> m (r, b)

foldProducer :: Fold a b -> Producer a m r -> m (r, b)
foldProducer Fold a b
fld Producer a m r
p =
  do
    (b
x, r
r) <- (forall x.
 (x -> a -> x) -> x -> (x -> b) -> Producer a m r -> m (b, r))
-> Fold a b -> Producer a m r -> m (b, r)
forall a b r.
(forall x. (x -> a -> x) -> x -> (x -> b) -> r) -> Fold a b -> r
L.purely forall x.
(x -> a -> x) -> x -> (x -> b) -> Producer a m r -> m (b, r)
forall (m :: * -> *) x a b r.
Monad m =>
(x -> a -> x) -> x -> (x -> b) -> Producer a m r -> m (b, r)
P.fold' Fold a b
fld Producer a m r
p
    (r, b) -> m (r, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (r
r, b
x)

foldProducerM ::
    forall a b m r .
    Monad m
    => FoldM m a b -> Producer a m r -> m (r, b)

foldProducerM :: FoldM m a b -> Producer a m r -> m (r, b)
foldProducerM FoldM m a b
fld Producer a m r
p =
  do
    (b
x, r
r) <- (forall x.
 (x -> a -> m x) -> m x -> (x -> m b) -> Producer a m r -> m (b, r))
-> FoldM m a b -> Producer a m r -> m (b, r)
forall a (m :: * -> *) b r.
(forall x. (x -> a -> m x) -> m x -> (x -> m b) -> r)
-> FoldM m a b -> r
L.impurely forall x.
(x -> a -> m x) -> m x -> (x -> m b) -> Producer a m r -> m (b, r)
forall (m :: * -> *) x a b r.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Producer a m r -> m (b, r)
P.foldM' FoldM m a b
fld Producer a m r
p
    (r, b) -> m (r, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (r
r, b
x)

foldVectorM ::
    forall v m a .
    (L.PrimMonad m, L.Vector v a)
    => FoldM m a (v a)

foldVectorM :: FoldM m a (v a)
foldVectorM = FoldM m a (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
FoldM m a (v a)
L.vectorM