{-# OPTIONS_GHC -Wno-redundant-constraints -Wno-orphans #-}

module DeferredFolds.Defs.UnfoldlM where

import qualified Data.ByteString.Internal as ByteString
import qualified Data.ByteString.Short.Internal as ShortByteString
import DeferredFolds.Prelude hiding (foldM, mapM_)
import qualified DeferredFolds.Prelude as A
import DeferredFolds.Types

deriving instance (Functor m) => Functor (UnfoldlM m)

instance (Monad m) => Applicative (UnfoldlM m) where
  pure :: forall a. a -> UnfoldlM m a
pure a
x =
    forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (\x -> a -> m x
step x
init -> x -> a -> m x
step x
init a
x)
  <*> :: forall a b. UnfoldlM m (a -> b) -> UnfoldlM m a -> UnfoldlM m b
(<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance (Monad m) => Alternative (UnfoldlM m) where
  empty :: forall a. UnfoldlM m a
empty =
    forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (forall a b. a -> b -> a
const forall (m :: * -> *) a. Monad m => a -> m a
return)
  {-# INLINE (<|>) #-}
  <|> :: forall a. UnfoldlM m a -> UnfoldlM m a -> UnfoldlM m a
(<|>) (UnfoldlM forall x. (x -> a -> m x) -> x -> m x
left) (UnfoldlM forall x. (x -> a -> m x) -> x -> m x
right) =
    forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (\x -> a -> m x
step x
init -> forall x. (x -> a -> m x) -> x -> m x
left x -> a -> m x
step x
init forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall x. (x -> a -> m x) -> x -> m x
right x -> a -> m x
step)

instance (Monad m) => Monad (UnfoldlM m) where
  return :: forall a. a -> UnfoldlM m a
return = forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  >>= :: forall a b. UnfoldlM m a -> (a -> UnfoldlM m b) -> UnfoldlM m b
(>>=) (UnfoldlM forall x. (x -> a -> m x) -> x -> m x
left) a -> UnfoldlM m b
rightK =
    forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> b -> m x
step x
init ->
      let newStep :: x -> a -> m x
newStep x
output a
x =
            case a -> UnfoldlM m b
rightK a
x of
              UnfoldlM forall x. (x -> b -> m x) -> x -> m x
right ->
                forall x. (x -> b -> m x) -> x -> m x
right x -> b -> m x
step x
output
       in forall x. (x -> a -> m x) -> x -> m x
left x -> a -> m x
newStep x
init

instance (Monad m) => MonadPlus (UnfoldlM m) where
  mzero :: forall a. UnfoldlM m a
mzero = forall (f :: * -> *) a. Alternative f => f a
empty
  mplus :: forall a. UnfoldlM m a -> UnfoldlM m a -> UnfoldlM m a
mplus = forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>)

instance MonadTrans UnfoldlM where
  lift :: forall (m :: * -> *) a. Monad m => m a -> UnfoldlM m a
lift m a
m = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (\x -> a -> m x
step x
init -> m a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= x -> a -> m x
step x
init)

instance (Monad m) => Semigroup (UnfoldlM m a) where
  <> :: UnfoldlM m a -> UnfoldlM m a -> UnfoldlM m a
(<>) = forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>)

instance (Monad m) => Monoid (UnfoldlM m a) where
  mempty :: UnfoldlM m a
mempty = forall (f :: * -> *) a. Alternative f => f a
empty
  mappend :: UnfoldlM m a -> UnfoldlM m a -> UnfoldlM m a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

instance Foldable (UnfoldlM Identity) where
  {-# INLINE foldMap #-}
  foldMap :: forall m a. Monoid m => (a -> m) -> UnfoldlM Identity a -> m
foldMap a -> m
inputMonoid = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' m -> a -> m
step forall a. Monoid a => a
mempty
    where
      step :: m -> a -> m
step m
monoid a
input = forall a. Monoid a => a -> a -> a
mappend m
monoid (a -> m
inputMonoid a
input)
  foldl :: forall b a. (b -> a -> b) -> b -> UnfoldlM Identity a -> b
foldl = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
  {-# INLINE foldl' #-}
  foldl' :: forall b a. (b -> a -> b) -> b -> UnfoldlM Identity a -> b
foldl' b -> a -> b
step b
init (UnfoldlM forall x. (x -> a -> Identity x) -> x -> Identity x
run) =
    forall a. Identity a -> a
runIdentity (forall x. (x -> a -> Identity x) -> x -> Identity x
run b -> a -> Identity b
identityStep b
init)
    where
      identityStep :: b -> a -> Identity b
identityStep b
state a
input = forall (m :: * -> *) a. Monad m => a -> m a
return (b -> a -> b
step b
state a
input)

instance (Eq a) => Eq (UnfoldlM Identity a) where
  == :: UnfoldlM Identity a -> UnfoldlM Identity a -> Bool
(==) UnfoldlM Identity a
left UnfoldlM Identity a
right = forall l. IsList l => l -> [Item l]
toList UnfoldlM Identity a
left forall a. Eq a => a -> a -> Bool
== forall l. IsList l => l -> [Item l]
toList UnfoldlM Identity a
right

instance (Show a) => Show (UnfoldlM Identity a) where
  show :: UnfoldlM Identity a -> String
show = forall a. Show a => a -> String
show forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall l. IsList l => l -> [Item l]
toList

instance IsList (UnfoldlM Identity a) where
  type Item (UnfoldlM Identity a) = a
  fromList :: [Item (UnfoldlM Identity a)] -> UnfoldlM Identity a
fromList [Item (UnfoldlM Identity a)]
list = forall (m :: * -> *) (foldable :: * -> *) a.
(Monad m, Foldable foldable) =>
foldable a -> UnfoldlM m a
foldable [Item (UnfoldlM Identity a)]
list
  toList :: UnfoldlM Identity a -> [Item (UnfoldlM Identity a)]
toList = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (:) []

-- | Check whether it's empty
{-# INLINE null #-}
null :: (Monad m) => UnfoldlM m input -> m Bool
null :: forall (m :: * -> *) input. Monad m => UnfoldlM m input -> m Bool
null (UnfoldlM forall x. (x -> input -> m x) -> x -> m x
run) = forall x. (x -> input -> m x) -> x -> m x
run (\Bool
_ input
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False) Bool
True

-- | Perform a monadic strict left fold
{-# INLINE foldlM' #-}
foldlM' :: (Monad m) => (output -> input -> m output) -> output -> UnfoldlM m input -> m output
foldlM' :: forall (m :: * -> *) output input.
Monad m =>
(output -> input -> m output)
-> output -> UnfoldlM m input -> m output
foldlM' output -> input -> m output
step output
init (UnfoldlM forall x. (x -> input -> m x) -> x -> m x
run) =
  forall x. (x -> input -> m x) -> x -> m x
run output -> input -> m output
step output
init

-- | A more efficient implementation of mapM_
{-# INLINE mapM_ #-}
mapM_ :: (Monad m) => (input -> m ()) -> UnfoldlM m input -> m ()
mapM_ :: forall (m :: * -> *) input.
Monad m =>
(input -> m ()) -> UnfoldlM m input -> m ()
mapM_ input -> m ()
step = forall (m :: * -> *) output input.
Monad m =>
(output -> input -> m output)
-> output -> UnfoldlM m input -> m output
foldlM' (forall a b. a -> b -> a
const input -> m ()
step) ()

-- | Same as 'mapM_' with arguments flipped
{-# INLINE forM_ #-}
forM_ :: (Monad m) => UnfoldlM m input -> (input -> m ()) -> m ()
forM_ :: forall (m :: * -> *) input.
Monad m =>
UnfoldlM m input -> (input -> m ()) -> m ()
forM_ = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) input.
Monad m =>
(input -> m ()) -> UnfoldlM m input -> m ()
mapM_

-- | Apply a Gonzalez fold
{-# INLINE fold #-}
fold :: Fold input output -> UnfoldlM Identity input -> output
fold :: forall input output.
Fold input output -> UnfoldlM Identity input -> output
fold (Fold x -> input -> x
step x
init x -> output
extract) = x -> output
extract forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' x -> input -> x
step x
init

-- | Apply a monadic Gonzalez fold
{-# INLINE foldM #-}
foldM :: (Monad m) => FoldM m input output -> UnfoldlM m input -> m output
foldM :: forall (m :: * -> *) input output.
Monad m =>
FoldM m input output -> UnfoldlM m input -> m output
foldM (FoldM x -> input -> m x
step m x
init x -> m output
extract) UnfoldlM m input
view =
  do
    x
initialState <- m x
init
    x
finalState <- forall (m :: * -> *) output input.
Monad m =>
(output -> input -> m output)
-> output -> UnfoldlM m input -> m output
foldlM' x -> input -> m x
step x
initialState UnfoldlM m input
view
    x -> m output
extract x
finalState

-- | Lift a fold input mapping function into a mapping of unfolds
{-# INLINE mapFoldMInput #-}
mapFoldMInput :: (Monad m) => (forall x. FoldM m b x -> FoldM m a x) -> UnfoldlM m a -> UnfoldlM m b
mapFoldMInput :: forall (m :: * -> *) b a.
Monad m =>
(forall x. FoldM m b x -> FoldM m a x)
-> UnfoldlM m a -> UnfoldlM m b
mapFoldMInput forall x. FoldM m b x -> FoldM m a x
newFoldM UnfoldlM m a
unfoldM = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> b -> m x
step x
init -> forall (m :: * -> *) input output.
Monad m =>
FoldM m input output -> UnfoldlM m input -> m output
foldM (forall x. FoldM m b x -> FoldM m a x
newFoldM (forall (m :: * -> *) a b x.
(x -> a -> m x) -> m x -> (x -> m b) -> FoldM m a b
FoldM x -> b -> m x
step (forall (m :: * -> *) a. Monad m => a -> m a
return x
init) forall (m :: * -> *) a. Monad m => a -> m a
return)) UnfoldlM m a
unfoldM

-- | Construct from any foldable
{-# INLINE foldable #-}
foldable :: (Monad m, Foldable foldable) => foldable a -> UnfoldlM m a
foldable :: forall (m :: * -> *) (foldable :: * -> *) a.
(Monad m, Foldable foldable) =>
foldable a -> UnfoldlM m a
foldable foldable a
foldable = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (\x -> a -> m x
step x
init -> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
A.foldlM x -> a -> m x
step x
init foldable a
foldable)

-- | Construct from a specification of how to execute a left-fold
{-# INLINE foldlRunner #-}
foldlRunner :: (Monad m) => (forall x. (x -> a -> x) -> x -> x) -> UnfoldlM m a
foldlRunner :: forall (m :: * -> *) a.
Monad m =>
(forall x. (x -> a -> x) -> x -> x) -> UnfoldlM m a
foldlRunner forall x. (x -> a -> x) -> x -> x
run = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (\x -> a -> m x
stepM x
state -> forall x. (x -> a -> x) -> x -> x
run (\m x
stateM a
a -> m x
stateM forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \x
state -> x -> a -> m x
stepM x
state a
a) (forall (m :: * -> *) a. Monad m => a -> m a
return x
state))

-- | Construct from a specification of how to execute a right-fold
{-# INLINE foldrRunner #-}
foldrRunner :: (Monad m) => (forall x. (a -> x -> x) -> x -> x) -> UnfoldlM m a
foldrRunner :: forall (m :: * -> *) a.
Monad m =>
(forall x. (a -> x -> x) -> x -> x) -> UnfoldlM m a
foldrRunner forall x. (a -> x -> x) -> x -> x
run = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (\x -> a -> m x
stepM -> forall x. (a -> x -> x) -> x -> x
run (\a
x x -> m x
k x
z -> x -> a -> m x
stepM x
z a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= x -> m x
k) forall (m :: * -> *) a. Monad m => a -> m a
return)

unfoldr :: (Monad m) => Unfoldr a -> UnfoldlM m a
unfoldr :: forall (m :: * -> *) a. Monad m => Unfoldr a -> UnfoldlM m a
unfoldr (Unfoldr forall x. (a -> x -> x) -> x -> x
unfoldr) = forall (m :: * -> *) a.
Monad m =>
(forall x. (a -> x -> x) -> x -> x) -> UnfoldlM m a
foldrRunner forall x. (a -> x -> x) -> x -> x
unfoldr

-- | Filter the values given a predicate
{-# INLINE filter #-}
filter :: (Monad m) => (a -> m Bool) -> UnfoldlM m a -> UnfoldlM m a
filter :: forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> UnfoldlM m a -> UnfoldlM m a
filter a -> m Bool
test (UnfoldlM forall x. (x -> a -> m x) -> x -> m x
run) = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM (\x -> a -> m x
step -> forall x. (x -> a -> m x) -> x -> m x
run (\x
state a
element -> a -> m Bool
test a
element forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. a -> a -> Bool -> a
bool (forall (m :: * -> *) a. Monad m => a -> m a
return x
state) (x -> a -> m x
step x
state a
element)))

-- | Ints in the specified inclusive range
{-# INLINE intsInRange #-}
intsInRange :: (Monad m) => Int -> Int -> UnfoldlM m Int
intsInRange :: forall (m :: * -> *). Monad m => Int -> Int -> UnfoldlM m Int
intsInRange Int
from Int
to =
  forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> Int -> m x
step x
init ->
    let loop :: x -> Int -> m x
loop !x
state Int
int =
          if Int
int forall a. Ord a => a -> a -> Bool
<= Int
to
            then do
              x
newState <- x -> Int -> m x
step x
state Int
int
              x -> Int -> m x
loop x
newState (forall a. Enum a => a -> a
succ Int
int)
            else forall (m :: * -> *) a. Monad m => a -> m a
return x
state
     in x -> Int -> m x
loop x
init Int
from

-- | TVar contents
{-# INLINE tVarValue #-}
tVarValue :: TVar a -> UnfoldlM STM a
tVarValue :: forall a. TVar a -> UnfoldlM STM a
tVarValue TVar a
var = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> a -> STM x
step x
state -> do
  a
a <- forall a. TVar a -> STM a
readTVar TVar a
var
  x -> a -> STM x
step x
state a
a

-- | Change the base monad using invariant natural transformations
{-# INLINE hoist #-}
hoist :: (forall a. m a -> n a) -> (forall a. n a -> m a) -> UnfoldlM m a -> UnfoldlM n a
hoist :: forall (m :: * -> *) (n :: * -> *) a.
(forall a. m a -> n a)
-> (forall a. n a -> m a) -> UnfoldlM m a -> UnfoldlM n a
hoist forall a. m a -> n a
trans1 forall a. n a -> m a
trans2 (UnfoldlM forall x. (x -> a -> m x) -> x -> m x
unfold) = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> a -> n x
step x
init ->
  forall a. m a -> n a
trans1 (forall x. (x -> a -> m x) -> x -> m x
unfold (\x
a a
b -> forall a. n a -> m a
trans2 (x -> a -> n x
step x
a a
b)) x
init)

-- | Bytes of a bytestring
{-# INLINEABLE byteStringBytes #-}
byteStringBytes :: ByteString -> UnfoldlM IO Word8
byteStringBytes :: ByteString -> UnfoldlM IO Word8
byteStringBytes (ByteString.PS ForeignPtr Word8
fp Int
off Int
len) =
  forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> Word8 -> IO x
step x
init ->
    forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr ->
      let endPtr :: Ptr Word8
endPtr = forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr (Int
off forall a. Num a => a -> a -> a
+ Int
len)
          iterate :: x -> Ptr Word8 -> IO x
iterate !x
state !Ptr Word8
ptr =
            if Ptr Word8
ptr forall a. Eq a => a -> a -> Bool
== Ptr Word8
endPtr
              then forall (m :: * -> *) a. Monad m => a -> m a
return x
state
              else do
                Word8
x <- forall a. Storable a => Ptr a -> IO a
peek Ptr Word8
ptr
                x
newState <- x -> Word8 -> IO x
step x
state Word8
x
                x -> Ptr Word8 -> IO x
iterate x
newState (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr Int
1)
       in x -> Ptr Word8 -> IO x
iterate x
init (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr Int
off)

-- | Bytes of a short bytestring
{-# INLINE shortByteStringBytes #-}
shortByteStringBytes :: (Monad m) => ShortByteString -> UnfoldlM m Word8
shortByteStringBytes :: forall (m :: * -> *).
Monad m =>
ShortByteString -> UnfoldlM m Word8
shortByteStringBytes (ShortByteString.SBS ByteArray#
ba#) = forall (m :: * -> *) prim.
(Monad m, Prim prim) =>
PrimArray prim -> UnfoldlM m prim
primArray (forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
ba#)

-- | Elements of a prim array
{-# INLINE primArray #-}
primArray :: (Monad m, Prim prim) => PrimArray prim -> UnfoldlM m prim
primArray :: forall (m :: * -> *) prim.
(Monad m, Prim prim) =>
PrimArray prim -> UnfoldlM m prim
primArray PrimArray prim
pa = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> prim -> m x
f x
z -> forall a (m :: * -> *) b.
(Prim a, Monad m) =>
(b -> a -> m b) -> b -> PrimArray a -> m b
foldlPrimArrayM' x -> prim -> m x
f x
z PrimArray prim
pa

-- | Elements of a prim array coming paired with indices
{-# INLINE primArrayWithIndices #-}
primArrayWithIndices :: (Monad m, Prim prim) => PrimArray prim -> UnfoldlM m (Int, prim)
primArrayWithIndices :: forall (m :: * -> *) prim.
(Monad m, Prim prim) =>
PrimArray prim -> UnfoldlM m (Int, prim)
primArrayWithIndices PrimArray prim
pa = forall (m :: * -> *) a.
(forall x. (x -> a -> m x) -> x -> m x) -> UnfoldlM m a
UnfoldlM forall a b. (a -> b) -> a -> b
$ \x -> (Int, prim) -> m x
step x
state ->
  let !size :: Int
size = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray prim
pa
      iterate :: Int -> x -> m x
iterate Int
index !x
state =
        if Int
index forall a. Ord a => a -> a -> Bool
< Int
size
          then do
            x
newState <- x -> (Int, prim) -> m x
step x
state (Int
index, forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray prim
pa Int
index)
            Int -> x -> m x
iterate (forall a. Enum a => a -> a
succ Int
index) x
newState
          else forall (m :: * -> *) a. Monad m => a -> m a
return x
state
   in Int -> x -> m x
iterate Int
0 x
state