{-# LANGUAGE PatternGuards, BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Stream.Infinite.Skew
-- Copyright   :  (C) 2011 Edward Kmett,
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Anticausal streams implemented as non-empty skew binary random access lists
--
-- The Applicative zips streams, the monad diagonalizes
------------------------------------------------------------------------------


module Data.Stream.Infinite.Skew
    ( Stream
    , (<|)      -- O(1)
    , (!!)
    , tail      -- O(1)
    , uncons    -- O(1)
    , drop      -- O(log n)
    , dropWhile -- O(n)
    , span
    , break
    , split
    , splitW
    , repeat
    , insert    -- O(n)
    , insertBy
    , adjust    -- O(log n)
    , update    -- O(log n)
    , from
    , indexed
    , interleave
    ) where

import Control.Arrow (first)
import Control.Applicative hiding (empty)
import Control.Comonad
import Data.Distributive
import Data.Functor.Alt
import Data.Functor.Extend
import Data.Functor.Rep
import Data.Foldable
import Data.Traversable
import Data.Semigroup hiding (Last)
import Data.Semigroup.Foldable
import Data.Semigroup.Traversable
import Prelude hiding (null, head, tail, drop, dropWhile, length, foldr, last, span, repeat, replicate, (!!), break)
import Data.Boring (Boring (..), Absurd (..))

infixr 5 :<, <|

data Complete a
    = Tip a
    | Bin !Integer a !(Complete a) !(Complete a)
    deriving Int -> Complete a -> ShowS
forall a. Show a => Int -> Complete a -> ShowS
forall a. Show a => [Complete a] -> ShowS
forall a. Show a => Complete a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Complete a] -> ShowS
$cshowList :: forall a. Show a => [Complete a] -> ShowS
show :: Complete a -> String
$cshow :: forall a. Show a => Complete a -> String
showsPrec :: Int -> Complete a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Complete a -> ShowS
Show

instance Functor Complete where
  fmap :: forall a b. (a -> b) -> Complete a -> Complete b
fmap a -> b
f (Tip a
a) = forall a. a -> Complete a
Tip (a -> b
f a
a)
  fmap a -> b
f (Bin Integer
w a
a Complete a
l Complete a
r) = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
w (a -> b
f a
a) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Complete a
l) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Complete a
r)

instance Extend Complete where
  extended :: forall a b. (Complete a -> b) -> Complete a -> Complete b
extended Complete a -> b
f w :: Complete a
w@Tip {} = forall a. a -> Complete a
Tip (Complete a -> b
f Complete a
w)
  extended Complete a -> b
f w :: Complete a
w@(Bin Integer
n a
_ Complete a
l Complete a
r) = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
n (Complete a -> b
f Complete a
w) (forall (w :: * -> *) a b. Extend w => (w a -> b) -> w a -> w b
extended Complete a -> b
f Complete a
l) (forall (w :: * -> *) a b. Extend w => (w a -> b) -> w a -> w b
extended Complete a -> b
f Complete a
r)

instance Comonad Complete where
  extend :: forall a b. (Complete a -> b) -> Complete a -> Complete b
extend Complete a -> b
f w :: Complete a
w@Tip {} = forall a. a -> Complete a
Tip (Complete a -> b
f Complete a
w)
  extend Complete a -> b
f w :: Complete a
w@(Bin Integer
n a
_ Complete a
l Complete a
r) = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
n (Complete a -> b
f Complete a
w) (forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend Complete a -> b
f Complete a
l) (forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend Complete a -> b
f Complete a
r)
  extract :: forall a. Complete a -> a
extract (Tip a
a) = a
a
  extract (Bin Integer
_ a
a Complete a
_ Complete a
_) = a
a

instance Foldable Complete where
  foldMap :: forall m a. Monoid m => (a -> m) -> Complete a -> m
foldMap a -> m
f (Tip a
a) = a -> m
f a
a
  foldMap a -> m
f (Bin Integer
_ a
a Complete a
l Complete a
r) = a -> m
f a
a forall a. Monoid a => a -> a -> a
`mappend` forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Complete a
l forall a. Monoid a => a -> a -> a
`mappend` forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Complete a
r
  foldr :: forall a b. (a -> b -> b) -> b -> Complete a -> b
foldr a -> b -> b
f b
z (Tip a
a) = a -> b -> b
f a
a b
z
  foldr a -> b -> b
