module Data.Stream.Future.Skew
( Future(..)
, (<|)
, cons
, length
, head
, tail
, tails
, last
, uncons
, index
, drop
, dropWhile
, indexed
, from
, break
, span
, split
, splitW
, repeat
, replicate
, insert
, insertBy
, update
, adjust
, fromList
, toFuture
) where
import Control.Applicative hiding (empty)
import Control.Comonad
import Data.Functor.Alt
import Data.Functor.Apply
import Data.Foldable hiding (toList)
import Data.Traversable (Traversable, traverse)
import qualified Data.Traversable as Traversable
import Data.Semigroup hiding (Last)
import Data.Semigroup.Foldable
import Data.Semigroup.Traversable
import Data.Monoid (Monoid(mappend))
import Prelude hiding (null, head, tail, drop, dropWhile, length, foldr, last, span, repeat, replicate, break)
infixr 5 :<, <|
data Complete a
= Tip a
| Bin !Int 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
extend f w@Tip {} = Tip (f w)
extend f w@(Bin n _ l r) = Bin n (f w) (extend f l) (extend f r)
instance Comonad Complete where
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 -> Int
weight Tip{} = 1
weight (Bin w _ _ _) = w
data Future a
= Last !(Complete a)
| !(Complete a) :< Future a
instance Show a => Show (Future a) where
showsPrec d as = showParen (d >= 10) $
showString "fromList " . showsPrec 11 (toList as)
instance Functor Future where
fmap f (t :< ts) = fmap f t :< fmap f ts
fmap f (Last t) = Last (fmap f t)
instance Extend Future where
extend g (Last t) = Last (extendTree g t Last)
extend g (t :< ts) = extendTree g t (:< ts) :< extend g ts
instance Comonad Future where
extract = head
extendTree :: (Future a -> b) -> Complete a -> (Complete a -> Future a) -> Complete b
extendTree g w@Tip{} f = Tip (g (f w))
extendTree g w@(Bin n _ l r) f = Bin n (g (f w)) (extendTree g l (:< f r)) (extendTree g r f)
instance Apply Future where
Last (Tip f) <.> as = singleton (f (extract as))
fs <.> Last (Tip a) = singleton (extract fs a)
Last (Bin _ f lf rf) <.> Last (Bin _ a la ra) = f a <| (lf :< Last rf <.> la :< Last ra )
Last (Bin _ f lf rf) <.> Bin _ a la ra :< as = f a <| (lf :< Last rf <.> la :< ra :< as)
Last (Bin _ f lf rf) <.> Tip a :< as = f a <| (lf :< Last rf <.> as )
Bin _ f lf rf :< fs <.> Last (Bin _ a la ra) = f a <| (lf :< rf :< fs <.> la :< Last ra )
Bin _ f lf rf :< fs <.> Tip a :< as = f a <| (lf :< rf :< fs <.> as )
Bin _ f lf rf :< fs <.> Bin _ a la ra :< as = f a <| (lf :< rf :< fs <.> la :< ra :< as)
Tip f :< fs <.> Tip a :< as = f a <| (fs <.> as )
Tip f :< fs <.> Bin _ a la ra :< as = f a <| (fs <.> la :< ra :< as)
Tip f :< fs <.> Last (Bin _ a la ra) = f a <| (fs <.> la :< Last ra )
instance Applicative Future where
pure = repeat
(<*>) = (<.>)
instance Alt Future where
as <!> bs = foldr (<|) bs as
instance Foldable Future where
foldMap f (t :< ts) = foldMap f t `mappend` foldMap f ts
foldMap f (Last t) = foldMap f t
foldr f z (t :< ts) = foldr f (foldr f z ts) t
foldr f z (Last t) = foldr f z t
toList :: Future a -> [a]
toList = foldr (:) []
instance Foldable1 Future where
foldMap1 f (t :< ts) = foldMap1 f t <> foldMap1 f ts
foldMap1 f (Last t) = foldMap1 f t
instance Traversable Future where
traverse f (t :< ts) = (:<) <$> traverse f t <*> traverse f ts
traverse f (Last t) = Last <$> traverse f t
instance Traversable1 Future where
traverse1 f (t :< ts) = (:<) <$> traverse1 f t <.> traverse1 f ts
traverse1 f (Last t) = Last <$> traverse1 f t
repeat :: a -> Future a
repeat a0 = go a0 (Tip a0)
where
go :: a -> Complete a -> Future a
go a as | ass <- bin a as as = as :< go a ass
replicate :: Int -> a -> Future a
replicate n a
| n <= 0 = error "replicate: non-positive argument"
| otherwise = go 1 n a (Tip a) (\0 r -> r)
where
go :: Int -> Int -> b -> Complete b -> (Int -> Future b -> r) -> r
go !i !j b tb k
| j >= i2p1 = go i2p1 j b (Bin i2p1 b tb tb) k'
| j >= i2 = k (j i2) (tb :< Last tb)
| otherwise = k (j i) (Last tb)
where
i2 = i * 2
i2p1 = i2 + 1
k' r xs
| r >= i2 = k (r i2) (tb :< tb :< xs)
| r >= i = k (r i) (tb :< xs)
| otherwise = k r xs
mapWithIndex :: (Int -> a -> b) -> Future a -> Future b
mapWithIndex f0 as0 = spine f0 0 as0
where
spine f m (Last as) = Last (tree f m as)
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)
indexed :: Future a -> Future (Int, a)
indexed = mapWithIndex (,)
from :: Num a => a -> Future a
from a = mapWithIndex ((+) . fromIntegral) (pure a)
singleton :: a -> Future a
singleton a = Last (Tip a)
length :: Future a -> Int
length (Last t) = weight t
length (t :< ts) = weight t + length ts
(<|) :: a -> Future a -> Future a
a <| (l :< Last r)
| weight l == weight r = Last (bin a l r)
a <| (l :< r :< as)
| weight l == weight r = bin a l r :< as
a <| as = Tip a :< as
cons :: a -> Future a -> Future a
cons = (<|)
head :: Future a -> a
head (a :< _) = extract a
head (Last a) = extract a
tail :: Future a -> Maybe (Future a)
tail (Tip{} :< ts) = Just ts
tail (Bin _ _ l r :< ts) = Just (l :< r :< ts)
tail (Last Tip{}) = Nothing
tail (Last (Bin _ _ l r)) = Just (l :< Last r)
tails :: Future a -> Future (Future a)
tails = duplicate
last :: Future a -> a
last (_ :< as) = last as
last (Last as) = go as
where go (Tip a) = a
go (Bin _ _ _ r) = go r
uncons :: Future a -> (a, Maybe (Future a))
uncons (Last (Tip a)) = (a, Nothing)
uncons (Last (Bin _ a l r)) = (a, Just (l :< Last r))
uncons (Tip a :< as) = (a, Just as)
uncons (Bin _ a l r :< as) = (a, Just (l :< r :< as))
index :: Int -> Future a -> a
index i (Last t)
| i < weight t = indexComplete i t
| otherwise = error "index: out of range"
index i (t :< ts)
| i < w = indexComplete i t
| otherwise = index (i w) ts
where w = weight t
indexComplete :: Int -> Complete a -> a
indexComplete 0 (Tip a) = a
indexComplete i (Bin w a l r)
| i == 0 = a
| i <= w' = indexComplete (i1) l
| otherwise = indexComplete (i1w') r
where w' = div w 2
indexComplete _ _ = error "index: index out of range"
drop :: Int -> Future a -> Maybe (Future a)
drop 0 ts = Just ts
drop i (t :< ts) = case compare i w of
LT -> Just (dropComplete i t (:< ts))
EQ -> Just ts
GT -> drop (i w) ts
where w = weight t
drop i (Last t)
| i < w = Just (dropComplete i t Last)
| otherwise = Nothing
where w = weight t
dropComplete :: Int -> Complete a -> (Complete a -> Future a) -> Future 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 "drop: index out of range"
dropWhile :: (a -> Bool) -> Future a -> Maybe (Future a)
dropWhile p as
| p (head as) = tail as >>= dropWhile p
| otherwise = Just as
span :: (a -> Bool) -> Future a -> ([a], Maybe (Future a))
span p aas = case uncons aas of
(a, Just as) | p a, (ts, fs) <- span p as -> (a:ts, fs)
(a, Nothing) | p a -> ([a], Nothing)
(_, _) -> ([], Just aas)
break :: (a -> Bool) -> Future a -> ([a], Maybe (Future a))
break p = span (not . p)
split :: (a -> Bool) -> Future a -> ([a], Maybe (Future a))
split p l@(Last a)
| p (extract a) = ([], Just l)
| otherwise = splitComplete p a Last
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 -> Future a) -> ([a], Maybe (Future a))
splitComplete p t@(Tip a) f
| p a = ([], Just (f t))
| otherwise = ([a], Nothing)
splitComplete p t@(Bin _ a l r) f
| p a = ([], Just (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 :: (Future a -> Bool) -> Future a -> ([a], Maybe (Future a))
splitW p l@(Last a)
| p l = ([], Just l)
| otherwise = splitCompleteW p a Last
splitW p (a :< as)
| p as = splitCompleteW p a (:< as)
| (ts, fs) <- splitW p as = (foldr (:) ts a, fs)
splitCompleteW :: (Future a -> Bool) -> Complete a -> (Complete a -> Future a) -> ([a], Maybe (Future a))
splitCompleteW p t@(Tip a) f
| w <- f t, p w = ([], Just w)
| otherwise = ([a], Nothing)
splitCompleteW p t@(Bin _ a l r) f
| w <- f t, p w = ([], Just 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] -> Future a
fromList [] = error "fromList: empty list"
fromList (x:xs) = go x xs
where go a [] = singleton a
go a (b:bs) = a <| go b bs
toFuture :: [a] -> Maybe (Future a)
toFuture [] = Nothing
toFuture xs = Just (fromList xs)
insert :: Ord a => a -> Future a -> Future a
insert a as = case split (a<=) as of
(_, Nothing) -> foldr (<|) (singleton a) as
(ts, Just as') -> foldr (<|) (a <| as') ts
insertBy :: (a -> a -> Ordering) -> a -> Future a -> Future a
insertBy cmp a as = case split (\b -> cmp a b <= EQ) as of
(_, Nothing) -> foldr (<|) (singleton a) as
(ts, Just as') -> foldr (<|) (a <| as') ts
adjust :: Int -> (a -> a) -> Future a -> Future a
adjust !n f d@(Last a)
| n < weight a = Last (adjustComplete n f a)
| otherwise = d
adjust !n f (a :< as)
| n < w = adjustComplete n f a :< as
| otherwise = a :< adjust (n w) f as
where w = weight a
adjustComplete :: Int -> (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 :: Int -> a -> Future a -> Future a
update n = adjust n . const