{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}

module Data.Seqn.Internal.Stream
  ( Step(..)
  , Stream(..)
  , foldr
  , ifoldr
  , isPrefixOf
  , isSubsequenceOf
  , zipWith
  ) where

import Prelude hiding (foldr, zipWith)
import Data.Functor.Classes (Eq1(..), Ord1(..), eq1, compare1)

import qualified Data.Seqn.Internal.Util as U

-- Budget stream fusion
-- The pieces here are adopted from vector-stream. See vector-stream on Hackage
-- for a more complete implementation.

-- Always benchmark and check the Core when making changes to stream stuff!

data Stream a = forall s. Stream (s -> Step s a) s

data Step s a
  = Yield !a s
  | Done

instance Eq a => Eq (Stream a) where
  == :: Stream a -> Stream a -> Bool
(==) = Stream a -> Stream a -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1
  {-# INLINE (==) #-}

instance Eq1 Stream where
  liftEq :: forall a b. (a -> b -> Bool) -> Stream a -> Stream b -> Bool
liftEq a -> b -> Bool
f (Stream s -> Step s a
step1 s
s10) (Stream s -> Step s b
step2 s
s20) = s -> s -> Bool
go s
s10 s
s20
    where
      go :: s -> s -> Bool
go s
s1 s
s2 = case s -> Step s a
step1 s
s1 of
        Yield a
x1 s
s1' -> case s -> Step s b
step2 s
s2 of
          Yield b
x2 s
s2' -> a -> b -> Bool
f a
x1 b
x2 Bool -> Bool -> Bool
&& s -> s -> Bool
go s
s1' s
s2'
          Step s b
Done -> Bool
False
        Step s a
Done -> case s -> Step s b
step2 s
s2 of
          Yield b
_ s
_ -> Bool
False
          Step s b
Done -> Bool
True
  {-# INLINE liftEq #-}

instance Ord a => Ord (Stream a) where
  compare :: Stream a -> Stream a -> Ordering
compare = Stream a -> Stream a -> Ordering
forall (f :: * -> *) a. (Ord1 f, Ord a) => f a -> f a -> Ordering
compare1
  {-# INLINE compare #-}

instance Ord1 Stream where
  liftCompare :: forall a b.
(a -> b -> Ordering) -> Stream a -> Stream b -> Ordering
liftCompare a -> b -> Ordering
f (Stream s -> Step s a
step1 s
s10) (Stream s -> Step s b
step2 s
s20) = s -> s -> Ordering
go s
s10 s
s20
    where
      go :: s -> s -> Ordering
go s
s1 s
s2 = case s -> Step s a
step1 s
s1 of
        Yield a
x1 s
s1' -> case s -> Step s b
step2 s
s2 of
          Yield b
x2 s
s2' -> a -> b -> Ordering
f a
x1 b
x2 Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> s -> s -> Ordering
go s
s1' s
s2'
          Step s b
Done -> Ordering
GT
        Step s a
Done -> case s -> Step s b
step2 s
s2 of
          Yield b
_ s
_ -> Ordering
LT
          Step s b
Done -> Ordering
EQ
  {-# INLINE liftCompare #-}

foldr :: (a -> b -> b) -> b -> Stream a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Stream a -> b
foldr a -> b -> b
f b
z (Stream s -> Step s a
step s
s0) = s -> b
go s
s0
  where
    go :: s -> b
go s
s = case s -> Step s a
step s
s of
      Yield a
x s
s' -> a -> b -> b
f a
x (s -> b
go s
s')
      Step s a
Done -> b
z
{-# INLINE foldr #-}

ifoldr :: (Int -> a -> b -> b) -> b -> Int -> (Int -> Int) -> Stream a -> b
ifoldr :: forall a b.
(Int -> a -> b -> b) -> b -> Int -> (Int -> Int) -> Stream a -> b
ifoldr Int -> a -> b -> b
f b
z Int
i0 Int -> Int
istep (Stream s -> Step s a
step s
s0) = Int -> s -> b
go Int
i0 s
s0
  where
    go :: Int -> s -> b
go !Int
i s
s = case s -> Step s a
step s
s of
      Yield a
x s
s' -> Int -> a -> b -> b
f Int
i a
x (Int -> s -> b
go (Int -> Int
istep Int
i) s
s')
      Step s a
Done -> b
z
{-# INLINE ifoldr #-}

isPrefixOf :: Eq a => Stream a -> Stream a -> Bool
isPrefixOf :: forall a. Eq a => Stream a -> Stream a -> Bool
isPrefixOf (Stream s -> Step s a
step1 s
s10) (Stream s -> Step s a
step2 s
s20) = s -> s -> Bool
go s
s10 s
s20
  where
    go :: s -> s -> Bool
go s
s1 s
s2 = case s -> Step s a
step1 s
s1 of
      Yield a
x1 s
s1' -> case s -> Step s a
step2 s
s2 of
        Yield a
x2 s
s2' -> a
x1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x2 Bool -> Bool -> Bool
&& s -> s -> Bool
go s
s1' s
s2'
        Step s a
Done -> Bool
False
      Step s a
Done -> Bool
True
{-# INLINE isPrefixOf #-}

isSubsequenceOf :: Eq a => Stream a -> Stream a -> Bool
isSubsequenceOf :: forall a. Eq a => Stream a -> Stream a -> Bool
isSubsequenceOf (Stream s -> Step s a
step1 s
s10) (Stream s -> Step s a
step2 s
s20) = s -> s -> Bool
go1 s
s10 s
s20
  where
    go1 :: s -> s -> Bool
go1 s
s1 s
s2 = case s -> Step s a
step1 s
s1 of
      Yield a
x s
s1' -> a -> s -> s -> Bool
go2 a
x s
s1' s
s2
      Step s a
Done -> Bool
True
    go2 :: a -> s -> s -> Bool
go2 !a
x s
s1' s
s2 = case s -> Step s a
step2 s
s2 of
      Yield a
y s
s2'
        | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y -> s -> s -> Bool
go1 s
s1' s
s2'
        | Bool
otherwise -> a -> s -> s -> Bool
go2 a
x s
s1' s
s2'
      Step s a
Done -> Bool
False
{-# INLINE isSubsequenceOf #-}

zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
zipWith :: forall a b c. (a -> b -> c) -> Stream a -> Stream b -> Stream c
zipWith a -> b -> c
f (Stream s -> Step s a
step1 s
s10) (Stream s -> Step s b
step2 s
s20) = (S2 s s -> Step (S2 s s) c) -> S2 s s -> Stream c
forall a s. (s -> Step s a) -> s -> Stream a
Stream S2 s s -> Step (S2 s s) c
step (s -> s -> S2 s s
forall a b. a -> b -> S2 a b
U.S2 s
s10 s
s20)
  where
    step :: S2 s s -> Step (S2 s s) c
step (U.S2 s
s1 s
s2) = case s -> Step s a
step1 s
s1 of
      Yield a
x1 s
s1' -> case s -> Step s b
step2 s
s2 of
        Yield b
x2 s
s2' -> c -> S2 s s -> Step (S2 s s) c
forall s a. a -> s -> Step s a
Yield (a -> b -> c
f a
x1 b
x2) (s -> s -> S2 s s
forall a b. a -> b -> S2 a b
U.S2 s
s1' s
s2')
        Step s b
Done -> Step (S2 s s) c
forall s a. Step s a
Done
      Step s a
Done -> Step (S2 s s) c
forall s a. Step s a
Done
    {-# INLINE [0] step #-}
{-# INLINE zipWith #-}