module Data.Stream.Infinite.Skew
( Stream
, (<|)
, (!!)
, head
, tail
, tails
, uncons
, index
, drop
, dropWhile
, span
, break
, split
, splitW
, repeat
, insert
, insertBy
, adjust
, update
, fromList
, from
, indexed
, interleave
, tabulate
) 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.Foldable hiding (toList)
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)
infixr 5 :<, <|
data Complete a
= Tip a
| Bin !Integer a !(Complete a) !(Complete a)
deriving Show
instance Functor Complete where
fmap f (Tip a) = Tip (f a)
fmap f (Bin w a l r) = Bin w (f a) (fmap f l) (fmap f r)
instance Extend Complete where
extended f w@Tip {} = Tip (f w)
extended f w@(Bin n _ l r) = Bin n (f w) (extended f l) (extended f r)
instance Comonad Complete where
extend f w@Tip {} = Tip (f w)
extend f w@(Bin n _ l r) = Bin n (f w) (extend f l) (extend f r)
extract (Tip a) = a
extract (Bin _ a _ _) = a
instance Foldable Complete where
foldMap f (Tip a) = f a
foldMap f (Bin _ a l r) = f a `mappend` foldMap f l `mappend` foldMap f r
foldr f z (Tip a) = f a z
foldr f z (Bin _ a l r) = f a (foldr f (foldr f z r) l)
instance Foldable1 Complete where
foldMap1 f (Tip a) = f a
foldMap1 f (Bin _ a l r) = f a <> foldMap1 f l <> foldMap1 f r
instance Traversable Complete where
traverse f (Tip a) = Tip <$> f a
traverse f (Bin n a l r) = Bin n <$> f a <*> traverse f l <*> traverse f r
instance Traversable1 Complete where
traverse1 f (Tip a) = Tip <$> f a
traverse1 f (Bin n a l r) = Bin n <$> f a <.> traverse1 f l <.> traverse1 f r
bin :: a -> Complete a -> Complete a -> Complete a
bin a l r = Bin (1 + weight l + weight r) a l r
weight :: Complete a -> Integer
weight Tip{} = 1
weight (Bin w _ _ _) = w
data Stream a = !(Complete a) :< Stream a
instance Show a => Show (Stream a) where
showsPrec d as = showParen (d >= 10) $
showString "fromList " . showsPrec 11 (toList as)
instance Functor Stream where
fmap f (t :< ts) = fmap f t :< fmap f ts
instance Extend Stream where
extended = extend
instance Comonad Stream where
extend g0 (t :< ts) = go g0 t (:< ts) :< extend g0 ts
where
go :: (Stream a -> b) -> Complete a -> (Complete a -> Stream a) -> Complete b
go g w@Tip{} f = Tip (g (f w))
go g w@(Bin n _ l r) f = Bin n (g (f w)) (go g l (:< f r)) (go g r f)
extract = head
instance Apply Stream where
fs <.> as = mapWithIndex (\n f -> f (as !! n)) fs
as <. _ = as
_ .> bs = bs
instance ComonadApply Stream where
(<@>) = (<.>)
(<@) = (<.)
(@>) = (.>)
instance Applicative Stream where
pure = repeat
(<*>) = (<.>)
(<* ) = (<. )
( *>) = ( .>)
instance Alt Stream where
as <!> bs = tabulate $ \i -> case quotRem i 2 of
(q,0) -> as !! q
(q,_) -> bs !! q
instance Foldable Stream where
foldMap f (t :< ts) = foldMap f t `mappend` foldMap f ts
foldr f z (t :< ts) = foldr f (foldr f z ts) t
toList :: Stream a -> [a]
toList = foldr (:) []
instance Foldable1 Stream where
foldMap1 f (t :< ts) = foldMap1 f t <> foldMap1 f ts
instance Traversable Stream where
traverse f (t :< ts) = (:<) <$> traverse f t <*> traverse f ts
instance Traversable1 Stream where
traverse1 f (t :< ts) = (:<) <$> traverse1 f t <.> traverse1 f ts
instance Distributive Stream where
distribute w = tabulate (\i -> fmap (!! i) w)
instance Semigroup (Stream a) where
(<>) = (<!>)
instance Monad Stream where
return = pure
as >>= f = mapWithIndex (\i a -> f a !! i) as
interleave :: Stream a -> Stream a -> Stream a
interleave = (<!>)
repeat :: a -> Stream a
repeat b = go b (Tip b)
where
go :: a -> Complete a -> Stream a
go a as | ass <- bin a as as = as :< go a ass
mapWithIndex :: (Integer -> a -> b) -> Stream a -> Stream b
mapWithIndex f0 as0 = spine f0 0 as0
where
spine f m (a :< as) = tree f m a :< spine f (m + weight a) as
tree f m (Tip a) = Tip (f m a)
tree f m (Bin n a l r) = Bin n (f m a) (tree f (m + 1) l) (tree f (m + 1 + weight l) r)
tabulate :: (Integer -> a) -> Stream a
tabulate f = mapWithIndex (const . f) (pure ())
indexed :: Stream a -> Stream (Integer, a)
indexed = mapWithIndex (,)
from :: Num a => a -> Stream a
from a = mapWithIndex ((+) . fromIntegral) (pure a)
(<|) :: a -> Stream a -> Stream a
a <| (l :< r :< as)
| weight l == weight r = bin a l r :< as
a <| as = Tip a :< as
head :: Stream a -> a
head (a :< _) = extract a
tail :: Stream a -> Stream a
tail (Tip{} :< ts) = ts
tail (Bin _ _ l r :< ts) = l :< r :< ts
tails :: Stream a -> Stream (Stream a)
tails = duplicate
uncons :: Stream a -> (a, Stream a)
uncons (Tip a :< as) = (a, as)
uncons (Bin _ a l r :< as) = (a, l :< r :< as)
index :: Integer -> Stream a -> a
index i (t :< ts)
| i < 0 = error "index: negative index"
| i < w = indexComplete i t
| otherwise = index (i w) ts
where w = weight t
indexComplete :: Integer -> Complete a -> a
indexComplete 0 (Tip a) = a
indexComplete 0 (Bin _ a _ _) = a
indexComplete i (Bin w _ l r)
| i <= w' = indexComplete (i1) l
| otherwise = indexComplete (i1w') r
where w' = div w 2
indexComplete _ _ = error "indexComplete"
(!!) :: Stream a -> Integer -> a
(!!) = flip index
drop :: Integer -> Stream a -> Stream a
drop 0 ts = ts
drop i (t :< ts) = case compare i w of
LT -> dropComplete i t (:< ts)
EQ -> ts
GT -> drop (i w) ts
where w = weight t
dropComplete :: Integer -> Complete a -> (Complete a -> Stream a) -> Stream a
dropComplete 0 t f = f t
dropComplete 1 (Bin _ _ l r) f = l :< f r
dropComplete i (Bin w _ l r) f = case compare (i 1) w' of
LT -> dropComplete (i1) l (:< f r)
EQ -> f r
GT -> dropComplete (i1w') r f
where w' = div w 2
dropComplete _ _ _ = error "dropComplete"
dropWhile :: (a -> Bool) -> Stream a -> Stream a
dropWhile p as
| p (head as) = dropWhile p (tail as)
| otherwise = as
span :: (a -> Bool) -> Stream a -> ([a], Stream a)
span p as
| a <- head as, p a = first (a:) $ span p (tail as)
| otherwise = ([], as)
break :: (a -> Bool) -> Stream a -> ([a], Stream a)
break p = span (not . p)
split :: (a -> Bool) -> Stream a -> ([a], Stream a)
split p (a :< as)
| p (extract as) = splitComplete p a (:< as)
| (ts, fs) <- split p as = (foldr (:) ts a, fs)
splitComplete :: (a -> Bool) -> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitComplete _ t@Tip{} f = ([], f t)
splitComplete p t@(Bin _ a l r) f
| p a = ([], f t)
| p (extract r), (ts, fs) <- splitComplete p l (:< f r) = (a:ts, fs)
| (ts, fs) <- splitComplete p r f = (a:foldr (:) ts l, fs)
splitW :: (Stream a -> Bool) -> Stream a -> ([a], Stream a)
splitW p (a :< as)
| p as = splitCompleteW p a (:< as)
| (ts, fs) <- splitW p as = (foldr (:) ts a, fs)
splitCompleteW :: (Stream a -> Bool) -> Complete a -> (Complete a -> Stream a) -> ([a], Stream a)
splitCompleteW _ t@Tip{} f = ([], f t)
splitCompleteW p t@(Bin _ a l r) f
| w <- f t, p w = ([], w)
| w <- f r, p w, (ts, fs) <- splitCompleteW p l (:< w) = (a:ts, fs)
| (ts, fs) <- splitCompleteW p r f = (a:foldr (:) ts l, fs)
fromList :: [a] -> Stream a
fromList = foldr (<|) (error "fromList: finite list")
insert :: Ord a => a -> Stream a -> Stream a
insert a as | (ts, as') <- split (a<=) as = foldr (<|) (a <| as') ts
insertBy :: (a -> a -> Ordering) -> a -> Stream a -> Stream a
insertBy cmp a as | (ts, as') <- split (\b -> cmp a b <= EQ) as = foldr (<|) (a <| as') ts
adjust :: Integer -> (a -> a) -> Stream a -> Stream a
adjust !n f (a :< as)
| n < w = adjustComplete n f a :< as
| otherwise = a :< adjust (n w) f as
where w = weight a
adjustComplete :: Integer -> (a -> a) -> Complete a -> Complete a
adjustComplete 0 f (Tip a) = Tip (f a)
adjustComplete _ _ t@Tip{} = t
adjustComplete n f (Bin m a l r)
| n == 0 = Bin m (f a) l r
| n < w = Bin m a (adjustComplete (n 1) f l) r
| otherwise = Bin m a l (adjustComplete (n 1 w) f r)
where w = weight l
update :: Integer -> a -> Stream a -> Stream a
update n = adjust n . const