f b
z (Bin Integer
_ a
a Complete a
l Complete a
r) = a -> b -> b
f a
a (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
z Complete a
r) Complete a
l)
#if __GLASGOW_HASKELL__ >= 710
  length :: forall a. Complete a -> Int
length Tip{} = Int
1
  length (Bin Integer
n a
_ Complete a
_ Complete a
_) = forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n
  null :: forall a. Complete a -> Bool
null Complete a
_ = Bool
False
#endif

instance Foldable1 Complete where
  foldMap1 :: forall m a. Semigroup m => (a -> m) -> Complete a -> m
foldMap1 a -> m
f (Tip a
a) = a -> m
f a
a
  foldMap1 a -> m
f (Bin Integer
_ a
a Complete a
l Complete a
r) = a -> m
f a
a forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable1 t, Semigroup m) =>
(a -> m) -> t a -> m
foldMap1 a -> m
f Complete a
l forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable1 t, Semigroup m) =>
(a -> m) -> t a -> m
foldMap1 a -> m
f Complete a
r

instance Traversable Complete where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Complete a -> f (Complete b)
traverse a -> f b
f (Tip a
a) = forall a. a -> Complete a
Tip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a
  traverse a -> f b
f (Bin Integer
n a
a Complete a
l Complete a
r) = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Complete a
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Complete a
r

instance Traversable1 Complete where
  traverse1 :: forall (f :: * -> *) a b.
Apply f =>
(a -> f b) -> Complete a -> f (Complete b)
traverse1 a -> f b
f (Tip a
a) = forall a. a -> Complete a
Tip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a
  traverse1 a -> f b
f (Bin Integer
n a
a Complete a
l Complete a
r) = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a forall (f :: * -> *) a b. Apply f => f (a -> b) -> f a -> f b
<.> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable1 t, Apply f) =>
(a -> f b) -> t a -> f (t b)
traverse1 a -> f b
f Complete a
l forall (f :: * -> *) a b. Apply f => f (a -> b) -> f a -> f b
<.> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable1 t, Apply f) =>
(a -> f b) -> t a -> f (t b)
traverse1 a -> f b
f Complete a
r

bin :: a -> Complete a -> Complete a -> Complete a
bin :: forall a. a -> Complete a -> Complete a -> Complete a
bin a
a Complete a
l Complete a
r = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin (Integer
1 forall a. Num a => a -> a -> a
+ forall a. Complete a -> Integer
weight Complete a
l forall a. Num a => a -> a -> a
+ forall a. Complete a -> Integer
weight Complete a
r) a
a Complete a
l Complete a
r
{-# INLINE bin #-}

weight :: Complete a -> Integer
weight :: forall a. Complete a -> Integer
weight Tip{} = Integer
1
weight (Bin Integer
w a
_ Complete a
_ Complete a
_) = Integer
w
{-# INLINE weight #-}

-- A future is a non-empty skew binary random access list of nodes.
-- The last node, however, is allowed to contain fewer values.
data Stream a = !(Complete a) :< Stream a
--  deriving Show

instance Show a => Show (Stream a) where
  showsPrec :: Int -> Stream a -> ShowS
showsPrec Int
d Stream a
as = Bool -> ShowS -> ShowS
showParen (Int
d forall a. Ord a => a -> a -> Bool
>= Int
10) forall a b. (a -> b) -> a -> b
$
    String -> ShowS
showString String
"fromList " forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Stream a
as)

instance Functor Stream where
  fmap :: forall a b. (a -> b) -> Stream a -> Stream b
fmap a -> b
f (Complete a
t :< Stream a
ts) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Complete a
t forall a. Complete a -> Stream a -> Stream a
:< forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Stream a
ts

instance Extend Stream where
  extended :: forall a b. (Stream a -> b) -> Stream a -> Stream b
extended = forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend

instance Comonad Stream where
  extend :: forall a b. (Stream a -> b) -> Stream a -> Stream b
extend Stream a -> b
g0 (Complete a
t :< Stream a
ts) = forall a b.
(Stream a -> b)
-> Complete a -> (Complete a -> Stream a) -> Complete b
go Stream a -> b
g0 Complete a
t (forall a. Complete a -> Stream a -> Stream a
:< Stream a
ts) forall a. Complete a -> Stream a -> Stream a
:< forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend Stream a -> b
g0 Stream a
ts
    where
      go :: (Stream a -> b) -> Complete a -> (Complete a -> Stream a) -> Complete b
      go :: forall a b.
(Stream a -> b)
-> Complete a -> (Complete a -> Stream a) -> Complete b
go Stream a -> b
g w :: Complete a
w@Tip{}         Complete a -> Stream a
f = forall a. a -> Complete a
Tip (Stream a -> b
g (Complete a -> Stream a
f Complete a
w))
      go Stream a -> b
g w :: Complete a
w@(Bin Integer
n a
_ Complete a
l Complete a
r) Complete a -> Stream a
f = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
n (Stream a -> b
g (Complete a -> Stream a
f Complete a
w)) (forall a b.
(Stream a -> b)
-> Complete a -> (Complete a -> Stream a) -> Complete b
go Stream a -> b
g Complete a
l (forall a. Complete a -> Stream a -> Stream a
:< Complete a -> Stream a
f Complete a
r))  (forall a b.
(Stream a -> b)
-> Complete a -> (Complete a -> Stream a) -> Complete b
go Stream a -> b
g Complete a
r Complete a -> Stream a
f)
  extract :: forall a. Stream a -> a
extract (Complete a
a :< Stream a
_) = forall (w :: * -> *) a. Comonad w => w a -> a
extract Complete a
a

instance Apply Stream where
  Stream (a -> b)
fs <.> :: forall a b. Stream (a -> b) -> Stream a -> Stream b
<.> Stream a
as = forall a b. (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex (\Integer
n a -> b
f -> a -> b
f (Stream a
as forall a. Stream a -> Integer -> a
!! Integer
n)) Stream (a -> b)
fs
  Stream a
as <. :: forall a b. Stream a -> Stream b -> Stream a
<.  Stream b
_  = Stream a
as
  Stream a
_   .> :: forall a b. Stream a -> Stream b -> Stream b
.> Stream b
bs = Stream b
bs

instance ComonadApply Stream where
  <@> :: forall a b. Stream (a -> b) -> Stream a -> Stream b
(<@>) = forall (f :: * -> *) a b. Apply f => f (a -> b) -> f a -> f b
(<.>)
  <@ :: forall a b. Stream a -> Stream b -> Stream a
(<@) = forall (f :: * -> *) a b. Apply f => f a -> f b -> f a
(<.)
  @> :: forall a b. Stream a -> Stream b -> Stream b
(@>) = forall (f :: * -> *) a b. Apply f => f a -> f b -> f b
(.>)

instance Applicative Stream where
  pure :: forall a. a -> Stream a
pure = forall a. a -> Stream a
repeat
  <*> :: forall a b. Stream (a -> b) -> Stream a -> Stream b
(<*>) = forall (f :: * -> *) a b. Apply f => f (a -> b) -> f a -> f b
(<.>)
  (<* ) = (<. )
  ( *>) = ( .>)

instance Alt Stream where
  Stream a
as <!> :: forall a. Stream a -> Stream a -> Stream a
<!> Stream a
bs = forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate forall a b. (a -> b) -> a -> b
$ \Rep Stream
i -> case forall a. Integral a => a -> a -> (a, a)
quotRem Rep Stream
i Integer
2 of
    (Integer
q,Integer
0) -> Stream a
as forall a. Stream a -> Integer -> a
!! Integer
q
    (Integer
q,Integer
_) -> Stream a
bs forall a. Stream a -> Integer -> a
!! Integer
q

instance Foldable Stream where
  foldMap :: forall m a. Monoid m => (a -> m) -> Stream a -> m
foldMap a -> m
f (Complete a
t :< Stream a
ts) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Complete a
t forall a. Monoid a => a -> a -> a
`mappend` forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Stream a
ts
  foldr :: forall a b. (a -> b -> b) -> b -> Stream a -> b
foldr a -> b -> b
f b
z (Complete a
t :< Stream a
ts) = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
z Stream a
ts) Complete a
t
#if __GLASGOW_HASKELL__ >= 710
  length :: forall a. Stream a -> Int
length Stream a
_ = forall a. HasCallStack => String -> a
error String
"infinite length"
  null :: forall a. Stream a -> Bool
null Stream a
_ = Bool
False
#endif

instance Foldable1 Stream where
  foldMap1 :: forall m a. Semigroup m => (a -> m) -> Stream a -> m
foldMap1 a -> m
f (Complete a
t :< Stream a
ts) = forall (t :: * -> *) m a.
(Foldable1 t, Semigroup m) =>
(a -> m) -> t a -> m
foldMap1 a -> m
f Complete a
t forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable1 t, Semigroup m) =>
(a -> m) -> t a -> m
foldMap1 a -> m
f Stream a
ts

instance Traversable Stream where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Stream a -> f (Stream b)
traverse a -> f b
f (Complete a
t :< Stream a
ts) = forall a. Complete a -> Stream a -> Stream a
(:<) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Complete a
t forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Stream a
ts

instance Traversable1 Stream where
  traverse1 :: forall (f :: * -> *) a b.
Apply f =>
(a -> f b) -> Stream a -> f (Stream b)
traverse1 a -> f b
f (Complete a
t :< Stream a
ts) = forall a. Complete a -> Stream a -> Stream a
(:<) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable1 t, Apply f) =>
(a -> f b) -> t a -> f (t b)
traverse1 a -> f b
f Complete a
t forall (f :: * -> *) a b. Apply f => f (a -> b) -> f a -> f b
<.> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable1 t, Apply f) =>
(a -> f b) -> t a -> f (t b)
traverse1 a -> f b
f Stream a
ts

instance Distributive Stream where
  distribute :: forall (f :: * -> *) a. Functor f => f (Stream a) -> Stream (f a)
distribute f (Stream a)
w = forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate (\Rep Stream
i -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Stream a -> Integer -> a
!! Rep Stream
i) f (Stream a)
w)

instance Representable Stream where
  type Rep Stream = Integer
  tabulate :: forall a. (Rep Stream -> a) -> Stream a
tabulate Rep Stream -> a
f      = forall a b. (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex (forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rep Stream -> a
f) (forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
  index :: forall a. Stream a -> Rep Stream -> a
index (Complete a
t :< Stream a
ts) Rep Stream
i
    | Rep Stream
i forall a. Ord a => a -> a -> Bool
< Integer
0     = forall a. HasCallStack => String -> a
error String
"index: negative index"
    | Rep Stream
i forall a. Ord a => a -> a -> Bool
< Integer
w     = forall a. Integer -> Complete a -> a
indexComplete Rep Stream
i Complete a
t
    | Bool
otherwise = forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index Stream a
ts (Rep Stream
i forall a. Num a => a -> a -> a
- Integer
w)
    where w :: Integer
w = forall a. Complete a -> Integer
weight Complete a
t

-- | @since 3.3.1
instance Boring a => Boring (Stream a) where
  boring :: Stream a
boring = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Boring a => a
boring

-- | @since 3.3.1
instance Absurd a => Absurd (Stream a) where
  absurd :: forall b. Stream a -> b
absurd = forall a b. Absurd a => a -> b
absurd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (w :: * -> *) a. Comonad w => w a -> a
extract

instance Semigroup (Stream a) where
  <> :: Stream a -> Stream a -> Stream a
(<>) = forall (f :: * -> *) a. Alt f => f a -> f a -> f a
(<!>)

instance Monad Stream where
  return :: forall a. a -> Stream a
return = forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Stream a
as >>= :: forall a b. Stream a -> (a -> Stream b) -> Stream b
>>= a -> Stream b
f = forall a b. (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex (\Integer
i a
a -> a -> Stream b
f a
a forall a. Stream a -> Integer -> a
!! Integer
i) Stream a
as

interleave :: Stream a -> Stream a -> Stream a
interleave :: forall a. Stream a -> Stream a -> Stream a
interleave = forall (f :: * -> *) a. Alt f => f a -> f a -> f a
(<!>)

repeat :: a -> Stream a
repeat :: forall a. a -> Stream a
repeat a
b = forall a. a -> Complete a -> Stream a
go a
b (forall a. a -> Complete a
Tip a
b)
    where
      go :: a -> Complete a -> Stream a
      go :: forall a. a -> Complete a -> Stream a
go a
a Complete a
as | Complete a
ass <- forall a. a -> Complete a -> Complete a -> Complete a
bin a
a Complete a
as Complete a
as = Complete a
as forall a. Complete a -> Stream a -> Stream a
:< forall a. a -> Complete a -> Stream a
go a
a Complete a
ass

mapWithIndex :: (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex :: forall a b. (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex Integer -> a -> b
f0 Stream a
as0 = forall {a} {a}.
(Integer -> a -> a) -> Integer -> Stream a -> Stream a
spine Integer -> a -> b
f0 Integer
0 Stream a
as0
  where
    spine :: (Integer -> a -> a) -> Integer -> Stream a -> Stream a
spine Integer -> a -> a
f Integer
m (Complete a
a :< Stream a
as) = forall {a} {a}.
(Integer -> a -> a) -> Integer -> Complete a -> Complete a
tree Integer -> a -> a
f Integer
m Complete a
a forall a. Complete a -> Stream a -> Stream a
:< (Integer -> a -> a) -> Integer -> Stream a -> Stream a
spine Integer -> a -> a
f (Integer
m forall a. Num a => a -> a -> a
+ forall a. Complete a -> Integer
weight Complete a
a) Stream a
as
    tree :: (Integer -> a -> a) -> Integer -> Complete a -> Complete a
tree Integer -> a -> a
f Integer
m (Tip a
a) = forall a. a -> Complete a
Tip (Integer -> a -> a
f Integer
m a
a)
    tree Integer -> a -> a
f Integer
m (Bin Integer
n a
a Complete a
l Complete a
r) = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
n (Integer -> a -> a
f Integer
m a
a) ((Integer -> a -> a) -> Integer -> Complete a -> Complete a
tree Integer -> a -> a
f (Integer
m forall a. Num a => a -> a -> a
+ Integer
1) Complete a
l) ((Integer -> a -> a) -> Integer -> Complete a -> Complete a
tree Integer -> a -> a
f (Integer
m forall a. Num a => a -> a -> a
+ Integer
1 forall a. Num a => a -> a -> a
+ forall a. Complete a -> Integer
weight Complete a
l) Complete a
r)

indexed :: Stream a -> Stream (Integer, a)
indexed :: forall a. Stream a -> Stream (Integer, a)
indexed = forall a b. (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex (,)

from :: Num a => a -> Stream a
from :: forall a. Num a => a -> Stream a
from a
a = forall a b. (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex (forall a. Num a => a -> a -> a
(+) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a)

-- | /O(1)/ cons
(<|) :: a -> Stream a -> Stream a
a
a <| :: forall a. a -> Stream a -> Stream a
<| (Complete a
l :< Complete a
r :< Stream a
as)
  | forall a. Complete a -> Integer
weight Complete a
l forall a. Eq a => a -> a -> Bool
== forall a. Complete a -> Integer
weight Complete a
r = forall a. a -> Complete a -> Complete a -> Complete a
bin a
a Complete a
l Complete a
r forall a. Complete a -> Stream a -> Stream a
:< Stream a
as
a
a <| Stream a
as = forall a. a -> Complete a
Tip a
a forall a. Complete a -> Stream a -> Stream a
:< Stream a
as
{-# INLINE (<|) #-}

-- | /O(1)/.
tail :: Stream a -> Stream a
tail :: forall a. Stream a -> Stream a
tail (Tip{} :< Stream a
ts) = Stream a
ts
tail (Bin Integer
_ a
_ Complete a
l Complete a
r :< Stream a
ts) = Complete a
l forall a. Complete a -> Stream a -> Stream a
:< Complete a
r forall a. Complete a -> Stream a -> Stream a
:< Stream a
ts
{-# INLINE tail #-}

-- | /O(1)/.
uncons :: Stream a -> (a, Stream a)
uncons :: forall a. Stream a -> (a, Stream a)
uncons (Tip a
a       :< Stream a
as)  = (a
a, Stream a
as)
uncons (Bin Integer
_ a
a Complete a
l Complete a
r :< Stream a
as)  = (a
a, Complete a
l forall a. Complete a -> Stream a -> Stream a
:< Complete a
r forall a. Complete a -> Stream a -> Stream a
:< Stream a
as)
{-# INLINE uncons #-}

indexComplete :: Integer -> Complete a -> a
indexComplete :: forall a. Integer -> Complete a -> a
indexComplete Integer
0 (Tip a
a) = a
a
indexComplete Integer
0 (Bin Integer
_ a
a Complete a
_ Complete a
_) = a
a
indexComplete Integer
i (Bin Integer
w a
_ Complete a
l Complete a
r)
  | Integer
i forall a. Ord a => a -> a -> Bool
<= Integer
w'   = forall a. Integer -> Complete a -> a
indexComplete (Integer
iforall a. Num a => a -> a -> a
-Integer
1) Complete a
l
  | Bool
otherwise = forall a. Integer -> Complete a -> a
indexComplete (Integer
iforall a. Num a => a -> a -> a
-Integer
1forall a. Num a => a -> a -> a
-Integer
w') Complete a
r
  where w' :: Integer
w' = forall a. Integral a => a -> a -> a
div Integer
w Integer
2
indexComplete Integer
_ Complete a
_ = forall a. HasCallStack => String -> a
error String
"indexComplete"

-- | /O(log n)/.
(!!) :: Stream a -> Integer -> a
!! :: forall a. Stream a -> Integer -> a
(!!) = forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index

-- | /O(log n)/.
drop :: Integer -> Stream a -> Stream a
drop :: forall a. Integer -> Stream a -> Stream a
drop Integer
0 Stream a
ts = Stream a
ts
drop Integer
i (Complete a
t :< Stream a
ts) = case forall a. Ord a => a -> a -> Ordering
compare Integer
i Integer
w of
  Ordering
LT -> forall a.
Integer -> Complete a -> (Complete a -> Stream a) -> Stream a
dropComplete Integer
i Complete a
t (forall a. Complete a -> Stream a -> Stream a
:< Stream a
ts)
  Ordering
EQ -> Stream a
ts
  Ordering
GT -> forall a. Integer -> Stream a -> Stream a
drop (Integer
i forall a. Num a => a -> a -> a
- Integer
w) Stream a
ts
  where w :: Integer
w = forall a. Complete a -> Integer
weight Complete a
t

dropComplete :: Integer -> Complete a -> (Complete a -> Stream a) -> Stream a
dropComplete :: forall a.
Integer -> Complete a -> (Complete a -> Stream a) -> Stream a
dropComplete Integer
0 Complete a
t Complete a -> Stream a
f             = Complete a -> Stream a
f Complete a
t
dropComplete Integer
1 (Bin Integer
_ a
_ Complete a
l Complete a
r) Complete a -> Stream a
f = Complete a
l forall a. Complete a -> Stream a -> Stream a
:< Complete a -> Stream a
f Complete a
r
dropComplete Integer
i (Bin Integer
w a
_ Complete a
l Complete a
r) Complete a -> Stream a
f = case forall a. Ord a => a -> a -> Ordering
compare (Integer
i forall a. Num a => a -> a -> a
- Integer
1) Integer
w' of
    Ordering
LT -> forall a.
Integer -> Complete a -> (Complete a -> Stream a) -> Stream a
dropComplete (Integer
iforall a. Num a => a -> a -> a
-Integer
1) Complete a
l (forall a. Complete a -> Stream a -> Stream a
:< Complete a -> Stream a
f Complete a
r)
    Ordering
EQ -> Complete a -> Stream a
f Complete a
r
    Ordering
GT -> forall a.
Integer -> Complete a -> (Complete a -> Stream a) -> Stream a
dropComplete (Integer
iforall a. Num a => a -> a -> a
-Integer
1forall a. Num a => a -> a -> a
-Integer
w') Complete a
r Complete a -> Stream a
f
    where w' :: Integer
w' = forall a. Integral a => a -> a -> a
div Integer
w Integer
2
dropComplete Integer
_ Complete a
_ Complete a -> Stream a
_ = forall a. HasCallStack => String -> a
error String
"dropComplete"

-- | /O(n)/.
dropWhile :: (a -> Bool) -> Stream a -> Stream a
dropWhile :: forall a. (a -> Bool) -> Stream a -> Stream a
dropWhile a -> Bool
p Stream a
as
  | a -> Bool
p (forall (w :: * -> *) a. Comonad w => w a -> a
extract Stream a
as) = forall a. (a -> Bool) -> Stream a -> Stream a
dropWhile a -> Bool
p (forall a. Stream a -> Stream a
tail Stream a
as)
  | Bool
otherwise   = Stream a
as

-- | /O(n)/
span :: (a -> Bool) -> Stream a -> ([a], Stream a)
span :: forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
span a -> Bool
p Stream a
as
  | a
a <- forall (w :: * -> *) a. Comonad w => w a -> a
extract Stream a
as, a -> Bool
p a
a = forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (a
aforall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
span a -> Bool
p (forall a. Stream a -> Stream a
tail Stream a
as)
  | Bool
otherwise = ([], Stream a
as)

-- | /O(n)/
break :: (a -> Bool) -> Stream a -> ([a], Stream a)
break :: forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
break a -> Bool
p = forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
span (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
p)

-- | /(O(n), O(log n))/ split at _some_ edge where function goes from False to True.
-- best used with a monotonic function
split :: (a -> Bool) -> Stream a -> ([a], Stream a)
split :: forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
split a -> Bool
p (Complete a
a :< Stream a
as)
  | a -> Bool
p (forall (w :: * -> *) a. Comonad w => w a -> a
extract Stream a
as) = forall a.
(a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitComplete a -> Bool
p Complete a
a (forall a. Complete a -> Stream a -> Stream a
:< Stream a
as)
  | ([a]
ts, Stream a
fs) <- forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
split a -> Bool
p Stream a
as = (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (:) [a]
ts Complete a
a, Stream a
fs)

-- for use when we know the split occurs within a given tree
splitComplete :: (a -> Bool) -> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitComplete :: forall a.
(a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitComplete a -> Bool
_ t :: Complete a
t@Tip{} Complete a -> Stream a
f = ([], Complete a -> Stream a
f Complete a
t)
splitComplete a -> Bool
p t :: Complete a
t@(Bin Integer
_ a
a Complete a
l Complete a
r) Complete a -> Stream a
f
  | a -> Bool
p a
a                                                   = ([], Complete a -> Stream a
f Complete a
t)
  | a -> Bool
p (forall (w :: * -> *) a. Comonad w => w a -> a
extract Complete a
r), ([a]
ts, Stream a
fs) <- forall a.
(a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitComplete a -> Bool
p Complete a
l (forall a. Complete a -> Stream a -> Stream a
:< Complete a -> Stream a
f Complete a
r) = (a
aforall a. a -> [a] -> [a]
:[a]
ts, Stream a
fs)
  |                ([a]
ts, Stream a
fs) <- forall a.
(a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitComplete a -> Bool
p Complete a
r Complete a -> Stream a
f        = (a
aforall a. a -> [a] -> [a]
:forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (:) [a]
ts Complete a
l, Stream a
fs)

-- | /(O(n), O(log n))/ split at _some_ edge where function goes from False to True.
-- best used with a monotonic function
--
-- > splitW p xs = (map extract &&& fmap (fmap extract)) . split p . duplicate
splitW :: (Stream a -> Bool) -> Stream a -> ([a], Stream a)
splitW :: forall a. (Stream a -> Bool) -> Stream a -> ([a], Stream a)
splitW Stream a -> Bool
p (Complete a
a :< Stream a
as)
  | Stream a -> Bool
p Stream a
as                    = forall a.
(Stream a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitCompleteW Stream a -> Bool
p Complete a
a (forall a. Complete a -> Stream a -> Stream a
:< Stream a
as)
  | ([a]
ts, Stream a
fs) <- forall a. (Stream a -> Bool) -> Stream a -> ([a], Stream a)
splitW Stream a -> Bool
p Stream a
as = (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (:) [a]
ts Complete a
a, Stream a
fs)

-- for use when we know the split occurs within a given tree
splitCompleteW :: (Stream a -> Bool) -> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitCompleteW :: forall a.
(Stream a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitCompleteW Stream a -> Bool
_ t :: Complete a
t@Tip{} Complete a -> Stream a
f = ([], Complete a -> Stream a
f Complete a
t)
splitCompleteW Stream a -> Bool
p t :: Complete a
t@(Bin Integer
_ a
a Complete a
l Complete a
r) Complete a -> Stream a
f
  | Stream a
w <- Complete a -> Stream a
f Complete a
t, Stream a -> Bool
p Stream a
w                                        = ([], Stream a
w)
  | Stream a
w <- Complete a -> Stream a
f Complete a
r, Stream a -> Bool
p Stream a
w, ([a]
ts, Stream a
fs) <- forall a.
(Stream a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitCompleteW Stream a -> Bool
p Complete a
l (forall a. Complete a -> Stream a -> Stream a
:< Stream a
w) = (a
aforall a. a -> [a] -> [a]
:[a]
ts, Stream a
fs)
  |                ([a]
ts, Stream a
fs) <- forall a.
(Stream a -> Bool)
-> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitCompleteW Stream a -> Bool
p Complete a
r Complete a -> Stream a
f      = (a
aforall a. a -> [a] -> [a]
:forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (:) [a]
ts Complete a
l, Stream a
fs)

-- | /O(n)/
insert :: Ord a => a -> Stream a -> Stream a
insert :: forall a. Ord a => a -> Stream a -> Stream a
insert a
a Stream a
as | ([a]
ts, Stream a
as') <- forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
split (a
aforall a. Ord a => a -> a -> Bool
<=) Stream a
as = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. a -> Stream a -> Stream a
(<|) (a
a forall a. a -> Stream a -> Stream a
<| Stream a
as') [a]
ts

-- | /O(n)/. Finds the split in O(log n), but then has to recons
insertBy :: (a -> a -> Ordering) -> a -> Stream a -> Stream a
insertBy :: forall a. (a -> a -> Ordering) -> a -> Stream a -> Stream a
insertBy a -> a -> Ordering
cmp a
a Stream a
as | ([a]
ts, Stream a
as') <- forall a. (a -> Bool) -> Stream a -> ([a], Stream a)
split (\a
b -> a -> a -> Ordering
cmp a
a a
b forall a. Ord a => a -> a -> Bool
<= Ordering
EQ) Stream a
as = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. a -> Stream a -> Stream a
(<|) (a
a forall a. a -> Stream a -> Stream a
<| Stream a
as') [a]
ts

-- | /O(log n)/ Change the value of the nth entry in the future
adjust :: Integer -> (a -> a) -> Stream a -> Stream a
adjust :: forall a. Integer -> (a -> a) -> Stream a -> Stream a
adjust !Integer
n a -> a
f (Complete a
a :< Stream a
as)
  | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
w = forall a. Integer -> (a -> a) -> Complete a -> Complete a
adjustComplete Integer
n a -> a
f Complete a
a forall a. Complete a -> Stream a -> Stream a
:< Stream a
as
  | Bool
otherwise = Complete a
a forall a. Complete a -> Stream a -> Stream a
:< forall a. Integer -> (a -> a) -> Stream a -> Stream a
adjust (Integer
n forall a. Num a => a -> a -> a
- Integer
w) a -> a
f Stream a
as
  where w :: Integer
w = forall a. Complete a -> Integer
weight Complete a
a

adjustComplete :: Integer -> (a -> a) -> Complete a -> Complete a
adjustComplete :: forall a. Integer -> (a -> a) -> Complete a -> Complete a
adjustComplete Integer
0 a -> a
f (Tip a
a) = forall a. a -> Complete a
Tip (a -> a
f a
a)
adjustComplete Integer
_ a -> a
_ t :: Complete a
t@Tip{} = Complete a
t
adjustComplete Integer
n a -> a
f (Bin Integer
m a
a Complete a
l Complete a
r)
  | Integer
n forall a. Eq a => a -> a -> Bool
== Integer
0 = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
m (a -> a
f a
a) Complete a
l Complete a
r
  | Integer
n forall a. Ord a => a -> a -> Bool
<= Integer
w = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
m a
a (forall a. Integer -> (a -> a) -> Complete a -> Complete a
adjustComplete (Integer
n forall a. Num a => a -> a -> a
- Integer
1) a -> a
f Complete a
l) Complete a
r
  | Bool
otherwise = forall a. Integer -> a -> Complete a -> Complete a -> Complete a
Bin Integer
m a
a Complete a
l (forall a. Integer -> (a -> a) -> Complete a -> Complete a
adjustComplete (Integer
n forall a. Num a => a -> a -> a
- Integer
1 forall a. Num a => a -> a -> a
- Integer
w) a -> a
f Complete a
r)
  where w :: Integer
w = forall a. Complete a -> Integer
weight Complete a
l

update :: Integer -> a -> Stream a -> Stream a
update :: forall a. Integer -> a -> Stream a -> Stream a
update Integer
n = forall a. Integer -> (a -> a) -> Stream a -> Stream a
adjust Integer
n forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const