diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,3 @@ +### 0.1.0.0 -- 2024-06-16 + +* First version. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2024, Soumik Sarkar + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,43 @@ +# seqn + +[![Hackage](https://img.shields.io/hackage/v/seqn?logo=haskell&color=blue)](https://hackage.haskell.org/package/seqn) +[![Haskell-CI](https://github.com/meooow25/seqn/actions/workflows/haskell-ci.yml/badge.svg)](https://github.com/meooow25/seqn/actions/workflows/haskell-ci.yml) + +`seqn` offers two sequence types: + +* `Seq`, an immutable sequence supporting operations such as index, insert, + delete, split, append, in logarithmic time. `Seq` is well-suited to use cases + where there are frequent changes to the structure of the sequence. + +* `MSeq`, a sequence like `Seq`, which additionally supports constant time + access to the accumulated "measure" of all its elements. See the documentation + for `MSeq` for more about measures. + +`seqn` also offers a priority-queue structure, `PQueue`, with logarithmic time +queue operations. + +## Documentation + +Please find the documentation on Hackage: [seqn](https://hackage.haskell.org/package/seqn) + +## Alternatives + +The following structures are similar to `Seq` and `MSeq`, but may be better +suited to some use cases. + +* Alternatives to `Seq`: + * [`Data.Sequence.Seq`](https://hackage.haskell.org/package/containers-0.7/docs/Data-Sequence.html#t:Seq) + from `containers` + * [`Data.RRBVector.Vector`](https://hackage.haskell.org/package/rrb-vector-0.2.1.0/docs/Data-RRBVector.html#t:Vector) + from `rrb-vector` +* Alternatives to `MSeq`: + * [`Data.FingerTree.FingerTree`](https://hackage.haskell.org/package/fingertree-0.1.5.0/docs/Data-FingerTree.html#t:FingerTree) + from `fingertree` + +For a detailed comparison, [see here](https://github.com/meooow25/seqn/tree/master/bench). + +## Acknowledgements + +The interface and implementation of `seqn` is largely influenced by +the libraries [`containers`](https://hackage.haskell.org/package/containers) and +[`fingertree`](https://hackage.haskell.org/package/fingertree). diff --git a/seqn.cabal b/seqn.cabal new file mode 100644 --- /dev/null +++ b/seqn.cabal @@ -0,0 +1,91 @@ +cabal-version: 3.0 +name: seqn +version: 0.1.0.0 +synopsis: Sequences and measured sequences +license: BSD-3-Clause +license-file: LICENSE +author: Soumik Sarkar +maintainer: soumiksarkar.3120@gmail.com +copyright: (c) 2024 Soumik Sarkar +category: Data Structures +homepage: https://github.com/meooow25/seqn +bug-reports: https://github.com/meooow25/seqn/issues +build-type: Simple + +description: + Sequences and measured sequences with logarithmic time index, split, append. + +extra-doc-files: + README.md + CHANGELOG.md + +tested-with: + GHC == 8.8.4 + , GHC == 8.10.7 + , GHC == 9.0.2 + , GHC == 9.2.8 + , GHC == 9.4.8 + , GHC == 9.6.4 + , GHC == 9.8.1 + +source-repository head + type: git + location: https://github.com/meooow25/seqn.git + +common warnings + ghc-options: + -Wall -Wcompat -Widentities -Wredundant-constraints -Wunused-packages + +library + import: warnings + + exposed-modules: + Data.Seqn.Seq + Data.Seqn.MSeq + Data.Seqn.PQueue + Data.Seqn.Internal.Seq + Data.Seqn.Internal.Tree + Data.Seqn.Internal.MSeq + Data.Seqn.Internal.MTree + Data.Seqn.Internal.PQueue + Data.Seqn.Internal.Util + other-modules: + Data.Seqn.Internal.Stream + Data.Seqn.Internal.KMP + + build-depends: + base >= 4.13.0.0 && < 5 + , deepseq >= 1.4.4.0 && < 1.7 + , indexed-traversable >= 0.1 && < 0.2 + , primitive >= 0.7.3.0 && < 0.10 + , samsort >= 0.1.0.0 && < 0.2 + , transformers >= 0.5.6.2 && < 0.7 + + hs-source-dirs: src + default-language: Haskell2010 + +test-suite seqn-test + import: warnings + default-language: Haskell2010 + + main-is: Main.hs + + other-modules: + ListLikeTests + ListExtra + MSeq + PQueue + Seq + TestUtil + + type: exitcode-stdio-1.0 + hs-source-dirs: test + build-depends: + base + , indexed-traversable + , QuickCheck + , quickcheck-classes-base + , seqn + , tasty + , tasty-quickcheck + , transformers diff --git a/src/Data/Seqn/Internal/KMP.hs b/src/Data/Seqn/Internal/KMP.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/KMP.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE BangPatterns #-} + +module Data.Seqn.Internal.KMP + ( Table + , State + , build + , step + ) where + +import Data.Primitive.Array (Array, indexArray, sizeofArray) +import Data.Primitive.PrimArray + ( PrimArray + , indexPrimArray + , newPrimArray + , readPrimArray + , runPrimArray + , writePrimArray + ) + +-- Knuth–Morris–Pratt algorithm +-- See https://en.wikipedia.org/wiki/Knuth%E2%80%93Morris%E2%80%93Pratt_algorithm +-- +-- In Table xa pa, +-- * xa is the pattern. +-- * pa is the prefix function. pa!i is the length of longest proper prefix of +-- xa that ends at index i of xa. + +data Table a = Table + {-# UNPACK #-} !(Array a) + {-# UNPACK #-} !(PrimArray Int) + +newtype State a = State Int + +-- Precondition: 0 < length xa +build :: Eq a => Array a -> (Table a, State a) +build xa + | n <= 0 = error "non-positive length" + | otherwise = (Table xa pa, State 0) + where + n = sizeofArray xa + !pa = runPrimArray $ do + pma <- newPrimArray n + writePrimArray pma 0 0 + for_ 1 (n-1) $ \i -> do + let go j | indexArray xa i == indexArray xa j = pure (j+1) + go 0 = pure 0 + go j = readPrimArray pma (j-1) >>= go + readPrimArray pma (i-1) >>= go >>= writePrimArray pma i + pure pma +{-# INLINABLE build #-} + +step :: Eq a => Table a -> State a -> a -> (Bool, State a) +step (Table xa pa) (State i) x = go i + where + go j | indexArray xa j == x = + if j+1 == sizeofArray xa + then (,) True $! State (indexPrimArray pa j) + else (False, State (j+1)) + go 0 = (False, State 0) + go j = go (indexPrimArray pa (j-1)) +{-# INLINABLE step #-} + +for_ :: Applicative f => Int -> Int -> (Int -> f a) -> f () +for_ !i1 !i2 f = go i1 + where + go i = if i > i2 then pure () else f i *> go (i+1) +{-# INLINE for_ #-} diff --git a/src/Data/Seqn/Internal/MSeq.hs b/src/Data/Seqn/Internal/MSeq.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/MSeq.hs @@ -0,0 +1,1432 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- This is an internal module. You probably don't need to import this. Use +-- "Data.Seqn.MSeq" instead. +-- +-- = WARNING +-- +-- Definitions in this module allow violating invariants that would otherwise be +-- guaranteed by "Data.Seqn.MSeq". Use at your own risk! +-- +module Data.Seqn.Internal.MSeq + ( + -- * MSeq + MSeq(..) + + -- * Construct + , empty + , singleton + , fromList + , fromRevList + , replicate + , replicateA + , generate + , generateA + , unfoldr + , unfoldl + , unfoldrM + , unfoldlM + , concatMap + , mfix + + -- * Convert + , toRevList + + -- * Index + , lookup + , index + , (!?) + , (!) + , update + , adjust + , insertAt + , deleteAt + + -- * Slice + , cons + , snoc + , uncons + , unsnoc + , take + , drop + , slice + , splitAt + , takeEnd + , dropEnd + , splitAtEnd + + -- * Filter + , filter + , mapMaybe + , mapEither + , filterA + , mapMaybeA + , mapEitherA + , takeWhile + , dropWhile + , span + , break + , takeWhileEnd + , dropWhileEnd + , spanEnd + , breakEnd + + -- * Transform + , map + , liftA2 + , traverse + , imap + , itraverse + , reverse + , intersperse + , scanl + , scanr + , sort + , sortBy + + -- * Search and test + , findEnd + , findIndex + , findIndexEnd + , infixIndices + , binarySearchFind + , isPrefixOf + , isSuffixOf + , isInfixOf + , isSubsequenceOf + + -- * Zip and unzip + , zipWith + , zipWith3 + , zipWithM + , zipWith3M + , unzipWith + , unzipWith3 + + -- * Measured queries + , summaryMay + , summary + , binarySearchPrefix + , binarySearchSuffix + + -- * Force + , liftRnf2 + + -- * Internal + , fromMTree + + -- * Testing + , valid + , debugShowsPrec + ) where + +import Prelude hiding (break, concatMap, drop, dropWhile, filter, liftA2, lookup, map, replicate, reverse, scanl, scanr, span, splitAt, take, takeWhile, traverse, unzip, unzip3, zip, zip3, zipWith, zipWith3) +import qualified Control.Applicative as Ap +import Control.Applicative.Backwards (Backwards(..)) +import Control.DeepSeq (NFData(..)) +import Data.Bifunctor (Bifunctor(..)) +import Data.Coerce (coerce) +import qualified Data.Foldable as F +import qualified Data.Foldable.WithIndex as IFo +import Data.Functor.Classes (Eq1(..), Ord1(..), Show1(..)) +import Data.Functor.Const (Const(..)) +import Data.Functor.Identity (Identity(..)) +import Data.List.NonEmpty (NonEmpty(..)) +import qualified Data.Monoid as Monoid +import qualified Data.Primitive.Array as A +import Data.Semigroup (Semigroup(..)) +import qualified Data.SamSort as Sam +import qualified GHC.Exts as X +import Text.Read (Read(..)) +import qualified Text.Read as Read + +import Data.Seqn.Internal.MTree (Measured(..), MTree(..)) +import qualified Data.Seqn.Internal.MTree as T +import qualified Data.Seqn.Internal.Util as U +import qualified Data.Seqn.Internal.Stream as Stream +import Data.Seqn.Internal.Stream (Step(..), Stream(..)) +import qualified Data.Seqn.Internal.KMP as KMP + +-------- +-- Seq +-------- + +-- | A sequence with elements of type @a@. An instance of @'Measured' a@ is +-- required for most operations. +data MSeq a + = MTree !a !(MTree a) + | MEmpty +-- See Note [Seq structure] in Data.Seqn.Internal.Seq + +-------------- +-- Instances +-------------- + +instance Eq a => Eq (MSeq a) where + t1 == t2 = compareLength t1 t2 == EQ && stream t1 == stream t2 + {-# INLINABLE (==) #-} + +-- | Lexicographical ordering +instance Ord a => Ord (MSeq a) where + compare t1 t2 = compare (stream t1) (stream t2) + {-# INLINABLE compare #-} + +instance Show a => Show (MSeq a) where + showsPrec _ t = shows (F.toList t) + {-# INLINABLE showsPrec #-} + +instance (Measured a, Read a) => Read (MSeq a) where + readPrec = fmap fromList readListPrec + {-# INLINABLE readPrec #-} + + readListPrec = Read.readListPrecDefault + {-# INLINABLE readListPrec #-} + +instance Eq1 MSeq where + liftEq f t1 t2 = compareLength t1 t2 == EQ && liftEq f (stream t1) (stream t2) + {-# INLINE liftEq #-} + +instance Ord1 MSeq where + liftCompare f t1 t2 = liftCompare f (stream t1) (stream t2) + {-# INLINE liftCompare #-} + +instance Show1 MSeq where + liftShowsPrec _ sl _ t = sl (F.toList t) + {-# INLINE liftShowsPrec #-} + +-- | +-- [@length@]: \(O(1)\). +-- +-- Folds are \(O(n)\). +instance Foldable MSeq where + fold = foldMap id + {-# INLINABLE fold #-} + + foldMap f = \case + MTree x xs -> f x <> T.foldMap f xs + MEmpty -> mempty + {-# INLINE foldMap #-} + + foldMap' f = F.foldl' (\z x -> z <> f x) mempty + {-# INLINE foldMap' #-} + + foldr f z = Stream.foldr f z . stream + {-# INLINE foldr #-} + + foldl f z = Stream.foldr (flip f) z . streamEnd + {-# INLINE foldl #-} + + foldl' f !z = \case + MTree x xs -> T.foldl' f (f z x) xs + MEmpty -> z + {-# INLINE foldl' #-} + + foldr' f !z = \case + MTree x xs -> f x $! T.foldr' f z xs + MEmpty -> z + {-# INLINE foldr' #-} + + null = \case + MTree _ _ -> False + MEmpty -> True + + length = \case + MTree _ xs -> 1 + T.size xs + MEmpty -> 0 + +instance Measured a => X.IsList (MSeq a) where + type Item (MSeq a) = a + fromList = fromList + {-# INLINE fromList #-} + + toList = F.toList + {-# INLINE toList #-} + +instance IFo.FoldableWithIndex Int MSeq where + ifoldMap f = \case + MTree x xs -> f 0 x <> T.ifoldMap f 1 xs + MEmpty -> mempty + {-# INLINE ifoldMap #-} + + ifoldr f z = Stream.ifoldr f z 0 (+1) . stream + {-# INLINE ifoldr #-} + + ifoldl f z = \t -> + Stream.ifoldr (flip . f) z (length t - 1) (subtract 1) (streamEnd t) + {-# INLINE ifoldl #-} + + ifoldr' f !z = \case + MTree x xs -> f 0 x $! T.ifoldr' f z (T.size xs) xs + MEmpty -> z + {-# INLINE ifoldr' #-} + + ifoldl' f !z = \case + MTree x xs -> T.ifoldl' f (f 0 z x) 1 xs + MEmpty -> z + {-# INLINE ifoldl' #-} + +-- | +-- [@(<>)@]: \(O(\left| \log n_1 - \log n_2 \right|)\). Concatenates two +-- sequences. +-- +-- [@stimes@]: \(O(\log c)\). @stimes c xs@ is @xs@ repeating @c@ times. If +-- @c < 0@, 'empty' is returned. +instance Measured a => Semigroup (MSeq a) where + MTree x xs <> MTree y ys = MTree x (T.link y xs ys) + l <> MEmpty = l + MEmpty <> r = r + {-# INLINABLE (<>) #-} + + stimes !c = \case + t@(MTree x xs) + | c <= 0 -> MEmpty + | fromIntegral c * toi (length t) > toi (maxBound :: Int) -> + error "MSeq.stimes: result size too large" + | otherwise -> MTree x (stimesLoop (c'-1) x xs xs) + where + c' = fromIntegral c :: Int + toi :: Int -> Integer + toi = fromIntegral + MEmpty -> MEmpty + {-# INLINABLE stimes #-} + -- See Note [Complexity of stimes] in Data.Seqn.Internal.Seq + + sconcat (x:|xs) = mconcat (x:xs) + +stimesLoop :: Measured a => Int -> a -> MTree a -> MTree a -> MTree a +stimesLoop c !x !xs !acc + | c <= 0 = acc + | c `mod` 2 == 0 = stimesLoop (c `div` 2) x (T.bin x xs xs) acc + | otherwise = stimesLoop (c `div` 2) x (T.bin x xs xs) (T.link x xs acc) +{-# INLINABLE stimesLoop #-} + +-- | +-- [@mempty@]: The empty sequence. +instance Measured a => Monoid (MSeq a) where + mempty = MEmpty + + mconcat = concatMap id + {-# INLINE mconcat #-} -- Inline for fusion + +instance (NFData (Measure a), NFData a) => NFData (MSeq a) where + rnf = \case + MTree x xs -> rnf x `seq` rnf xs + MEmpty -> () + {-# INLINABLE rnf #-} + +-------------- +-- Construct +-------------- + +-- | The empty sequence. +empty :: MSeq a +empty = MEmpty + +-- | A singleton sequence. +singleton :: a -> MSeq a +singleton x = MTree x T.MTip + +-- | \(O(n)\). Create an @MSeq@ from a list. +fromList :: Measured a => [a] -> MSeq a +fromList = ltrFinish . F.foldl' ltrPush Nil +{-# INLINE fromList #-} +-- See Note [fromList implementation] + +-- | \(O(n)\). Create an @MSeq@ from a reversed list. +fromRevList :: Measured a => [a] -> MSeq a +fromRevList = rtlFinish . F.foldl' (flip rtlPush) Nil +{-# INLINE fromRevList #-} +-- See Note [fromList implementation] + +-- | \(O(\log n)\). A sequence with a repeated element. +-- If the length is negative, 'empty' is returned. +replicate :: Measured a => Int -> a -> MSeq a +replicate !n x = stimes n (MTree x MTip) +{-# INLINABLE replicate #-} + +-- | \(O(n)\). Generate a sequence from a length and an applicative action. +-- If the length is negative, 'empty' is returned. +replicateA :: (Measured a, Applicative f) => Int -> f a -> f (MSeq a) +replicateA !n m = generateA n (const m) +{-# INLINABLE replicateA #-} + +-- | \(O(n)\). Generate a sequence from a length and a generator. +-- If the length is negative, 'empty' is returned. +generate :: Measured a => Int -> (Int -> a) -> MSeq a +generate = + (coerce :: (Int -> (Int -> Identity a) -> Identity (MSeq a)) + -> Int -> (Int -> a) -> MSeq a) + generateA +{-# INLINE generate #-} + +-- | \(O(n)\). Generate a sequence from a length and an applicative generator. +-- If the length is negative, 'empty' is returned. +generateA + :: (Measured a, Applicative f) => Int -> (Int -> f a) -> f (MSeq a) +generateA n f + | n <= 0 = pure MEmpty + | otherwise = Ap.liftA2 MTree (f 0) (T.generateA f 1 (n-1)) +{-# INLINE generateA #-} + +-- | \(O(n)\). Unfold a sequence from left to right. +unfoldr :: Measured a => (b -> Maybe (a, b)) -> b -> MSeq a +unfoldr = + (coerce :: ((b -> Identity (Maybe (a, b))) -> b -> Identity (MSeq a)) + -> (b -> Maybe (a, b)) -> b -> MSeq a) + unfoldrM +{-# INLINE unfoldr #-} + +-- | \(O(n)\). Unfold a sequence monadically from left to right. +unfoldrM :: (Measured a, Monad m) => (b -> m (Maybe (a, b))) -> b -> m (MSeq a) +unfoldrM f = go Nil + where + go !b z = f z >>= \case + Nothing -> pure $! ltrFinish b + Just (x, z') -> go (ltrPush b x) z' +{-# INLINE unfoldrM #-} + +-- | \(O(n)\). Unfold a sequence from right to left. +unfoldl :: Measured a => (b -> Maybe (b, a)) -> b -> MSeq a +unfoldl = + (coerce :: ((b -> Identity (Maybe (b, a))) -> b -> Identity (MSeq a)) + -> (b -> Maybe (b, a)) -> b -> MSeq a) + unfoldlM +{-# INLINE unfoldl #-} + +-- | \(O(n)\). Unfold a sequence monadically from right to left. +unfoldlM :: (Measured a, Monad m) => (b -> m (Maybe (b, a))) -> b -> m (MSeq a) +unfoldlM f = go Nil + where + go !b z = f z >>= \case + Nothing -> pure $! rtlFinish b + Just (z', x) -> go (rtlPush x b) z' +{-# INLINE unfoldlM #-} + +-- | \(O \left(\sum_i \log n_i \right)\). +-- Map over a @Foldable@ and concatenate the results. +concatMap :: (Measured b, Foldable f) => (a -> MSeq b) -> f a -> MSeq b +concatMap f = ltrFinish . F.foldl' g Nil + where + g b x = case f x of + MEmpty -> b + MTree y ys -> ltrPushMany b y ys + {-# INLINE g #-} +{-# INLINE concatMap #-} +-- See Note [concatMap implementation] + +-- | Monadic fixed point. See "Control.Monad.Fix". +mfix :: Measured a => (a -> MSeq a) -> MSeq a +mfix f = + imap + (\i _ -> let x = index i (f x) in x) + (f (error "MSeq.mfix: f must be lazy")) +{-# INLINE mfix #-} + +------------ +-- Convert +------------ + +-- | \(O(n)\). Convert to a list in reverse. +-- +-- To convert to a list without reversing, use +-- @Data.Foldable.'Data.Foldable.toList'@. +toRevList :: MSeq a -> [a] +toRevList t = X.build $ \lcons lnil -> F.foldl (flip lcons) lnil t +{-# INLINE toRevList #-} + +---------- +-- Index +---------- + +-- | \(O(\log n)\). Look up the element at an index. +lookup :: Int -> MSeq a -> Maybe a +lookup !i (MTree x xs) + | i < 0 || T.size xs < i = Nothing + | i == 0 = Just x + | otherwise = Just $! T.index (i-1) xs +lookup _ MEmpty = Nothing +{-# INLINE lookup #-} + +-- | \(O(\log n)\). Look up the element at an index. Calls @error@ if the index +-- is out of bounds. +index :: Int -> MSeq a -> a +index !i = \case + MTree x xs + | i == 0 -> x + | otherwise -> T.index (i-1) xs + MEmpty -> error "MSeq.index: out of bounds" + +-- | \(O(\log n)\). Infix version of 'lookup'. +(!?) :: MSeq a -> Int -> Maybe a +(!?) = flip lookup + +-- | \(O(\log n)\). Infix version of 'index'. Calls @error@ if the index is out +-- of bounds. +(!) :: MSeq a -> Int -> a +(!) = flip index + +-- | \(O(\log n)\). Update an element at an index. If the index is out of +-- bounds, the sequence is returned unchanged. +update :: Measured a => Int -> a -> MSeq a -> MSeq a +update i x = adjust (const x) i +{-# INLINABLE update #-} + +-- | \(O(\log n)\). Adjust the element at an index. If the index is out of +-- bounds, the sequence is returned unchanged. +adjust :: Measured a => (a -> a) -> Int -> MSeq a -> MSeq a +adjust f !i t = case t of + MTree x xs + | i < 0 || T.size xs < i -> t + | i == 0 -> MTree (f x) xs + | otherwise -> MTree x (runIdentity (T.adjustF (Identity U.#. f) (i-1) xs)) + MEmpty -> MEmpty +{-# INLINE adjust #-} + +-- | \(O(\log n)\). Insert an element at an index. If the index is out of +-- bounds, the element is added to the closest end of the sequence. +insertAt :: Measured a => Int -> a -> MSeq a -> MSeq a +insertAt !i y t = case t of + MTree x xs + | i <= 0 -> cons y t + | otherwise -> MTree x (T.insertAt (i-1) y xs) + MEmpty -> singleton y +{-# INLINABLE insertAt #-} + +-- | \(O(\log n)\). Delete an element at an index. If the index is out of +-- bounds, the sequence is returned unchanged. +deleteAt :: Measured a => Int -> MSeq a -> MSeq a +deleteAt !i t = case t of + MTree x xs + | i < 0 || T.size xs < i -> t + | i == 0 -> fromMTree xs + | otherwise -> MTree x (T.deleteAt (i-1) xs) + MEmpty -> MEmpty +{-# INLINABLE deleteAt #-} + +---------- +-- Slice +---------- + +-- | \(O(\log n)\). Append a value to the beginning of a sequence. +cons :: Measured a => a -> MSeq a -> MSeq a +cons x (MTree y ys) = MTree x (T.cons y ys) +cons x MEmpty = singleton x +{-# INLINABLE cons #-} + +-- | \(O(\log n)\). Append a value to the end of a sequence. +snoc :: Measured a => MSeq a -> a -> MSeq a +snoc (MTree y ys) x = MTree y (T.snoc ys x) +snoc MEmpty x = singleton x +{-# INLINABLE snoc #-} + +-- | \(O(\log n)\). The head and tail of a sequence. +uncons :: Measured a => MSeq a -> Maybe (a, MSeq a) +uncons (MTree x xs) = Just . (,) x $! fromMTree xs +uncons MEmpty = Nothing +{-# INLINE uncons #-} + +-- | \(O(\log n)\). The init and last of a sequence. +unsnoc :: Measured a => MSeq a -> Maybe (MSeq a, a) +unsnoc (MTree x xs) = case T.unsnoc xs of + U.SNothing -> Just (MEmpty, x) + U.SJust (U.S2 ys y) -> Just (MTree x ys, y) +unsnoc MEmpty = Nothing +{-# INLINE unsnoc #-} + +-- | \(O(\log n)\). Take a number of elements from the beginning of a sequence. +take :: Measured a => Int -> MSeq a -> MSeq a +take !i t@(MTree x xs) + | i <= 0 = MEmpty + | T.size xs < i = t + | otherwise = MTree x (getConst (T.splitAtF (i-1) xs)) +take _ MEmpty = MEmpty +{-# INLINABLE take #-} + +-- | \(O(\log n)\). Drop a number of elements from the beginning of a sequence. +drop :: Measured a => Int -> MSeq a -> MSeq a +drop !i t@(MTree _ xs) + | i <= 0 = t + | T.size xs < i = MEmpty + | otherwise = case U.unTagged (T.splitAtF (i-1) xs) of + U.S2 x' xs' -> MTree x' xs' +drop _ MEmpty = MEmpty +{-# INLINABLE drop #-} + +-- | \(O(\log n)\). The slice of a sequence between two indices (inclusive). +slice :: Measured a => (Int, Int) -> MSeq a -> MSeq a +slice (i,j) = drop i . take (j+1) +{-# INLINABLE slice #-} + +-- | \(O(\log n)\). Take a number of elements from the end of a sequence. +takeEnd :: Measured a => Int -> MSeq a -> MSeq a +takeEnd n t = drop (length t - n) t +{-# INLINABLE takeEnd #-} + +-- | \(O(\log n)\). Drop a number of elements from the end of a sequence. +dropEnd :: Measured a => Int -> MSeq a -> MSeq a +dropEnd n t = take (length t - n) t +{-# INLINABLE dropEnd #-} + +-- | \(O(\log n)\). Split a sequence at a given index. +-- +-- @splitAt n xs == ('take' n xs, 'drop' n xs)@ +splitAt :: Measured a => Int -> MSeq a -> (MSeq a, MSeq a) +splitAt !i t@(MTree x xs) + | i <= 0 = (MEmpty, t) + | T.size xs < i = (t, MEmpty) + | otherwise = case T.splitAtF (i-1) xs of + U.S2 xs1 (U.S2 x' xs2) -> (MTree x xs1, MTree x' xs2) +splitAt _ MEmpty = (MEmpty, MEmpty) +{-# INLINABLE splitAt #-} + +-- | \(O(\log n)\). Split a sequence at a given index from the end. +-- +-- @splitAtEnd n xs == ('dropEnd' n xs, 'takeEnd' n xs)@ +splitAtEnd :: Measured a => Int -> MSeq a -> (MSeq a, MSeq a) +splitAtEnd i s = splitAt (length s - i) s +{-# INLINABLE splitAtEnd #-} + +----------- +-- Filter +----------- + +-- | \(O(n)\). Keep elements that satisfy a predicate. +filter :: Measured a => (a -> Bool) -> MSeq a -> MSeq a +filter = + (coerce :: ((a -> Identity Bool) -> MSeq a -> Identity (MSeq a)) + -> (a -> Bool) -> MSeq a -> MSeq a) + filterA +{-# INLINE filter #-} + +-- | \(O(n)\). Map over elements and collect the @Just@s. +mapMaybe :: Measured b => (a -> Maybe b) -> MSeq a -> MSeq b +mapMaybe = + (coerce :: ((a -> Identity (Maybe b)) -> MSeq a -> Identity (MSeq b)) + -> (a -> Maybe b) -> MSeq a -> MSeq b) + mapMaybeA +{-# INLINE mapMaybe #-} + +-- | \(O(n)\). Map over elements and split the @Left@s and @Right@s. +mapEither + :: (Measured b, Measured c) => (a -> Either b c) -> MSeq a -> (MSeq b, MSeq c) +mapEither = + (coerce :: ((a -> Identity (Either b c)) -> MSeq a -> Identity (MSeq b, MSeq c)) + -> (a -> Either b c) -> MSeq a -> (MSeq b, MSeq c)) + mapEitherA +{-# INLINE mapEither #-} + +-- | \(O(n)\). Keep elements that satisfy an applicative predicate. +filterA :: (Measured a, Applicative f) => (a -> f Bool) -> MSeq a -> f (MSeq a) +filterA f = mapMaybeA (\x -> fmap (\b -> if b then Just x else Nothing) (f x)) +{-# INLINE filterA #-} + +-- | \(O(n)\). Traverse over elements and collect the @Just@s. +mapMaybeA + :: (Measured b, Applicative f) => (a -> f (Maybe b)) -> MSeq a -> f (MSeq b) +mapMaybeA f = \case + MTree x xs -> Ap.liftA2 (maybe fromMTree MTree) (f x) (T.mapMaybeA f xs) + MEmpty -> pure MEmpty +{-# INLINE mapMaybeA #-} + +-- | \(O(n)\). Traverse over elements and split the @Left@s and @Right@s. +mapEitherA + :: (Measured b, Measured c, Applicative f) + => (a -> f (Either b c)) -> MSeq a -> f (MSeq b, MSeq c) +mapEitherA f = \case + MTree x xs -> (\g -> Ap.liftA2 g (f x) (T.mapEitherA f xs)) $ \mx xs' -> + case mx of + Left x' -> unS2 $ bimap (MTree x') fromMTree xs' + Right x' -> unS2 $ bimap fromMTree (MTree x') xs' + MEmpty -> pure (MEmpty, MEmpty) + where + unS2 (U.S2 x y) = (x, y) +{-# INLINE mapEitherA #-} + +-- | \(O(i + \log n)\). The longest prefix of elements that satisfy a predicate. +-- \(i\) is the length of the prefix. +takeWhile :: Measured a => (a -> Bool) -> MSeq a -> MSeq a +takeWhile p t = IFo.ifoldr (\i x z -> if p x then z else take i t) t t +{-# INLINE takeWhile #-} + +-- | \(O(i + \log n)\). The remainder after removing the longest prefix of +-- elements that satisfy a predicate. +-- \(i\) is the length of the prefix. +dropWhile :: Measured a => (a -> Bool) -> MSeq a -> MSeq a +dropWhile p t = IFo.ifoldr (\i x z -> if p x then z else drop i t) MEmpty t +{-# INLINE dropWhile #-} + +-- | \(O(i + \log n)\). The longest prefix of elements that satisfy a predicate, +-- together with the remainder of the sequence. +-- \(i\) is the length of the prefix. +-- +-- @span p xs == ('takeWhile' p xs, 'dropWhile' p xs)@ +span :: Measured a => (a -> Bool) -> MSeq a -> (MSeq a, MSeq a) +span p t = IFo.ifoldr (\i x z -> if p x then z else splitAt i t) (t, MEmpty) t +{-# INLINE span #-} + +-- | \(O(i + \log n)\). The longest prefix of elements that /do not/ satisfy a +-- predicate, together with the remainder of the sequence. \(i\) is the length +-- of the prefix. +-- +-- @break p == 'span' (not . p)@ +break :: Measured a => (a -> Bool) -> MSeq a -> (MSeq a, MSeq a) +break p = span (not . p) +{-# INLINE break #-} + +-- | \(O(i + \log n)\). The longest suffix of elements that satisfy a predicate. +-- \(i\) is the length of the suffix. +takeWhileEnd :: Measured a => (a -> Bool) -> MSeq a -> MSeq a +takeWhileEnd p t = IFo.ifoldl (\i z x -> if p x then z else drop (i+1) t) t t +{-# INLINE takeWhileEnd #-} + +-- | \(O(i + \log n)\). The remainder after removing the longest suffix of +-- elements that satisfy a predicate. +-- \(i\) is the length of the suffix. +dropWhileEnd :: Measured a => (a -> Bool) -> MSeq a -> MSeq a +dropWhileEnd p t = + IFo.ifoldl (\i z x -> if p x then z else take (i+1) t) MEmpty t +{-# INLINE dropWhileEnd #-} + +-- | \(O(i + \log n)\). The longest suffix of elements that satisfy a predicate, +-- together with the remainder of the sequence. +-- \(i\) is the length of the suffix. +-- +-- @spanEnd p xs == ('dropWhileEnd' p xs, 'takeWhileEnd' p xs)@ +spanEnd :: Measured a => (a -> Bool) -> MSeq a -> (MSeq a, MSeq a) +spanEnd p t = + IFo.ifoldl (\i z x -> if p x then z else splitAt (i+1) t) (MEmpty, t) t +{-# INLINE spanEnd #-} + +-- | \(O(i + \log n)\). The longest suffix of elements that /do not/ satisfy a +-- predicate, together with the remainder of the sequence. +-- \(i\) is the length of the suffix. +-- +-- @breakEnd p == 'spanEnd' (not . p)@ +breakEnd :: Measured a => (a -> Bool) -> MSeq a -> (MSeq a, MSeq a) +breakEnd p = spanEnd (not . p) +{-# INLINE breakEnd #-} + +-------------- +-- Transform +-------------- + +-- Note [Functor MSeq] +-- ~~~~~~~~~~~~~~~~~~~ +-- MSeq cannot be a Functor because of the Measured constraint on the element +-- type. So class methods which require Functor are provided as standalone. +-- +-- This problem has a decent solution in the form of the Mono* classes from the +-- mono-traversable package. I would use it here if it had not decided to +-- provide instances for all popular packages, giving it a ridiculous dependency +-- footprint of split, unordered-containers, and vector! + +-- | \(O(n)\). Map over a sequence. +map :: Measured b => (a -> b) -> MSeq a -> MSeq b +map = + (coerce :: ((a -> Identity b) -> MSeq a -> Identity (MSeq b)) + -> (a -> b) -> MSeq a -> MSeq b) + traverse +{-# INLINE map #-} + +-- | \(O(n_1 n_2)\). Cartesian product of two sequences. +liftA2 :: Measured c => (a -> b -> c) -> MSeq a -> MSeq b -> MSeq c +liftA2 f t1 t2 = case t2 of + MEmpty -> MEmpty + MTree x MTip -> map (`f` x) t1 + _ -> concatMap (\x -> map (f x) t2) t1 +{-# INLINE liftA2 #-} + +-- | \(O(n)\). Traverse a sequence. +traverse + :: (Measured b, Applicative f) => (a -> f b) -> MSeq a -> f (MSeq b) +traverse f = \case + MEmpty -> pure MEmpty + MTree x xs -> Ap.liftA2 MTree (f x) (T.traverse f xs) +{-# INLINE traverse #-} + +-- | \(O(n)\). Map over a sequence with index. +imap :: Measured b => (Int -> a -> b) -> MSeq a -> MSeq b +imap = + (coerce :: ((Int -> a -> Identity b) -> MSeq a -> Identity (MSeq b)) + -> (Int -> a -> b) -> MSeq a -> MSeq b) + itraverse +{-# INLINE imap #-} + +-- | \(O(n)\). Traverse a sequence with index. +itraverse + :: (Measured b, Applicative f) => (Int -> a -> f b) -> MSeq a -> f (MSeq b) +itraverse f = \case + MEmpty -> pure MEmpty + MTree x xs -> Ap.liftA2 MTree (f 0 x) (T.itraverse f 1 xs) +{-# INLINE itraverse #-} + +-- | \(O(n)\). Reverse a sequence. +reverse :: Measured a => MSeq a -> MSeq a +reverse (MTree x xs) = case T.uncons (rev xs) of + U.SNothing -> MTree x MTip + U.SJust (U.S2 x' xs') -> MTree x' (T.snoc xs' x) + where + rev T.MTip = T.MTip + rev (T.MBin sz _ y l r) = T.binn sz y (rev r) (rev l) +reverse MEmpty = MEmpty +{-# INLINABLE reverse #-} + +-- | \(O(n)\). Intersperse an element between the elements of a sequence. +intersperse :: Measured a => a -> MSeq a -> MSeq a +intersperse y (MTree x xs) = case T.unsnoc (go xs) of + U.SNothing -> error "intersperse: impossible" + U.SJust (U.S2 xs' _) -> MTree x xs' + where + go T.MTip = T.singleton y + go (T.MBin sz _ z l r) = T.binn (sz*2+1) z (go l) (go r) + -- No need to balance, x <= 3y => 2x+1 <= 3(2y+1) +intersperse _ MEmpty = MEmpty +{-# INLINABLE intersperse #-} + +-- | \(O(n)\). Like 'Data.Foldable.foldl'' but keeps all intermediate values. +scanl :: Measured b => (b -> a -> b) -> b -> MSeq a -> MSeq b +scanl f !z0 = + cons z0 . + flip U.evalSState z0 . + traverse (\x -> U.sState (\z -> let z' = f z x in U.S2 z' z')) +{-# INLINE scanl #-} +-- See Note [SState for scans] in Data.Seqn.Internal.Seq + +-- | \(O(n)\). Like 'Data.Foldable.foldr'' but keeps all intermediate values. +scanr :: Measured b => (a -> b -> b) -> b -> MSeq a -> MSeq b +scanr f !z0 = + flip snoc z0 . + flip U.evalSState z0 . + forwards . + traverse + (\x -> Backwards (U.sState (\z -> let z' = f x z in U.S2 z' z'))) +{-# INLINE scanr #-} + +-- | \(O(n \log n)\). Sort a sequence. +sort :: (Ord a, Measured a) => MSeq a -> MSeq a +sort = sortBy compare +{-# INLINABLE sort #-} + +-- | \(O(n \log n)\). Sort a sequence using a comparison function. +sortBy :: Measured a => (a -> a -> Ordering) -> MSeq a -> MSeq a +sortBy cmp xs = imap (\i _ -> A.indexArray xa i) xs + where + n = length xs + xa = A.createArray n errorElement $ \ma@(A.MutableArray ma#) -> do + IFo.ifoldr (\i x z -> A.writeArray ma i x *> z) (pure ()) xs + Sam.sortArrayBy cmp ma# 0 n +{-# INLINABLE sortBy #-} +-- See Note [Inlinable sortBy] in Data.Seqn.Internal.Seq + +-------------------- +-- Search and test +-------------------- + +-- | \(O(n)\). The last element satisfying a predicate. +-- +-- To get the first element, use @Data.Foldable.'Data.Foldable.find'@. +findEnd :: (a -> Bool) -> MSeq a -> Maybe a +findEnd f = + Monoid.getLast . foldMap (\x -> Monoid.Last (if f x then Just x else Nothing)) +{-# INLINE findEnd #-} + +-- | \(O(n)\). The index of the first element satisfying a predicate. +findIndex :: (a -> Bool) -> MSeq a -> Maybe Int +findIndex f = + Monoid.getFirst . + IFo.ifoldMap (\i x -> Monoid.First (if f x then Just i else Nothing)) +{-# INLINE findIndex #-} + +-- | \(O(n)\). The index of the last element satisfying a predicate. +findIndexEnd :: (a -> Bool) -> MSeq a -> Maybe Int +findIndexEnd f = + Monoid.getLast . + IFo.ifoldMap (\i x -> Monoid.Last (if f x then Just i else Nothing)) +{-# INLINE findIndexEnd #-} + +-- | \(O(n_1 + n_2)\). Indices in the second sequence where the first sequence +-- begins as a substring. Includes overlapping occurences. +infixIndices :: Eq a => MSeq a -> MSeq a -> [Int] +infixIndices t1 t2 + | null t1 = [0 .. length t2] + | compareLength t1 t2 == GT = [] + | otherwise = X.build $ \lcons lnil -> + let n1 = length t1 + t1a = infixIndicesMkArray n1 t1 + !(!mt, !mt0) = KMP.build t1a + f !i x k = \ !m -> case KMP.step mt m x of + (b,m') -> + if b + then lcons (i-n1+1) (k m') + else k m' + in IFo.ifoldr f (\ !_ -> lnil) t2 mt0 +{-# INLINE infixIndices #-} -- Inline for fusion + +infixIndicesMkArray :: Int -> MSeq a -> A.Array a +infixIndicesMkArray !n !t = A.createArray n errorElement $ \ma -> + IFo.ifoldr (\i x z -> A.writeArray ma i x *> z) (pure ()) t + +-- | \(O(\log n)\). Binary search for an element in a sequence. +-- +-- Given a function @f@ this function returns an arbitrary element @x@, if it +-- exists, such that @f x = EQ@. @f@ must be monotonic on the sequence— +-- specifically @fmap f@ must result in a sequence which has many (possibly +-- zero) @LT@s, followed by many @EQ@s, followed by many @GT@s. +binarySearchFind :: (a -> Ordering) -> MSeq a -> Maybe a +binarySearchFind f = \case + MEmpty -> Nothing + MTree x xs -> case f x of + LT -> go xs + EQ -> Just x + GT -> Nothing + where + go MTip = Nothing + go (MBin _ _ y l r) = case f y of + LT -> go r + EQ -> Just y + GT -> go l +{-# INLINE binarySearchFind #-} + +-- | \(O(\min(n_1,n_2))\). Whether the first sequence is a prefix of the second. +isPrefixOf :: Eq a => MSeq a -> MSeq a -> Bool +isPrefixOf t1 t2 = + compareLength t1 t2 /= GT && Stream.isPrefixOf (stream t1) (stream t2) +{-# INLINABLE isPrefixOf #-} + +-- | \(O(\min(n_1,n_2))\). Whether the first sequence is a suffix of the second. +isSuffixOf :: Eq a => MSeq a -> MSeq a -> Bool +isSuffixOf t1 t2 = + compareLength t1 t2 /= GT && Stream.isPrefixOf (streamEnd t1) (streamEnd t2) +{-# INLINABLE isSuffixOf #-} + +-- | \(O(n_1 + n_2)\). Whether the first sequence is a substring of the second. +isInfixOf :: Eq a => MSeq a -> MSeq a -> Bool +isInfixOf t1 t2 = not (null (infixIndices t1 t2)) +{-# INLINABLE isInfixOf #-} + +-- | \(O(n_1 + n_2)\). Whether the first sequence is a subsequence of the second. +isSubsequenceOf :: Eq a => MSeq a -> MSeq a -> Bool +isSubsequenceOf t1 t2 = + compareLength t1 t2 /= GT && Stream.isSubsequenceOf (stream t1) (stream t2) +{-# INLINABLE isSubsequenceOf #-} + +------------------ +-- Zip and unzip +------------------ + +-- | \(O(\min(n_1,n_2))\). Zip two sequences with a function. The result is +-- as long as the shorter sequence. +zipWith :: Measured c => (a -> b -> c) -> MSeq a -> MSeq b -> MSeq c +zipWith = + (coerce :: ((a -> b -> Identity c) -> MSeq a -> MSeq b -> Identity (MSeq c)) + -> (a -> b -> c) -> MSeq a -> MSeq b -> MSeq c) + zipWithM +{-# INLINE zipWith #-} + +-- | \(O(\min(n_1,n_2,n_3))\). Zip three sequences with a function. The result +-- is as long as the shortest sequence. +zipWith3 + :: Measured d => (a -> b -> c -> d) -> MSeq a -> MSeq b -> MSeq c -> MSeq d +zipWith3 = + (coerce :: ((a -> b -> c -> Identity d) -> MSeq a -> MSeq b -> MSeq c -> Identity (MSeq d)) + -> (a -> b -> c -> d) -> MSeq a -> MSeq b -> MSeq c -> MSeq d) + zipWith3M +{-# INLINE zipWith3 #-} + +-- | \(O(\min(n_1,n_2))\). Zip two sequences with a monadic function. +zipWithM + :: (Measured c, Monad m) => (a -> b -> m c) -> MSeq a -> MSeq b -> m (MSeq c) +zipWithM f t1 t2 = zipWithStreamM f t1 (stream t2) +{-# INLINE zipWithM #-} + +-- | \(O(\min(n_1,n_2,n_3))\). Zip three sequences with a monadic function. +zipWith3M + :: (Measured d, Monad m) + => (a -> b -> c -> m d) -> MSeq a -> MSeq b -> MSeq c -> m (MSeq d) +zipWith3M f t1 t2 t3 = + zipWithStreamM + (\x (U.S2 y z) -> f x y z) + t1 + (Stream.zipWith U.S2 (stream t2) (stream t3)) +{-# INLINE zipWith3M #-} + +zipWithStreamM + :: (Measured c, Monad m) + => (a -> b -> m c) -> MSeq a -> Stream b -> m (MSeq c) +zipWithStreamM f t strm = case t of + MEmpty -> pure MEmpty + MTree x xs -> case strm of + Stream step s -> case step s of + Done -> pure MEmpty + Yield y s1 -> + Ap.liftA2 MTree (f x y) (T.zipWithStreamM f xs (Stream step s1)) +{-# INLINE zipWithStreamM #-} + +-- | \(O(n)\). Map over a sequence and unzip the result. +unzipWith + :: (Measured b, Measured c) + => (a -> (b, c)) -> MSeq a -> (MSeq b, MSeq c) +unzipWith f t = case t of + MTree x xs -> + case (f x, T.unzipWithA (Identity U.#. f) xs) of + ((x1,x2), Identity (U.S2 xs1 xs2)) -> + let !t1 = MTree x1 xs1 + !t2 = MTree x2 xs2 + in (t1,t2) + MEmpty -> (MEmpty, MEmpty) +{-# INLINE unzipWith #-} + +-- | \(O(n)\). Map over a sequence and unzip the result. +unzipWith3 + :: (Measured b, Measured c, Measured d) + => (a -> (b, c, d)) -> MSeq a -> (MSeq b, MSeq c, MSeq d) +unzipWith3 f t = case t of + MTree x xs -> + case (f x, T.unzipWith3A (Identity U.#. f) xs) of + ((x1,x2,x3), Identity (U.S3 xs1 xs2 xs3)) -> + let !t1 = MTree x1 xs1 + !t2 = MTree x2 xs2 + !t3 = MTree x3 xs3 + in (t1,t2,t3) + MEmpty -> (MEmpty, MEmpty, MEmpty) +{-# INLINE unzipWith3 #-} + +--------------------- +-- Measured queries +--------------------- + +-- | \(O(1)\). The summary is the fold of measures of all elements in the +-- sequence. Returns @Nothing@ if the sequence is empty. +-- +-- @summaryMay == 'foldMap' (Just . 'measure')@ +summaryMay :: Measured a => MSeq a -> Maybe (Measure a) +summaryMay = \case + MTree x xs -> Just $! measure x T.<<> xs + MEmpty -> Nothing + +-- | \(O(1)\). The summary is the fold of measures of all elements in the +-- sequence. +-- +-- @summary == 'foldMap' 'measure'@ +summary :: (Measured a, Monoid (Measure a)) => MSeq a -> Measure a +summary = \case + MTree x xs -> measure x T.<<> xs + MEmpty -> mempty + +-- | \(O(\log n)\). Perform a binary search on the summaries of the non-empty +-- prefixes of the sequence. +-- +-- @binarySearchPrefix p xs@ for a monotonic predicate @p@ returns two adjacent +-- indices @i@ and @j@, @0 <= i < j < length xs@. +-- +-- * @i@ is the greatest index such that +-- @p (fromJust (summaryMay (take (i+1) xs)))@ +-- is @False@, or @Nothing@ if there is no such index. +-- * @j@ is the least index such that +-- @p (fromJust (summaryMay (take (j+1) xs)))@ +-- is @True@, or @Nothing@ if there is no such index. +-- +-- ==== __Examples__ +-- +-- @ +-- import "Data.Monoid" (Sum(..)) +-- +-- newtype A = A Int deriving Show +-- +-- instance Measured A where +-- type Measure A = Sum Int +-- measure (A x) = Sum x +-- @ +-- +-- >>> let xs = fromList [A 1, A 2, A 3, A 4] +-- +-- The summaries of the prefixes of @xs@ by index are: +-- +-- * @0: measure (A 1) = Sum 1@. +-- * @1: measure (A 1) <> measure (A 2) = Sum 3@. +-- * @2: measure (A 1) <> measure (A 2) <> measure (A 3) = Sum 6@. +-- * @3: measure (A 1) <> measure (A 2) <> measure (A 3) <> measure (A 4) = Sum 10@. +-- +-- >>> binarySearchPrefix (> Sum 4) xs +-- (Just 1,Just 2) +-- +-- @ +-- ╭──────────┬──────────┬──────────┬──────────╮ +-- index: │ 0 │ 1 │ 2 │ 3 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- prefix summary: │ Sum 1 │ Sum 3 │ Sum 6 | Sum 10 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- (> Sum 4): │ False │ False │ True │ True │ +-- ╰──────────┴──────────┴──────────┴──────────╯ +-- result: ( Just 1 , Just 2 ) +-- @ +-- +-- >>> binarySearchPrefix (> Sum 20) xs +-- (Just 3,Nothing) +-- +-- @ +-- ╭──────────┬──────────┬──────────┬──────────╮ +-- index: │ 0 │ 1 │ 2 │ 3 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- prefix summary: │ Sum 1 │ Sum 3 │ Sum 6 | Sum 10 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- (> Sum 20): │ False │ False │ False │ False │ +-- ╰──────────┴──────────┴──────────┴──────────╯ +-- result: ( Just 3 , Nothing ) +-- @ +-- +binarySearchPrefix + :: Measured a => (Measure a -> Bool) -> MSeq a -> (Maybe Int, Maybe Int) +binarySearchPrefix p = \case + MTree x xs + | p v -> (Nothing, Just 0) + | p (v T.<<> xs) -> let !i = go 1 v xs in (Just (i-1), Just i) + | otherwise -> let !i = T.size xs in (Just i, Nothing) + where + v = measure x + MEmpty -> (Nothing, Nothing) + where + go !i !vup = \case + MBin _ _ x l r + | p v -> go i vup l + | p v' -> i + T.size l + | otherwise -> go (i + T.size l + 1) v' r + where + v = vup T.<<> l + v' = v <> measure x + MTip -> error "MSeq.binarySearchPrefix: bad p" +{-# INLINE binarySearchPrefix #-} + +-- | \(O(\log n)\). Perform a binary search on the summaries of the non-empty +-- suffixes of the sequence. +-- +-- @binarySearchSuffix p xs@ for a monotonic predicate @p@ returns two adjacent +-- indices @i@ and @j@, @0 <= i < j < length xs@. +-- +-- * @i@ is the greatest index such that +-- @p (fromJust (summaryMay (drop i xs)))@ is +-- @True@, or @Nothing@ if there is no such index. +-- * @j@ is the least index such that +-- @p (fromJust (summaryMay (drop j xs)))@ is +-- @False@, or @Nothing@ if there is no such index +-- +-- ==== __Examples__ +-- +-- @ +-- import "Data.Monoid" (Sum(..)) +-- +-- newtype A = A Int deriving Show +-- +-- instance Measured A where +-- type Measure A = Sum Int +-- measure (A x) = Sum x +-- @ +-- +-- >>> let xs = fromList [A 1, A 2, A 3, A 4] +-- +-- The summaries of the suffixes of @xs@ by index are: +-- +-- * @0: measure (A 1) <> measure (A 2) <> measure (A 3) <> measure (A 4) = Sum 10@. +-- * @1: measure (A 2) <> measure (A 3) <> measure (A 4) = Sum 9@. +-- * @2: measure (A 3) <> measure (A 4) = Sum 7@. +-- * @3: measure (A 4) = Sum 4@. +-- +-- >>> binarySearchSuffix (> Sum 4) xs +-- (Just 2,Just 3) +-- +-- @ +-- ╭──────────┬──────────┬──────────┬──────────╮ +-- index: │ 0 │ 1 │ 2 │ 3 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- suffix summary: │ Sum 10 │ Sum 9 │ Sum 7 | Sum 4 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- (> Sum 4): │ True │ True │ True │ False │ +-- ╰──────────┴──────────┴──────────┴──────────╯ +-- result: ( Just 2 , Just 3 ) +-- @ +-- +-- >>> binarySearchSuffix (> Sum 20) xs +-- (Nothing,Just 0) +-- +-- @ +-- ╭──────────┬──────────┬──────────┬──────────╮ +-- index: │ 0 │ 1 │ 2 │ 3 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- suffix summary: │ Sum 10 │ Sum 9 │ Sum 7 | Sum 4 │ +-- ├──────────┼──────────┼──────────┼──────────┤ +-- (> Sum 20): │ False │ False │ False │ False │ +-- ╰──────────┴──────────┴──────────┴──────────╯ +-- result: ( Nothing , Just 0 ) +-- @ +-- +binarySearchSuffix + :: Measured a => (Measure a -> Bool) -> MSeq a -> (Maybe Int, Maybe Int) +binarySearchSuffix p = \case + MTree x xs -> case xs of + MBin _ rv rx rl rr + | p rv -> let !i = goR rx rl rr + in if i == 0 + then let !j = T.size xs in (Just j, Nothing) + else let !j = T.size xs - i in (Just j, Just (j+1)) + | p v -> (Just 0, Just 1) + | otherwise -> (Nothing, Just 0) + where + v = measure x <> rv + MTip + | p (measure x) -> (Just 0, Nothing) + | otherwise -> (Nothing, Just 0) + MEmpty -> (Nothing, Nothing) + where + goR !x !l r = case r of + MBin rsz rv rx rl rr + | p rv -> goR rx rl rr + | p v -> rsz + | otherwise -> go (1 + rsz) v l + where + v = measure x <> rv + MTip + | p v -> 0 + | otherwise -> go 1 v l + where + v = measure x + go !i !vup = \case + MBin _ _ x l r + | p v -> go i vup r + | p v' -> i + T.size r + | otherwise -> go (1 + T.size r + i) v' l + where + v = r T.<>> vup + v' = measure x <> v + MTip -> error "MSeq.binarySearchSuffix: bad p" +{-# INLINE binarySearchSuffix #-} + +---------- +-- Force +---------- + +-- | Reduce a sequence to normal form, given functions to reduce its contents. +liftRnf2 :: (Measure a -> ()) -> (a -> ()) -> MSeq a -> () +liftRnf2 g f = \case + MTree x xs -> f x `seq` T.liftRnf2 g f xs + MEmpty -> () +{-# INLINE liftRnf2 #-} + +-------- +-- Util +-------- + +fromMTree :: Measured a => MTree a -> MSeq a +fromMTree t = case T.uncons t of + U.SNothing -> MEmpty + U.SJust (U.S2 x xs) -> MTree x xs +{-# INLINE fromMTree #-} + +-- See Note [compareLength] +compareLength :: MSeq a -> MSeq b -> Ordering +compareLength l r = case l of + MTree _ xs -> case r of + MTree _ ys -> compareSize xs ys + MEmpty -> GT + MEmpty -> case r of + MTree _ _ -> LT + MEmpty -> EQ +{-# INLINE compareLength #-} + +compareSize :: MTree a -> MTree b -> Ordering +compareSize l r = case l of + MBin szl _ _ _ _ -> case r of + MBin szr _ _ _ _ -> compare szl szr + MTip -> GT + MTip -> case r of + MBin _ _ _ _ _ -> LT + MTip -> EQ +{-# INLINE compareSize #-} + +---------- +-- Build +---------- + +-- WARNING +-- +-- The functions below are similar but they should not be mixed together! All of +-- them operate on Stack, but what the Stack means is not the same between +-- functions. +-- +-- left-to-right, 1 element at a time: ltrPush, ltrFinish +-- left-to-right, many elements at a time: ltrPushMany, ltrFinish +-- right-to-left, 1 element at a time: rtlPush, rtlFinish + +-- See Note [fromList implementation] in Data.Seqn.Internal.Seq +-- See Note [concatMap implementation] in Data.Seqn.Internal.Seq + +ltrPush :: Measured a => Stack a -> a -> Stack a +ltrPush stk y = case stk of + Push x MTip stk' -> ltrPushLoop stk' x 1 (T.singleton y) + _ -> Push y MTip stk +{-# INLINABLE ltrPush #-} + +ltrPushLoop :: Measured a => Stack a -> a -> Int -> MTree a -> Stack a +ltrPushLoop stk y !ysz ys = case stk of + Push x xs@(MBin xsz _ _ _ _) stk' + | xsz == ysz -> ltrPushLoop stk' x sz (T.binn sz y xs ys) + where + sz = xsz+xsz+1 + _ -> Push y ys stk +{-# INLINABLE ltrPushLoop #-} + +rtlPush :: Measured a => a -> Stack a -> Stack a +rtlPush x = \case + Push y MTip stk' -> rtlPushLoop x 1 (T.singleton y) stk' + stk -> Push x MTip stk +{-# INLINABLE rtlPush #-} + +rtlPushLoop :: Measured a => a -> Int -> MTree a -> Stack a -> Stack a +rtlPushLoop x !xsz xs = \case + Push y ys@(MBin ysz _ _ _ _) stk' + | xsz == ysz -> rtlPushLoop x sz (T.binn sz y xs ys) stk' + where + sz = xsz+xsz+1 + stk -> Push x xs stk +{-# INLINABLE rtlPushLoop #-} + +ltrPushMany :: Measured a => Stack a -> a -> MTree a -> Stack a +ltrPushMany stk y ys = case stk of + Push x xs stk' + | ysz > xsz `div` 2 -> ltrPushManyLoop stk' x xsz xs y ysz ys + | otherwise -> Push y ys stk + where + xsz = 1 + T.size xs + ysz = 1 + T.size ys + Nil -> Push y ys Nil +{-# INLINABLE ltrPushMany #-} + +ltrPushManyLoop + :: Measured a + => Stack a -> a -> Int -> MTree a -> a -> Int -> MTree a -> Stack a +ltrPushManyLoop stk y !ysz ys z !zsz zs = case stk of + Push x xs@(MBin xsz1 _ _ _ _) stk' + | xsz < zsz + -> ltrPushManyLoop stk' x (xsz + ysz) (T.link y xs ys) z zsz zs + | yzsz > xsz `div` 2 + -> ltrPushManyLoop stk' x xsz xs y yzsz (T.link z ys zs) + | otherwise + -> Push y (T.link z ys zs) stk + where + xsz = 1+xsz1 + yzsz = ysz+zsz + _ -> Push y (T.link z ys zs) stk +{-# INLINABLE ltrPushManyLoop #-} + +ltrFinish :: Measured a => Stack a -> MSeq a +ltrFinish = wrapUpStack + MEmpty + U.S2 + (\(U.S2 y ys) x xs -> U.S2 x (T.link y xs ys)) + (\(U.S2 y ys) -> MTree y ys) +{-# INLINABLE ltrFinish #-} + +rtlFinish :: Measured a => Stack a -> MSeq a +rtlFinish = wrapUpStack + MEmpty + U.S2 + (\(U.S2 x xs) y ys -> U.S2 x (T.link y xs ys)) + (\(U.S2 x xs) -> MTree x xs) +{-# INLINABLE rtlFinish #-} + +----------- +-- Stream +----------- + +-- See Note [Streams] in Data.Seqn.Internal.Seq + +stream :: MSeq a -> Stream a +stream !t = Stream step s + where + s = case t of + MTree x xs -> Push x xs Nil + MEmpty -> Nil + step = \case + Nil -> Done + Push x xs stk -> let !stk' = down xs stk in Yield x stk' + {-# INLINE [0] step #-} +{-# INLINE stream #-} + +streamEnd :: MSeq a -> Stream a +streamEnd !t = Stream step s + where + s = case t of + MTree x xs -> Push x xs Nil + MEmpty -> Nil + step = \case + Nil -> Done + Push x xs stk -> case rDown x xs stk of + U.S2 y stk' -> Yield y stk' + {-# INLINE [0] step #-} +{-# INLINE streamEnd #-} + +down :: MTree a -> Stack a -> Stack a +down (MBin _ _ x l r) stk = down l (Push x r stk) +down MTip stk = stk + +rDown :: a -> MTree a -> Stack a -> U.S2 a (Stack a) +rDown !y (MBin _ _ x l r) stk = rDown x r (Push y l stk) +rDown y MTip stk = U.S2 y stk + +---------- +-- Stack +---------- + +-- This is used in various places. What it stores depends on the specific use +-- case. +data Stack a + = Push !a !(MTree a) !(Stack a) + | Nil + +wrapUpStack + :: c -- empty + -> (a -> MTree a -> b) -- initial + -> (b -> a -> MTree a -> b) -- fold fun + -> (b -> c) -- finish + -> Stack a + -> c +wrapUpStack z0 f0 f fin = go + where + go Nil = z0 + go (Push x xs stk) = go1 (f0 x xs) stk + go1 !z Nil = fin z + go1 z (Push x xs stk) = go1 (f z x xs) stk +{-# INLINE wrapUpStack #-} + +------------ +-- Testing +------------ + +valid :: (Measured a, Eq (Measure a)) => MSeq a -> Bool +valid = \case + MTree _ xs -> T.valid xs + MEmpty -> True + +debugShowsPrec :: (Show a, Show (Measure a)) => Int -> MSeq a -> ShowS +debugShowsPrec p = \case + MTree x xs -> + showParen (p > 10) $ + showString "MTree " . + showsPrec 11 x . + showString " " . + T.debugShowsPrec 11 xs + MEmpty -> showString "MEmpty" + +---------- +-- Error +---------- + +errorElement :: a +errorElement = error "MSeq: errorElement" diff --git a/src/Data/Seqn/Internal/MTree.hs b/src/Data/Seqn/Internal/MTree.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/MTree.hs @@ -0,0 +1,740 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- This is an internal module. You probably don't need to import this. Use +-- "Data.Seqn.MSeq" instead. +-- +-- = WARNING +-- +-- Definitions in this module allow violating invariants that would otherwise be +-- guaranteed by "Data.Seqn.MSeq". Use at your own risk! +-- +module Data.Seqn.Internal.MTree + ( + -- * Measured + Measured(..) + + -- * MTree + , MTree(..) + + -- * Basic + , singleton + , size + , (<>>) + , (<<>) + , bin + , binn + + -- * Folds + , foldMap + , foldl' + , foldr' + , ifoldl' + , ifoldr' + , traverse + , ifoldMap + , itraverse + + -- * Construct + , generateA + + -- * Index + , index + , adjustF + , insertAt + , deleteAt + + -- * Slice + , cons + , snoc + , uncons + , unconsSure + , unsnoc + , unsnocSure + , splitAtF + + -- * Transform + , mapMaybeA + , mapEitherA + + -- * Force + , liftRnf2 + + -- * Zip and unzip + , zipWithStreamM + , unzipWithA + , unzipWith3A + + -- * Tree helpers + , fold + , foldSimple + , link + , merge + , glue + , balanceL + , balanceR + + -- * Testing + , valid + , debugShowsPrec + ) where + +import Prelude hiding (foldMap, foldl', traverse) +import qualified Control.Applicative as Ap +import Control.DeepSeq (NFData(..)) +import Data.Bifunctor (Bifunctor(..)) +import Data.Coerce (coerce) + +import Data.Seqn.Internal.Stream (Stream(..), Step(..)) +import qualified Data.Seqn.Internal.Util as U + +------------- +-- Measured +------------- + +-- | Types that have a combinable property, called the measure. +class Semigroup (Measure a) => Measured a where + type Measure a + + -- | Calculate the measure of a value. + measure :: a -> Measure a + +---------- +-- MTree +---------- + +data MTree a + = MBin {-# UNPACK #-} !Int !(Measure a) !a !(MTree a) !(MTree a) + | MTip + +-------------- +-- Instances +-------------- + +instance (NFData (Measure a), NFData a) => NFData (MTree a) where + rnf = \case + MBin _ v x l r -> rnf v `seq` rnf x `seq` rnf l `seq` rnf r + MTip -> () + {-# INLINABLE rnf #-} + +liftRnf2 :: (Measure a -> ()) -> (a -> ()) -> MTree a -> () +liftRnf2 g f = go + where + go (MBin _ v x l r) = g v `seq` f x `seq` go l `seq` go r + go MTip = () +{-# INLINE liftRnf2 #-} + +-------------- +-- Basic ops +-------------- + +singleton :: Measured a => a -> MTree a +singleton x = MBin 1 (measure x) x MTip MTip +{-# INLINE singleton #-} + +size :: MTree a -> Int +size (MBin n _ _ _ _) = n +size MTip = 0 +{-# INLINE size #-} + +infixr 6 <>> +infixr 6 <<> + +(<>>) :: Measured a => MTree a -> Measure a -> Measure a +MBin _ v _ _ _ <>> x = v <> x +MTip <>> x = x +{-# INLINE (<>>) #-} + +(<<>) :: Measured a => Measure a -> MTree a -> Measure a +x <<> MBin _ v _ _ _ = x <> v +x <<> MTip = x +{-# INLINE (<<>) #-} + +-- O(1). Link two trees with a value in between. Precondition: The trees are +-- balanced wrt each other. +bin :: Measured a => a -> MTree a -> MTree a -> MTree a +bin x l r = MBin (size l + size r + 1) (l <>> measure x <<> r) x l r +{-# INLINE bin #-} + +-- O(1). Link two trees with a value in between and a known total size. +-- Precondition: The trees are balanced wrt each other. +binn :: Measured a => Int -> a -> MTree a -> MTree a -> MTree a +binn n x l r = MBin n (l <>> measure x <<> r) x l r +{-# INLINE binn #-} + +---------- +-- Folds +---------- + +-- See Note [Folds] in Data.Seqn.Internal.Seq + +foldl' :: (b -> a -> b) -> b -> MTree a -> b +foldl' f !z0 = \case + MBin _ _ x l r -> go z0 x l r + MTip -> z0 + where + go !z !x l r = case l of + MBin _ _ lx ll lr -> case r of + MBin _ _ rx rl rr -> + let !z' = go z lx ll lr + in go (f z' x) rx rl rr + MTip -> + let !z' = go z lx ll lr + in f z' x + MTip -> case r of + MBin _ _ rx rl rr -> go (f z x) rx rl rr + MTip -> f z x +{-# INLINE foldl' #-} + +ifoldl' :: (Int -> b -> a -> b) -> b -> Int -> MTree a -> b +ifoldl' f !z0 !i0 = \case + MBin _ _ x l r -> go z0 i0 x l r + MTip -> z0 + where + go !z !i !x l r = case l of + MBin lsz _ lx ll lr -> case r of + MBin _ _ rx rl rr -> + let !z' = go z i lx ll lr + in go (f (i+lsz) z' x) (i+lsz+1) rx rl rr + MTip -> + let !z' = go z i lx ll lr + in f (i+lsz) z' x + MTip -> case r of + MBin _ _ rx rl rr -> go (f i z x) (i+1) rx rl rr + MTip -> f i z x +{-# INLINE ifoldl' #-} + +foldr' :: (a -> b -> b) -> b -> MTree a -> b +foldr' f !z0 = \case + MBin _ _ x l r -> go z0 x l r + MTip -> z0 + where + go !z !x l r = case l of + MBin _ _ lx ll lr -> case r of + MBin _ _ rx rl rr -> + let !z' = go z rx rl rr + in go (f x z') lx ll lr + MTip -> go (f x z) lx ll lr + MTip -> case r of + MBin _ _ rx rl rr -> f x $! go z rx rl rr + MTip -> f x z +{-# INLINE foldr' #-} + +ifoldr' :: (Int -> a -> b -> b) -> b -> Int -> MTree a -> b +ifoldr' f !z0 !i0 = \case + MBin _ _ x l r -> go z0 i0 x l r + MTip -> z0 + where + go !z !i !x l r = case l of + MBin _ _ lx ll lr -> case r of + MBin rsz _ rx rl rr -> + let !z' = go z i rx rl rr + in go (f (i-rsz) x z') (i-rsz-1) lx ll lr + MTip -> go (f i x z) (i-1) lx ll lr + MTip -> case r of + MBin rsz _ rx rl rr -> f (i-rsz) x $! go z i rx rl rr + MTip -> f i x z +{-# INLINE ifoldr' #-} + +fold + :: b + -> (Int -> a -> b -> b -> b) + -> (Int -> a -> b -> b) + -> (Int -> a -> b -> b) + -> (a -> b) + -> MTree a + -> b +fold tip glr gl gr g = \case + MBin sz _ x l r -> go sz x l r + MTip -> tip + where + go !sz !x l r = case l of + MBin lsz _ lx ll lr -> case r of + MBin rsz _ rx rl rr -> glr sz x (go lsz lx ll lr) (go rsz rx rl rr) + MTip -> gl sz x (go lsz lx ll lr) + MTip -> case r of + MBin rsz _ rx rl rr -> gr sz x (go rsz rx rl rr) + MTip -> g x +{-# INLINE fold #-} + +foldSimple :: b -> (Int -> a -> b -> b -> b) -> MTree a -> b +foldSimple tip f = fold tip f gl gr g + where + gl !sz x ml = f sz x ml tip + {-# INLINE gl #-} + gr !sz x mr = f sz x tip mr + {-# INLINE gr #-} + g x = f 1 x tip tip + {-# INLINE g #-} +{-# INLINE foldSimple #-} + +foldMap :: forall a m. Monoid m => (a -> m) -> MTree a -> m +foldMap f = coerce (fold @m @a) (mempty @m) glr gl gr f + where + glr (_ :: Int) x l r = l <> f x <> r + {-# INLINE glr #-} + gl (_ :: Int) x l = l <> f x + {-# INLINE gl #-} + gr (_ :: Int) x r = f x <> r + {-# INLINE gr #-} +{-# INLINE foldMap #-} + +traverse :: (Measured b, Applicative f) => (a -> f b) -> MTree a -> f (MTree b) +traverse f = fold (pure MTip) glr gl gr g + where + glr !sz x ml mr = Ap.liftA3 (flip (binn sz)) ml (f x) mr + {-# INLINE glr #-} + gl !sz x ml = Ap.liftA2 (\l' x' -> binn sz x' l' MTip) ml (f x) + {-# INLINE gl #-} + gr !sz x mr = Ap.liftA2 (\x' r' -> binn sz x' MTip r') (f x) mr + {-# INLINE gr #-} + g x = fmap singleton (f x) + {-# INLINE g #-} +{-# INLINE traverse #-} + +ifoldMap :: Monoid m => (Int -> a -> m) -> Int -> MTree a -> m +ifoldMap f !i0 = \case + MBin _ _ x l r -> go i0 x l r + MTip -> mempty + where + go !i x l r = case l of + MBin lsz _ lx ll lr -> case r of + MBin _ _ rx rl rr -> + go i lx ll lr <> f (i+lsz) x <> go (i+lsz+1) rx rl rr + MTip -> go i lx ll lr <> f (i+lsz) x + MTip -> case r of + MBin _ _ rx rl rr -> f i x <> go (i+1) rx rl rr + MTip -> f i x +{-# INLINE ifoldMap #-} + +itraverse + :: (Measured b, Applicative f) + => (Int -> a -> f b) -> Int -> MTree a -> f (MTree b) +itraverse f !i0 = \case + MBin sz _ x l r -> go i0 sz x l r + MTip -> pure MTip + where + go !i !sz x l r = case l of + MBin lsz _ lx ll lr -> case r of + MBin rsz _ rx rl rr -> + Ap.liftA3 + (flip (binn sz)) + (go i lsz lx ll lr) + (f (i+lsz) x) + (go (i+lsz+1) rsz rx rl rr) + MTip -> + Ap.liftA2 + (\l' x' -> binn sz x' l' MTip) + (go i lsz lx ll lr) + (f (i+lsz) x) + MTip -> case r of + MBin rsz _ rx rl rr -> + Ap.liftA2 + (\x' r' -> binn sz x' MTip r') + (f i x) + (go (i+1) rsz rx rl rr) + MTip -> + fmap singleton (f i x) +{-# INLINE itraverse #-} + +----------------- +-- Construction +----------------- + +generateA + :: (Measured a, Applicative f) + => (Int -> f a) -> Int -> Int -> f (MTree a) +generateA f = go + where + go !i n + | n <= 0 = pure MTip + | otherwise = + Ap.liftA3 + (flip (binn n)) + (go i lsz) + (f (i+lsz)) + (go (i+lsz+1) (n-lsz-1)) + where + lsz = (n-1) `div` 2 +{-# INLINE generateA #-} + +------------- +-- Indexing +------------- + +-- Precondition: 0 <= i < size xs +index :: Int -> MTree a -> a +index !i = \case + MBin _ _ x l r -> case compare i szl of + LT -> index i l + EQ -> x + GT -> index (i-szl-1) r + where + szl = size l + MTip -> errorOutOfBounds "MTree.index" + +-- Precondition: 0 <= i < size xs +adjustF + :: (Measured a, Functor f) + => (a -> f a) -> Int -> MTree a -> f (MTree a) +adjustF f = go + where + go !i = \case + MBin sz _ x l r -> case compare i szl of + LT -> fmap (\l' -> binn sz x l' r) (go i l) + EQ -> fmap (\x' -> binn sz x' l r) (f x) + GT -> fmap (binn sz x l) (go (i-szl-1) r) + where + szl = size l + MTip -> errorOutOfBounds "MTree.adjustF" +{-# INLINE adjustF #-} + +-- Inserts at ends if not in bounds +insertAt :: Measured a => Int -> a -> MTree a -> MTree a +insertAt !i x (MBin _ _ y l r) + | i <= szl = balanceL y (insertAt i x l) r + | otherwise = balanceR y l (insertAt (i-szl-1) x r) + where + szl = size l +insertAt _ x MTip = singleton x +{-# INLINABLE insertAt #-} + +-- Precondition: 0 <= i < size xs +deleteAt :: Measured a => Int -> MTree a -> MTree a +deleteAt !i (MBin _ _ x l r) = case compare i szl of + LT -> balanceR x (deleteAt i l) r + EQ -> glue l r + GT -> balanceL x l (deleteAt (i-szl-1) r) + where + szl = size l +deleteAt _ MTip = errorOutOfBounds "MTree.deleteAt" +{-# INLINABLE deleteAt #-} + +---------- +-- Slice +---------- + +cons :: Measured a => a -> MTree a -> MTree a +cons x MTip = singleton x +cons x (MBin _ _ y l r) = balanceL y (cons x l) r +{-# INLINABLE cons #-} + +snoc :: Measured a => MTree a -> a -> MTree a +snoc MTip x = singleton x +snoc (MBin _ _ y l r) x = balanceR y l (snoc r x) +{-# INLINABLE snoc #-} + +uncons :: Measured a => MTree a -> U.SMaybe (U.S2 a (MTree a)) +uncons (MBin _ _ x l r) = U.SJust (unconsSure x l r) +uncons MTip = U.SNothing +{-# INLINE uncons #-} + +unconsSure :: Measured a => a -> MTree a -> MTree a -> U.S2 a (MTree a) +unconsSure x (MBin _ _ lx ll lr) r = case unconsSure lx ll lr of + U.S2 y l' -> U.S2 y (balanceR x l' r) +unconsSure x MTip r = U.S2 x r +{-# INLINABLE unconsSure #-} + +unsnoc :: Measured a => MTree a -> U.SMaybe (U.S2 (MTree a) a) +unsnoc (MBin _ _ x l r) = U.SJust $ unsnocSure x l r +unsnoc MTip = U.SNothing +{-# INLINE unsnoc #-} + +unsnocSure :: Measured a => a -> MTree a -> MTree a -> U.S2 (MTree a) a +unsnocSure x l (MBin _ _ rx rl rr) = case unsnocSure rx rl rr of + U.S2 r' y -> U.S2 (balanceL x l r') y +unsnocSure x l MTip = U.S2 l x +{-# INLINABLE unsnocSure #-} + +-- Precondition: 0 <= i < size xs +splitAtF + :: (Measured a, U.Biapplicative f) + => Int -> MTree a -> f (MTree a) (U.S2 a (MTree a)) +splitAtF = go + where + go !i (MBin _ _ x l r) = case compare i szl of + LT -> second (second (\lr -> link x lr r)) (go i l) + EQ -> U.bipure l (U.S2 x r) + GT -> first (link x l) (go (i-szl-1) r) + where + szl = size l + go _ MTip = errorOutOfBounds "MTree.splitAtF" +{-# INLINE splitAtF #-} + +-------------- +-- Transform +-------------- + +mapMaybeA + :: (Applicative f, Measured b) + => (a -> f (Maybe b)) -> MTree a -> f (MTree b) +mapMaybeA f = foldSimple tip g + where + tip = pure MTip + {-# INLINE tip #-} + g _ x ml mr = (\h -> Ap.liftA3 h ml (f x) mr) $ \l my r -> + case my of + Nothing -> merge l r + Just y -> link y l r + {-# INLINE g #-} +{-# INLINE mapMaybeA #-} + +mapEitherA + :: (Applicative f, Measured b, Measured c) + => (a -> f (Either b c)) -> MTree a -> f (U.S2 (MTree b) (MTree c)) +mapEitherA f = foldSimple tip g + where + tip = pure (U.bipure MTip MTip) + {-# INLINE tip #-} + g _ x ml mr = (\h -> Ap.liftA3 h ml (f x) mr) $ \l my r -> + case my of + Left y -> U.biliftA2 (link y) merge l r + Right y -> U.biliftA2 merge (link y) l r + {-# INLINE g #-} +{-# INLINE mapEitherA #-} + +------------------ +-- Zip and unzip +------------------ + +zipWithStreamM + :: (Measured c, Monad m) + => (a -> b -> m c) -> MTree a -> Stream b -> m (MTree c) +zipWithStreamM f t (Stream step s) = U.evalSStateT (foldSimple tip g t) s + where + tip = pure MTip + {-# INLINE tip #-} + g _ x ml mr = U.SStateT $ \s2 -> do + U.S2 s3 l <- U.runSStateT ml s2 + case step s3 of + Done -> pure $ U.S2 s3 l + Yield y s4 -> do + z <- f x y + U.S2 s5 r <- U.runSStateT mr s4 + pure $! U.S2 s5 (link z l r) + {-# INLINE g #-} +{-# INLINE zipWithStreamM #-} + +unzipWithA + :: (Measured b, Measured c, Applicative f) + => (a -> f (b, c)) -> MTree a -> f (U.S2 (MTree b) (MTree c)) +unzipWithA f = foldSimple tip g + where + tip = pure (U.S2 MTip MTip) + {-# INLINE tip #-} + g !sz x ml mr = Ap.liftA3 bin2 ml (f x) mr + where + bin2 (U.S2 l1 l2) (x1,x2) (U.S2 r1 r2) = + U.S2 (binn sz x1 l1 r1) (binn sz x2 l2 r2) + {-# INLINE g #-} +{-# INLINE unzipWithA #-} + +unzipWith3A + :: (Measured b, Measured c, Measured d, Applicative f) + => (a -> f (b, c, d)) + -> MTree a + -> f (U.S3 (MTree b) (MTree c) (MTree d)) +unzipWith3A f = foldSimple tip g + where + tip = pure (U.S3 MTip MTip MTip) + {-# INLINE tip #-} + g !sz x ml mr = Ap.liftA3 bin3 ml (f x) mr + where + bin3 (U.S3 l1 l2 l3) (x1,x2,x3) (U.S3 r1 r2 r3) = + U.S3 (binn sz x1 l1 r1) (binn sz x2 l2 r2) (binn sz x3 l3 r3) + {-# INLINE g #-} +{-# INLINE unzipWith3A #-} + +----------- +-- Errors +----------- + +errorOutOfBounds :: String -> a +errorOutOfBounds name = error (name ++ ": out of bounds") + +------------ +-- Balance +------------ + +-- O(|log n1 - log n2|). Link two trees with a value in between. +link :: Measured a => a -> MTree a -> MTree a -> MTree a +link !x MTip r = cons x r +link x l MTip = snoc l x +link x l@(MBin ls lv lx ll lr) r@(MBin rs rv rx rl rr) + | delta*ls < rs = balanceL rx (linkL x ls l rl) rr + | delta*rs < ls = balanceR lx ll (linkR x lr rs r) + | otherwise = MBin (1+ls+rs) (lv <> measure x <> rv) x l r +{-# INLINE link #-} + +linkL :: Measured a => a -> Int -> MTree a -> MTree a -> MTree a +linkL !x !ls !l r = case r of + MBin rs rv rx rl rr + | delta*ls < rs -> balanceL rx (linkL x ls l rl) rr + | otherwise -> MBin (1+ls+rs) (l <>> measure x <> rv) x l r + MTip -> error "MTree.linkL: impossible" +{-# INLINABLE linkL #-} + +linkR :: Measured a => a -> MTree a -> Int -> MTree a -> MTree a +linkR !x l !rs !r = case l of + MBin ls lv lx ll lr + | delta*rs < ls -> balanceR lx ll (linkR x lr rs r) + | otherwise -> MBin (1+ls+rs) (lv <> measure x <<> r) x l r + MTip -> error "MTree.linkR: impossible" +{-# INLINABLE linkR #-} + +-- O(log (n1 + n2)). Link two trees. +merge :: Measured a => MTree a -> MTree a -> MTree a +merge MTip r = r +merge l MTip = l +merge l@(MBin ls _ lx ll lr) r@(MBin rs _ rx rl rr) + | ls < rs = case unsnocSure lx ll lr of U.S2 l' mx -> link mx l' r + | otherwise = case unconsSure rx rl rr of U.S2 mx r' -> link mx l r' +{-# INLINE merge #-} + +-- O(log (n1 + n2)). Link two trees. Precondition: The trees must be balanced +-- wrt each other. +glue :: Measured a => MTree a -> MTree a -> MTree a +glue MTip r = r +glue l MTip = l +glue l@(MBin ls _ lx ll lr) r@(MBin rs _ rx rl rr) + | ls > rs = case unsnocSure lx ll lr of U.S2 l' m -> balanceR m l' r + | otherwise = case unconsSure rx rl rr of U.S2 m r' -> balanceL m l r' +{-# INLINE glue #-} + +-- See Note [Balance] in Data.Seqn.Internal.Tree +delta, ratio :: Int +delta = 3 +ratio = 2 + +-- O(1). Restores balance with at most one right rotation. Precondition: One +-- right rotation must be enough to restore balance. This is the case when the +-- left tree might have been inserted to or the right tree deleted from. +balanceL :: Measured a => a -> MTree a -> MTree a -> MTree a +balanceL !x l r = case r of + MTip -> case l of + MTip -> MBin 1 v x MTip MTip + MBin _ lv lx ll lr -> case lr of + MTip -> case ll of + MTip -> MBin 2 (lv <> v) x l MTip + MBin _ _ _ _ _ -> + MBin 3 (lv <> v) lx ll (MBin 1 v x MTip MTip) + MBin _ lrv lrx _ _ -> case ll of + MTip -> + MBin 3 + (lv <> v) + lrx + (MBin 1 (measure lx) lx MTip MTip) + (MBin 1 v x MTip MTip) + MBin _ _ _ _ _ -> + MBin 4 (lv <> v) lx ll (MBin 2 (lrv <> measure x) x lr MTip) + MBin rs rv _ _ _ -> case l of + MTip -> MBin (1+rs) (v <> rv) x MTip r + MBin ls lv lx ll lr + | ls > delta*rs -> case (ll, lr) of + (MBin lls llv _ _ _, MBin lrs lrv lrx lrl lrr) + | lrs < ratio*lls -> + MBin (1+ls+rs) + (lv <> v <> rv) + lx + ll + (MBin (1+rs+lrs) (lrv <> v <> rv) x lr r) + | otherwise -> + MBin (1+ls+rs) + (lv <> v <> rv) + lrx + (MBin (1+lls+size lrl) (llv <> measure lx <<> lrl) lx ll lrl) + (MBin (1+rs+size lrr) (lrr <>> v <> rv) x lrr r) + _ -> error "MTree.balanceL: impossible" + | otherwise -> MBin (1+ls+rs) (lv <> v <> rv) x l r + where + v = measure x +{-# INLINABLE balanceL #-} + +-- O(1). Restores balance with at most one left rotation. Precondition: One left +-- rotation must be enough to restore balance. This is the case when the right +-- tree might have been inserted to or the left tree deleted from. +balanceR :: Measured a => a -> MTree a -> MTree a -> MTree a +balanceR !x l r = case l of + MTip -> case r of + MTip -> MBin 1 v x MTip MTip + MBin _ rv rx rl rr -> case rl of + MTip -> case rr of + MTip -> MBin 2 (v <> rv) x MTip r + MBin _ _ _ _ _ -> MBin 3 (v <> rv) rx (MBin 1 v x MTip MTip) rr + MBin _ rlv rlx _ _ -> case rr of + MTip -> + MBin 3 + (v <> rv) + rlx + (MBin 1 v x MTip MTip) + (MBin 1 (measure rx) rx MTip MTip) + MBin _ _ _ _ _ -> + MBin 4 (v <> rv) rx (MBin 2 (v <> rlv) x MTip rl) rr + MBin ls lv _ _ _ -> case r of + MTip -> MBin (1+ls) (lv <> v) x l MTip + MBin rs rv rx rl rr + | rs > delta*ls -> case (rl, rr) of + (MBin rls rlv rlx rll rlr, MBin rrs rrv _ _ _) + | rls < ratio*rrs -> + MBin (1+ls+rs) + (lv <> v <> rv) + rx + (MBin (1+ls+rls) (lv <> v <> rlv) x l rl) + rr + | otherwise -> + MBin (1+ls+rs) + (lv <> v <> rv) + rlx + (MBin (1+ls+size rll) (lv <> v <<> rll) x l rll) + (MBin (1+rrs+size rlr) (rlr <>> measure rx <> rrv) rx rlr rr) + _ -> error "MTree.balanceR: impossible" + | otherwise -> MBin (1+ls+rs) (lv <> v <> rv) x l r + where + v = measure x +{-# INLINABLE balanceR #-} + +------------ +-- Testing +------------ + +valid :: (Measured a, Eq (Measure a)) => MTree a -> Bool +valid s = balanceOk s && sizeOk s && measureOk s + where + balanceOk = \case + MBin _ _ _ l r -> ok && balanceOk l && balanceOk r + where + ok = size l + size r <= 1 || + (size l <= delta * size r && size r <= delta * size l) + MTip -> True + + sizeOk = \case + MBin sz _ _ l r -> sizeOk l && sizeOk r && size l + size r + 1 == sz + MTip -> True + + measureOk = \case + MBin _ v x l r -> + measureOk l && measureOk r && l <>> measure x <<> r == v + MTip -> True + +debugShowsPrec :: (Show a, Show (Measure a)) => Int -> MTree a -> ShowS +debugShowsPrec p = \case + MBin sz v x l r -> + showParen (p > 10) $ + showString "MBin " . + shows sz . + showString " " . + showsPrec 11 v . + showString " " . + showsPrec 11 x . + showString " " . + debugShowsPrec 11 l . + showString " " . + debugShowsPrec 11 r + MTip -> showString "MTip" diff --git a/src/Data/Seqn/Internal/PQueue.hs b/src/Data/Seqn/Internal/PQueue.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/PQueue.hs @@ -0,0 +1,328 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- This is an internal module. You probably don't need to import this. Use +-- "Data.Seqn.PQueue" instead. +-- +-- = WARNING +-- +-- Definitions in this module allow violating invariants that would otherwise be +-- guaranteed by "Data.Seqn.PQueue". Use at your own risk! +-- +module Data.Seqn.Internal.PQueue + ( + -- * PQueue + PQueue(..) + , Elem(..) + , Min(..) + , empty + , singleton + , fromList + , concatMap + , insert + , min + , minView + , toSortedList + + -- * Entry + , Entry(..) + , entryPrio + , entryValue + ) where + +import Prelude hiding (concatMap, min) +import Data.Coerce (coerce) +import Control.DeepSeq (NFData(..), NFData1(..)) +import qualified Data.Foldable as F +import qualified Data.Foldable.WithIndex as IFo +import Data.Functor.Classes (Eq1(..), Ord1(..), Show1(..)) +import qualified GHC.Exts as X + +import Data.Seqn.MSeq (Measured(..)) +import Data.Seqn.Internal.MSeq (MSeq(..)) +import qualified Data.Seqn.Internal.MSeq as MSeq +import Data.Seqn.Internal.MTree (MTree(..)) +import qualified Data.Seqn.Internal.MTree as T +import qualified Data.Seqn.Internal.Util as U + +newtype Min a = Min a + +-- Note: We do not use Data.Semigroup.Min because we need a left-biased (<>) +-- for the FIFO property. Data.Semigroup.Min simply delegates to the min +-- function of the underlying type, which is not required to be left-biased. + +instance Ord a => Semigroup (Min a) where + x@(Min x') <> y@(Min y') = if x' <= y' then x else y + {-# INLINE (<>) #-} + +instance NFData a => NFData (Min a) where + rnf = (coerce :: (a -> ()) -> Min a -> ()) rnf + {-# INLINABLE rnf #-} + +instance NFData1 Min where + liftRnf = coerce + +newtype Elem a = Elem a + deriving newtype (Eq, Ord, Show, Read, NFData) + +instance Ord a => Measured (Elem a) where + type Measure (Elem a) = Min a + measure = coerce + +-- | A minimum priority queue. +-- +-- @PQueue@ can be used as a maximum priority queue by wrapping its elements +-- with t'Data.Ord.Down'. +newtype PQueue a = PQueue (MSeq (Elem a)) + deriving newtype + ( + -- | Insertion order. + Eq + + -- | Lexicographical ordering, in insertion order. + , Ord + + -- | + -- [@(<>)@]: \(O(\left| \log n_1 - \log n_2 \right|)\). Concatenate + -- two @PQueue@s. + , Semigroup + + -- | + -- [@mempty@]: The empty queue. + , Monoid + + , Show + , Read + ) + +instance Eq1 PQueue where + liftEq = + (coerce :: ((Elem a -> Elem b -> Bool) -> MSeq (Elem a) -> MSeq (Elem b) -> Bool) + -> (a -> b -> Bool) -> PQueue a -> PQueue b -> Bool) + liftEq + {-# INLINE liftEq #-} + +instance Ord1 PQueue where + liftCompare = + (coerce :: ((Elem a -> Elem b -> Ordering) -> MSeq (Elem a) -> MSeq (Elem b) -> Ordering) + -> (a -> b -> Ordering) -> PQueue a -> PQueue b -> Ordering) + liftCompare + {-# INLINE liftCompare #-} + +instance Show1 PQueue where + liftShowsPrec _ sl _ s = sl (F.toList s) + {-# INLINE liftShowsPrec #-} + +-- | +-- [length]: \(O(1)\). +-- +-- Folds in insertion order. +instance Foldable PQueue where + foldMap = + (coerce :: ((Elem a -> m) -> MSeq (Elem a) -> m) + -> (a -> m) -> PQueue a -> m) + foldMap + {-# INLINE foldMap #-} + + foldr = + (coerce :: ((Elem a -> b -> b) -> b -> MSeq (Elem a) -> b) + -> (a -> b -> b) -> b -> PQueue a -> b) + foldr + {-# INLINE foldr #-} + + foldl' = + (coerce :: ((b -> Elem a -> b) -> b -> MSeq (Elem a) -> b) + -> (b -> a -> b) -> b -> PQueue a -> b) + F.foldl' + {-# INLINE foldl' #-} + + foldl = + (coerce :: ((b -> Elem a -> b) -> b -> MSeq (Elem a) -> b) + -> (b -> a -> b) -> b -> PQueue a -> b) + F.foldl + {-# INLINE foldl #-} + + foldr' = + (coerce :: ((Elem a -> b -> b) -> b -> MSeq (Elem a) -> b) + -> (a -> b -> b) -> b -> PQueue a -> b) + F.foldr' + {-# INLINE foldr' #-} + + null = coerce (null @MSeq) + length = coerce (length @MSeq) + +-- | Folds in insertion order. +instance IFo.FoldableWithIndex Int PQueue where + ifoldMap = + (coerce :: ((Int -> Elem a -> m) -> MSeq (Elem a) -> m) + -> (Int -> a -> m) -> PQueue a -> m) + IFo.ifoldMap + {-# INLINE ifoldMap #-} + + ifoldr = + (coerce :: ((Int -> Elem a -> b -> b) -> b -> MSeq (Elem a) -> b) + -> (Int -> a -> b -> b) -> b -> PQueue a -> b) + IFo.ifoldr + {-# INLINE ifoldr #-} + + ifoldr' = + (coerce :: ((Int -> Elem a -> b -> b) -> b -> MSeq (Elem a) -> b) + -> (Int -> a -> b -> b) -> b -> PQueue a -> b) + IFo.ifoldr' + {-# INLINE ifoldr' #-} + + ifoldl' = + (coerce :: ((Int -> b -> Elem a -> b) -> b -> MSeq (Elem a) -> b) + -> (Int -> b -> a -> b) -> b -> PQueue a -> b) + IFo.ifoldl' + {-# INLINE ifoldl' #-} + + ifoldl = + (coerce :: ((Int -> b -> Elem a -> b) -> b -> MSeq (Elem a) -> b) + -> (Int -> b -> a -> b) -> b -> PQueue a -> b) + IFo.ifoldl + {-# INLINE ifoldl #-} + +instance Ord a => X.IsList (PQueue a) where + type Item (PQueue a) = a + + fromList = fromList + {-# INLINE fromList #-} + + toList = F.toList + {-# INLINE toList #-} + +instance NFData a => NFData (PQueue a) where + rnf = (coerce :: (MSeq (Elem a) -> ()) -> PQueue a -> ()) rnf + {-# INLINABLE rnf #-} + +instance NFData1 PQueue where + liftRnf f (PQueue t) = MSeq.liftRnf2 (liftRnf f) (coerce f) t + {-# INLINE liftRnf #-} + +--------------- +-- Operations +--------------- + +-- | The empty queue. +empty :: PQueue a +empty = PQueue MSeq.empty + +-- | A singleton queue. +singleton :: a -> PQueue a +singleton = coerce MSeq.singleton + +-- | \(O(n)\). Create a queue from a list. +fromList :: Ord a => [a] -> PQueue a +fromList = coerce MSeq.fromList +{-# INLINE fromList #-} + +-- | \(O \left(\sum_i \log n_i \right)\). +-- Map over a @Foldable@ and concatenate the results. +concatMap :: (Ord b, Foldable f) => (a -> PQueue b) -> f a -> PQueue b +concatMap = + (coerce :: ((a -> MSeq (Elem b)) -> f a -> MSeq (Elem b)) + -> (a -> PQueue b) -> f a -> PQueue b) + MSeq.concatMap +{-# INLINE concatMap #-} + +-- | \(O(\log n)\). Insert an element into the queue. +-- +-- Note: When inserting multiple elements, it is more efficient to concatenate +-- a fresh queue rather than repeatedly insert elements. +-- +-- @ +-- q <> fromList xs -- Good +-- foldl' (flip insert) q xs -- Worse +-- @ +insert :: Ord a => a -> PQueue a -> PQueue a +insert = coerce (flip MSeq.snoc) +{-# INLINABLE insert #-} + +-- | \(O(1)\). The minimum element in the queue. +min :: Ord a => PQueue a -> Maybe a +min = + (coerce :: (MSeq (Elem a) -> Maybe (Min a)) -> PQueue a -> Maybe a) + MSeq.summaryMay +{-# INLINE min #-} + +-- | \(O(\log n)\). The minimum element in the queue, with the rest of the +-- queue. +minView :: Ord a => PQueue a -> Maybe (a, PQueue a) +minView (PQueue t) = case t of + MEmpty -> Nothing + MTree x xs -> case minViewSure x xs of + U.S2 y t' -> Just (y, PQueue t') +{-# INLINE minView #-} + +minViewSure :: Ord a => Elem a -> MTree (Elem a) -> U.S2 a (MSeq (Elem a)) +minViewSure x@(Elem !x1) xs = case xs of + MTip -> U.S2 x1 MSeq.empty + MBin _ (Min v) _ _ _ + | x1 <= v -> U.S2 x1 (MSeq.fromMTree xs) + | otherwise -> U.S2 v (MTree x (deleteSure v xs)) +{-# INLINABLE minViewSure #-} + +deleteSure :: Ord a => a -> MTree (Elem a) -> MTree (Elem a) +deleteSure !k = \case + MBin _ _ x@(Elem x1) l r -> case l of + MTip + | x1 <= k -> r + | otherwise -> T.cons x (deleteSure k r) + MBin _ (Min v) _ _ _ + | v <= k -> T.balanceR x (deleteSure k l) r + | x1 <= k -> T.glue l r + | otherwise -> T.balanceL x l (deleteSure k r) + MTip -> error "PQueue.deleteSure: impossible" +{-# INLINABLE deleteSure #-} + +-- | \(O(n \log n)\). Convert to a sorted list. +toSortedList :: Ord a => PQueue a -> [a] +toSortedList q0 = X.build $ \lcons lnil -> + let go q = case minView q of + Nothing -> lnil + Just (x,q') -> lcons x (go q') + in go q0 +{-# INLINE toSortedList #-} + +---------- +-- Entry +---------- + +-- | A priority associated with a value. A @PQueue (Entry k a)@ may be used +-- when the priority is separate from the value. +data Entry k a = Entry !k a + deriving (Show, Read, Functor) + +-- | Compares by @k@ only. +instance Eq k => Eq (Entry k a) where + Entry k1 _ == Entry k2 _ = k1 == k2 + {-# INLINABLE (==) #-} + +-- | Compares by @k@ only. +instance Ord k => Ord (Entry k a) where + compare (Entry k1 _) (Entry k2 _) = compare k1 k2 + {-# INLINABLE compare #-} + + Entry k1 _ <= Entry k2 _ = k1 <= k2 + {-# INLINABLE (<=) #-} + +instance (NFData k, NFData a) => NFData (Entry k a) where + rnf (Entry k x) = rnf k `seq` rnf x + {-# INLINABLE rnf #-} + +-- | The priority. +entryPrio :: Entry k a -> k +entryPrio (Entry k _) = k + +-- | The value. +entryValue :: Entry k a -> a +entryValue (Entry _ x) = x diff --git a/src/Data/Seqn/Internal/Seq.hs b/src/Data/Seqn/Internal/Seq.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/Seq.hs @@ -0,0 +1,1765 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- This is an internal module. You probably don't need to import this. Use +-- "Data.Seqn.Seq" instead. +-- +-- = WARNING +-- +-- Definitions in this module allow violating invariants that would otherwise be +-- guaranteed by "Data.Seqn.Seq". Use at your own risk! +-- +module Data.Seqn.Internal.Seq + ( + -- * Seq + Seq(..) + + -- * Construct + , empty + , singleton + , fromList + , fromRevList + , replicate + , replicateA + , generate + , generateA + , unfoldr + , unfoldl + , unfoldrM + , unfoldlM + , concatMap + + -- * Convert + , toRevList + + -- * Index + , lookup + , index + , (!?) + , (!) + , update + , adjust + , insertAt + , deleteAt + + -- * Slice + , cons + , snoc + , uncons + , unsnoc + , take + , drop + , slice + , splitAt + , takeEnd + , dropEnd + , splitAtEnd + , tails + , inits + , chunksOf + + -- * Filter + , filter + , catMaybes + , mapMaybe + , mapEither + , filterA + , mapMaybeA + , mapEitherA + , takeWhile + , dropWhile + , span + , break + , takeWhileEnd + , dropWhileEnd + , spanEnd + , breakEnd + + -- * Transform + , reverse + , intersperse + , scanl + , scanr + , sort + , sortBy + + -- * Search and test + , findEnd + , findIndex + , findIndexEnd + , infixIndices + , binarySearchFind + , isPrefixOf + , isSuffixOf + , isInfixOf + , isSubsequenceOf + + -- * Zip and unzip + , zip + , zip3 + , zipWith + , zipWith3 + , zipWithM + , zipWith3M + , unzip + , unzip3 + , unzipWith + , unzipWith3 + + -- * Internal + , fromTree + + -- * Testing + , valid + , debugShowsPrec + ) where + +import Prelude hiding (concatMap, break, drop, dropWhile, filter, lookup, replicate, reverse, scanl, scanr, span, splitAt, take, takeWhile, traverse, unzip, unzip3, zip, zip3, zipWith, zipWith3) +import qualified Control.Applicative as Ap +import Control.Applicative.Backwards (Backwards(..)) +import Control.DeepSeq (NFData(..), NFData1(..)) +import Control.Monad (MonadPlus(..)) +import Control.Monad.Fix (MonadFix(..)) +import Control.Monad.Zip (MonadZip(..)) +import Data.Bifunctor (Bifunctor(..)) +import Data.Coerce (coerce) +import qualified Data.Foldable as F +import qualified Data.Foldable.WithIndex as IFo +import Data.Functor.Classes (Eq1(..), Ord1(..), Show1(..), Read1(..)) +import qualified Data.Functor.Classes as F1 +import Data.Functor.Const (Const(..)) +import Data.Functor.Identity (Identity(..)) +import qualified Data.Functor.WithIndex as IFu +import Data.List.NonEmpty (NonEmpty(..)) +import qualified Data.Monoid as Monoid +import qualified Data.Primitive.Array as A +import qualified Data.SamSort as Sam +import Data.Semigroup (Semigroup(..)) +import Data.String (IsString(..)) +import qualified Data.Traversable as Tr +import qualified Data.Traversable.WithIndex as ITr +import qualified GHC.Exts as X +import Text.Read (Read(..)) +import qualified Text.Read as Read + +import qualified Data.Seqn.Internal.KMP as KMP +import Data.Seqn.Internal.Stream (Stream(..), Step(..)) +import qualified Data.Seqn.Internal.Stream as Stream +import Data.Seqn.Internal.Tree (Tree(..)) +import qualified Data.Seqn.Internal.Tree as T +import qualified Data.Seqn.Internal.Util as U + +-------- +-- Seq +-------- + +-- | A sequence with elements of type @a@. +data Seq a + = Tree !a !(Tree a) + | Empty + +-- Note [Seq structure] +-- ~~~~~~~~~~~~~~~~~~~~ +-- A Seq is a weight-balanced binary tree, with a small twist: the first element +-- is kept aside from the tree. It can be viewed as a binary tree with a root +-- and a right child, but a missing left child. The motivation for this change +-- is that it improves the complexity of the append operation, from +-- O(log (n_1 + n_2)) to O(|log n_1 - log n_2|), while not affecting any of the +-- other operations. Is it worth the trouble? I think so. + +-------------- +-- Instances +-------------- + +instance Eq a => Eq (Seq a) where + t1 == t2 = compareLength t1 t2 == EQ && stream t1 == stream t2 + {-# INLINABLE (==) #-} + +-- | Lexicographical ordering +instance Ord a => Ord (Seq a) where + compare t1 t2 = compare (stream t1) (stream t2) + {-# INLINABLE compare #-} + +instance Show a => Show (Seq a) where + showsPrec _ t = shows (F.toList t) + {-# INLINABLE showsPrec #-} + +instance Read a => Read (Seq a) where + readPrec = fmap fromList readListPrec + {-# INLINABLE readPrec #-} + + readListPrec = Read.readListPrecDefault + {-# INLINABLE readListPrec #-} + +instance Eq1 Seq where + liftEq f t1 t2 = compareLength t1 t2 == EQ && liftEq f (stream t1) (stream t2) + {-# INLINE liftEq #-} + +instance Ord1 Seq where + liftCompare f t1 t2 = liftCompare f (stream t1) (stream t2) + {-# INLINE liftCompare #-} + +instance Show1 Seq where + liftShowsPrec _ sl _ t = sl (F.toList t) + {-# INLINE liftShowsPrec #-} + +instance Read1 Seq where + liftReadPrec _ = fmap fromList + liftReadListPrec = F1.liftReadListPrecDefault + +-- | +-- [@length@]: \(O(1)\). +-- +-- Folds are \(O(n)\). +instance Foldable Seq where + fold = foldMap id + {-# INLINABLE fold #-} + + foldMap f = Tr.foldMapDefault f + {-# INLINE foldMap #-} + + foldMap' f = F.foldl' (\z x -> z <> f x) mempty + {-# INLINE foldMap' #-} + + foldr f z = Stream.foldr f z . stream + {-# INLINE foldr #-} + + foldl f z = Stream.foldr (flip f) z . streamEnd + {-# INLINE foldl #-} + + foldl' f !z = \case + Tree x xs -> T.foldl' f (f z x) xs + Empty -> z + {-# INLINE foldl' #-} + + foldr' f !z = \case + Tree x xs -> f x $! T.foldr' f z xs + Empty -> z + {-# INLINE foldr' #-} + + null = \case + Tree _ _ -> False + Empty -> True + + length = \case + Tree _ xs -> 1 + T.size xs + Empty -> 0 + +-- | +-- [@fmap@]: \(O(n)\). +-- +-- [@(<$)@]: \(O(\log n)\). +instance Functor Seq where + fmap f = Tr.fmapDefault f + {-# INLINE fmap #-} + + x <$ xs = replicate (length xs) x + +instance Traversable Seq where + traverse f = \case + Empty -> pure Empty + Tree x xs -> Ap.liftA2 Tree (f x) (T.traverse f xs) + {-# INLINE traverse #-} + +-- | +-- [@(<>)@]: \(O(\left| \log n_1 - \log n_2 \right|)\). Concatenates two +-- sequences. +-- +-- [@stimes@]: \(O(\log c)\). @stimes c xs@ is @xs@ repeated @c@ times. If +-- @c < 0@, 'empty' is returned. +instance Semigroup (Seq a) where + Tree x xs <> Tree y ys = Tree x (T.link y xs ys) + l <> Empty = l + Empty <> r = r + + stimes !c = \case + t@(Tree x xs) + | c <= 0 -> Empty + | fromIntegral c * toi (length t) > toi (maxBound :: Int) -> + error "Seq.stimes: result size too large" + | otherwise -> Tree x (stimesLoop (c'-1) x xs xs) + where + c' = fromIntegral c :: Int + toi :: Int -> Integer + toi = fromIntegral + Empty -> Empty + {-# INLINABLE stimes #-} + + sconcat (x:|xs) = mconcat (x:xs) + +-- Note [Complexity of stimes] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- +-- Let stimesLoop be initially called with trees (xs and acc) of size (n-1). +-- +-- stimesLoop is called O(log c) times in total, since c halves on every call. +-- At any iteration, xs is made up of initial tree bin-ed with itself multiple +-- times, and acc is made up of the some of the xs linked together. +-- All operations in stimesLoop are O(1) except for link, which takes +-- O(log(size xs) - log(size acc)). +-- +-- The cost of the ith iteration is O(1) if 2^i is in c. +-- If not, a link is done with some cost depending on xs and acc. +-- For iteration i, the size of xs is 2^i n - 1. +-- The size of acc is (sum of 2^p_j * n) - 1 for the powers of 2 p_j < i in c. +-- Let there be k powers of 2 in c, i.e. c = \sum_{i=1}^k p_i. +-- Then the total cost of the links is +-- O(\sum_{i=2}^k (\log (2^{p_i} n - 1) - \log ((\sum_{j=1}^{i-1} 2^{p_j} n) - 1)))) +-- = O(\sum_{i=2}^k (\log (2^{p_i} n) - \log (\sum_{j=1}^{i-1} 2^{p_j} n))) +-- = O(\sum_{i=2}^k (\log (2^{p_i} n) - \log (2^{p_{i-1}} n))) +-- = O(\sum_{i=2}^k (p_i + \log n - p_{i-1} - \log n)) +-- = O(\sum_{i=2}^k (p_i - p_{i-1})) +-- = O(p_k - p_1) +-- = O(\log c) + +stimesLoop :: Int -> a -> Tree a -> Tree a -> Tree a +stimesLoop c !x !xs !acc + | c <= 0 = acc + | c `mod` 2 == 0 = stimesLoop (c `div` 2) x (T.bin x xs xs) acc + | otherwise = stimesLoop (c `div` 2) x (T.bin x xs xs) (T.link x xs acc) + +-- | +-- [@mempty@]: The empty sequence. +instance Monoid (Seq a) where + mempty = Empty + + mconcat = concatMap id + {-# INLINE mconcat #-} -- Inline for fusion + +instance NFData a => NFData (Seq a) where + rnf = \case + Tree x xs -> rnf x `seq` rnf xs + Empty -> () + {-# INLINABLE rnf #-} + +instance NFData1 Seq where + liftRnf f = \case + Tree x xs -> f x `seq` liftRnf f xs + Empty -> () + {-# INLINE liftRnf #-} + +-- | +-- [@liftA2@]: \(O(n_1 n_2)\). +-- +-- [@(<*)@]: \(O(n_1 \log n_2)\). +-- +-- [@(*>)@]: \(O(\log n_1)\). +instance Applicative Seq where + pure = singleton + + liftA2 f t1 t2 = case t2 of + Empty -> Empty + Tree x Tip -> fmap (`f` x) t1 + _ -> concatMap (\x -> fmap (f x) t2) t1 + {-# INLINE liftA2 #-} + + t1 <* t2 = case t2 of + Empty -> Empty + Tree _ Tip -> t1 + _ -> concatMap (replicate (length t2)) t1 + + s1 *> s2 = stimes (length s1) s2 + +instance Ap.Alternative Seq where + empty = Empty + (<|>) = (<>) + +instance Monad Seq where + t >>= f = concatMap f t + {-# INLINE (>>=) #-} + +instance MonadPlus Seq + +instance MonadFail Seq where + fail _ = Empty + +instance MonadFix Seq where + mfix f = + IFu.imap + (\i _ -> let x = index i (f x) in x) + (f (error "Seq.mfix: f must be lazy")) + {-# INLINE mfix #-} + +instance MonadZip Seq where + mzip = zip + mzipWith = zipWith + munzip = unzip + +instance (a ~ Char) => IsString (Seq a) where + fromString = fromList + +instance X.IsList (Seq a) where + type Item (Seq a) = a + fromList = fromList + {-# INLINE fromList #-} + + toList = F.toList + {-# INLINE toList #-} + +instance IFu.FunctorWithIndex Int Seq where + imap f = ITr.imapDefault f + {-# INLINE imap #-} + +instance IFo.FoldableWithIndex Int Seq where + ifoldMap f = ITr.ifoldMapDefault f + {-# INLINE ifoldMap #-} + + ifoldr f z = Stream.ifoldr f z 0 (+1) . stream + {-# INLINE ifoldr #-} + + ifoldl f z = \t -> + Stream.ifoldr (flip . f) z (length t - 1) (subtract 1) (streamEnd t) + {-# INLINE ifoldl #-} + + ifoldr' f !z = \case + Tree x xs -> f 0 x $! T.ifoldr' f z (T.size xs) xs + Empty -> z + {-# INLINE ifoldr' #-} + + ifoldl' f !z = \case + Tree x xs -> T.ifoldl' f (f 0 z x) 1 xs + Empty -> z + {-# INLINE ifoldl' #-} + +instance ITr.TraversableWithIndex Int Seq where + itraverse f = \case + Empty -> pure Empty + Tree x xs -> Ap.liftA2 Tree (f 0 x) (T.itraverse f 1 xs) + {-# INLINE itraverse #-} + +-------------- +-- Construct +-------------- + +-- | The empty sequence. +empty :: Seq a +empty = Empty + +-- | A singleton sequence. +singleton :: a -> Seq a +singleton x = Tree x Tip + +-- | \(O(n)\). Create a @Seq@ from a list. +-- +-- ==== __Examples__ +-- +-- >>> fromList [8,1,19,11,5,12,12] +-- [8,1,19,11,5,12,12] +fromList :: [a] -> Seq a +fromList = ltrFinish . F.foldl' ltrPush Nil +{-# INLINE fromList #-} +-- See Note [fromList implementation] + +-- | \(O(n)\). Create a @Seq@ from a reversed list. +-- +-- ==== __Examples__ +-- +-- >>> fromRevList "!olleH" +-- "Hello!" +fromRevList :: [a] -> Seq a +fromRevList = rtlFinish . F.foldl' (flip rtlPush) Nil +{-# INLINE fromRevList #-} +-- See Note [fromList implementation] + +-- | \(O(\log n)\). A sequence with a repeated element. +-- If the length is negative, 'empty' is returned. +-- +-- ==== __Examples__ +-- +-- >>> replicate 3 "ha" +-- ["ha","ha","ha"] +replicate :: Int -> a -> Seq a +replicate !n x = stimes n (Tree x Tip) + +-- | \(O(n)\). Generate a sequence from a length and an applicative action. +-- If the length is negative, 'empty' is returned. +-- +-- ==== __Examples__ +-- +-- >>> import System.Random (randomIO) +-- >>> import Data.Word (Word8) +-- >>> replicateA 5 (randomIO :: IO Word8) +-- [26,134,30,58,221] +replicateA :: Applicative f => Int -> f a -> f (Seq a) +replicateA !n m = generateA n (const m) +{-# INLINABLE replicateA #-} + +-- | \(O(n)\). Generate a sequence from a length and a generator. +-- If the length is negative, 'empty' is returned. +-- +-- ==== __Examples__ +-- +-- >>> generate 4 (10*) +-- [0,10,20,30] +generate :: Int -> (Int -> a) -> Seq a +generate = + (coerce :: (Int -> (Int -> Identity a) -> Identity (Seq a)) + -> Int -> (Int -> a) -> Seq a) + generateA +{-# INLINE generate #-} + +-- | \(O(n)\). Generate a sequence from a length and an applicative generator. +-- If the length is negative, 'empty' is returned. +generateA :: Applicative f => Int -> (Int -> f a) -> f (Seq a) +generateA n f + | n <= 0 = pure Empty + | otherwise = Ap.liftA2 Tree (f 0) (T.generateA f 1 (n-1)) +{-# INLINE generateA #-} + +-- | \(O(n)\). Unfold a sequence from left to right. +-- +-- ==== __Examples__ +-- +-- >>> let f (i,a,b) = if i >= 10 then Nothing else Just (a, (i+1, b, a+b)) +-- >>> unfoldr f (0,0,1) +-- [0,1,1,2,3,5,8,13,21,34] +unfoldr :: (b -> Maybe (a, b)) -> b -> Seq a +unfoldr = + (coerce :: ((b -> Identity (Maybe (a, b))) -> b -> Identity (Seq a)) + -> (b -> Maybe (a, b)) -> b -> Seq a) + unfoldrM +{-# INLINE unfoldr #-} + +-- | \(O(n)\). Unfold a sequence monadically from left to right. +unfoldrM :: Monad m => (b -> m (Maybe (a, b))) -> b -> m (Seq a) +unfoldrM f = go Nil + where + go !b z = f z >>= \case + Nothing -> pure $! ltrFinish b + Just (x, z') -> go (ltrPush b x) z' +{-# INLINE unfoldrM #-} + +-- | \(O(n)\). Unfold a sequence from right to left. +-- +-- ==== __Examples__ +-- +-- >>> let f i = if i <= 0 then Nothing else Just (i `div` 2, i) +-- >>> unfoldl f 1024 +-- [1,2,4,8,16,32,64,128,256,512,1024] +unfoldl :: (b -> Maybe (b, a)) -> b -> Seq a +unfoldl = + (coerce :: ((b -> Identity (Maybe (b, a))) -> b -> Identity (Seq a)) + -> (b -> Maybe (b, a)) -> b -> Seq a) + unfoldlM +{-# INLINE unfoldl #-} + +-- | \(O(n)\). Unfold a sequence monadically from right to left. +unfoldlM :: Monad m => (b -> m (Maybe (b, a))) -> b -> m (Seq a) +unfoldlM f = go Nil + where + go !b z = f z >>= \case + Nothing -> pure $! rtlFinish b + Just (z', x) -> go (rtlPush x b) z' +{-# INLINE unfoldlM #-} + +-- | \(O \left(\sum_i \log n_i \right)\). +-- Map over a @Foldable@ and concatenate the results. +-- +-- ==== __Examples__ +-- +-- >>> concatMap (uncurry replicate) [(1,'H'),(1,'e'),(2,'l'),(1,'o')] +-- "Hello" +concatMap :: Foldable f => (a -> Seq b) -> f a -> Seq b +concatMap f = ltrFinish . F.foldl' g Nil + where + g b x = case f x of + Empty -> b + Tree y ys -> ltrPushMany b y ys + {-# INLINE g #-} +{-# INLINE concatMap #-} +-- See Note [concatMap implementation] + +------------ +-- Convert +------------ + +-- | \(O(n)\). Convert to a list in reverse. +-- +-- To convert to a list without reversing, use +-- @Data.Foldable.'Data.Foldable.toList'@. +-- +-- ==== __Examples__ +-- +-- >>> toRevList (fromList "!olleH") +-- "Hello!" +toRevList :: Seq a -> [a] +toRevList t = X.build $ \lcons lnil -> F.foldl (flip lcons) lnil t +{-# INLINE toRevList #-} + +---------- +-- Index +---------- + +-- Precondition: 0 <= i < size xs +index_ :: Int -> Tree a -> a +index_ !i xs = getConst (T.adjustF Const i xs) + +-- | \(O(\log n)\). Look up the element at an index. +-- +-- ==== __Examples__ +-- +-- >>> lookup 3 (fromList "haskell") +-- Just 'k' +-- >>> lookup (-1) (singleton 7) +-- Nothing +lookup :: Int -> Seq a -> Maybe a +lookup !i (Tree x xs) + | i < 0 || T.size xs < i = Nothing + | i == 0 = Just x + | otherwise = Just $! index_ (i-1) xs +lookup _ Empty = Nothing +{-# INLINE lookup #-} + +-- | \(O(\log n)\). Look up the element at an index. Calls @error@ if the index +-- is out of bounds. +-- +-- ==== __Examples__ +-- +-- >>> index 3 (fromList "haskell") +-- 'k' +-- >>> index (-1) (singleton 7) +-- *** Exception: ... +index :: Int -> Seq a -> a +index !i = \case + Tree x xs + | i == 0 -> x + | otherwise -> index_ (i-1) xs + Empty -> error "Seq.index: out of bounds" + +-- | \(O(\log n)\). Infix version of 'lookup'. +(!?) :: Seq a -> Int -> Maybe a +(!?) = flip lookup + +-- | \(O(\log n)\). Infix version of 'index'. Calls @error@ if the index is out +-- of bounds. +(!) :: Seq a -> Int -> a +(!) = flip index + +-- | \(O(\log n)\). Update an element at an index. If the index is out of +-- bounds, the sequence is returned unchanged. +-- +-- ==== __Examples__ +-- +-- >>> update 3 'b' (fromList "bird") +-- "birb" +-- >>> update 3 True (singleton False) +-- [False] +update :: Int -> a -> Seq a -> Seq a +update i x = adjust (const x) i + +-- | \(O(\log n)\). Adjust the element at an index. If the index is out of +-- bounds the sequence is returned unchanged. +-- +-- ==== __Examples__ +-- +-- >>> adjust Data.List.reverse 1 (fromList ["Hello", "ereht"]) +-- ["Hello","there"] +-- >>> adjust (*100) (-1) (singleton 7) +-- [7] +adjust :: (a -> a) -> Int -> Seq a -> Seq a +adjust f !i t = case t of + Tree x xs + | i < 0 || T.size xs < i -> t + | i == 0 -> Tree (f x) xs + | otherwise -> Tree x (runIdentity (T.adjustF (Identity U.#. f) (i-1) xs)) + Empty -> Empty +{-# INLINE adjust #-} + +-- | \(O(\log n)\). Insert an element at an index. If the index is out of +-- bounds, the element is added to the closest end of the sequence. +-- +-- ==== __Examples__ +-- +-- >>> insertAt 1 'a' (fromList "ct") +-- "cat" +-- >>> insertAt (-10) 0 (fromList [5,6,7]) +-- [0,5,6,7] +-- >>> insertAt 10 0 (fromList [5,6,7]) +-- [5,6,7,0] +insertAt :: Int -> a -> Seq a -> Seq a +insertAt !i y t = case t of + Tree x xs + | i <= 0 -> cons y t + | otherwise -> Tree x (T.insertAt (i-1) y xs) + Empty -> singleton y + +-- | \(O(\log n)\). Delete an element at an index. If the index is out of +-- bounds, the sequence is returned unchanged. +-- +-- ==== __Examples__ +-- +-- >>> deleteAt 2 (fromList "cart") +-- "cat" +-- >>> deleteAt 10 (fromList [5,6,7]) +-- [5,6,7] +deleteAt :: Int -> Seq a -> Seq a +deleteAt !i t = case t of + Tree x xs + | i < 0 || T.size xs < i -> t + | i == 0 -> fromTree xs + | otherwise -> Tree x (T.deleteAt (i-1) xs) + Empty -> Empty + +------------ +-- Slicing +------------ + +-- | \(O(\log n)\). Append a value to the beginning of a sequence. +-- +-- ==== __Examples__ +-- +-- >>> cons 1 (fromList [2,3]) +-- [1,2,3] +cons :: a -> Seq a -> Seq a +cons x (Tree y ys) = Tree x (T.cons y ys) +cons x Empty = singleton x + +-- | \(O(\log n)\). Append a value to the end of a sequence. +-- +-- ==== __Examples__ +-- +-- >>> snoc (fromList [1,2]) 3 +-- [1,2,3] +snoc :: Seq a -> a -> Seq a +snoc (Tree y ys) x = Tree y (T.snoc ys x) +snoc Empty x = singleton x + +-- | \(O(\log n)\). The head and tail of a sequence. +-- +-- ==== __Examples__ +-- +-- >>> uncons (fromList [1,2,3]) +-- Just (1,[2,3]) +-- >>> uncons empty +-- Nothing +uncons :: Seq a -> Maybe (a, Seq a) +uncons (Tree x xs) = Just . (,) x $! fromTree xs +uncons Empty = Nothing +{-# INLINE uncons #-} + +-- | \(O(\log n)\). The init and last of a sequence. +-- +-- ==== __Examples__ +-- +-- >>> unsnoc (fromList [1,2,3]) +-- Just ([1,2],3) +-- >>> unsnoc empty +-- Nothing +unsnoc :: Seq a -> Maybe (Seq a, a) +unsnoc (Tree x xs) = case T.unsnoc xs of + U.SNothing -> Just (Empty, x) + U.SJust (U.S2 ys y) -> Just (Tree x ys, y) +unsnoc Empty = Nothing +{-# INLINE unsnoc #-} + +-- | \(O(\log n)\). Take a number of elements from the beginning of a sequence. +-- +-- ==== __Examples__ +-- +-- >>> take 3 (fromList "haskell") +-- "has" +-- >>> take (-1) (fromList [1,2,3]) +-- [] +-- >>> take 10 (fromList [1,2,3]) +-- [1,2,3] +take :: Int -> Seq a -> Seq a +take !i t@(Tree x xs) + | i <= 0 = Empty + | T.size xs < i = t + | otherwise = Tree x (getConst (T.splitAtF (i-1) xs)) +take _ Empty = Empty + +-- | \(O(\log n)\). Drop a number of elements from the beginning of a sequence. +-- +-- ==== __Examples__ +-- +-- >>> drop 3 (fromList "haskell") +-- "kell" +-- >>> drop (-1) (fromList [1,2,3]) +-- [1,2,3] +-- >>> drop 10 (fromList [1,2,3]) +-- [] +drop :: Int -> Seq a -> Seq a +drop !i t@(Tree _ xs) + | i <= 0 = t + | T.size xs < i = Empty + | otherwise = case U.unTagged (T.splitAtF (i-1) xs) of + U.S2 x' xs' -> Tree x' xs' +drop _ Empty = Empty + +-- | \(O(\log n)\). Take a number of elements from the end of a sequence. +takeEnd :: Int -> Seq a -> Seq a +takeEnd n t = drop (length t - n) t + +-- | \(O(\log n)\). Drop a number of elements from the end of a sequence. +dropEnd :: Int -> Seq a -> Seq a +dropEnd n t = take (length t - n) t + +-- | \(O(\log n)\). The slice of a sequence between two indices (inclusive). +-- +-- ==== __Examples__ +-- +-- >>> slice (1,3) (fromList "haskell") +-- "ask" +-- >>> slice (-10,2) (fromList [1,2,3,4,5]) +-- [1,2,3] +-- >>> slice (2,1) (fromList [1,2,3,4,5]) +-- [] +slice :: (Int, Int) -> Seq a -> Seq a +slice (i,j) = drop i . take (j+1) + +-- | \(O(\log n)\). Split a sequence at a given index. +-- +-- @splitAt n xs == ('take' n xs, 'drop' n xs)@ +-- +-- ==== __Examples__ +-- +-- >>> splitAt 3 (fromList "haskell") +-- ("has","kell") +-- >>> splitAt (-1) (fromList [1,2,3]) +-- ([],[1,2,3]) +-- >>> splitAt 10 (fromList [1,2,3]) +-- ([1,2,3],[]) +splitAt :: Int -> Seq a -> (Seq a, Seq a) +splitAt !i t@(Tree x xs) + | i <= 0 = (Empty, t) + | T.size xs < i = (t, Empty) + | otherwise = case T.splitAtF (i-1) xs of + U.S2 xs1 (U.S2 x' xs2) -> (Tree x xs1, Tree x' xs2) +splitAt _ Empty = (Empty, Empty) + +-- | \(O(\log n)\). Split a sequence at a given index from the end. +-- +-- @splitAtEnd n xs == ('dropEnd' n xs, 'takeEnd' n xs)@ +splitAtEnd :: Int -> Seq a -> (Seq a, Seq a) +splitAtEnd i s = splitAt (length s - i) s + +-- | \(O(n \log n)\). All suffixes of a sequence, longest first. +-- +-- ==== __Examples__ +-- +-- >>> tails (fromList [1,2,3]) +-- [[1,2,3],[2,3],[3],[]] +tails :: Seq a -> Seq (Seq a) +tails t0 = cons t0 (U.evalSState (Tr.traverse f t0) t0) + where + f _ = U.sState $ \t -> case uncons t of + Nothing -> U.S2 t t -- impossible + -- Could have been error but https://gitlab.haskell.org/ghc/ghc/-/issues/24806 + Just (_,t') -> U.S2 t' t' +-- See Note [Tails implementation] + +-- | \(O(n \log n)\). All prefixes of a sequence, shortest first. +-- +-- ==== __Examples__ +-- +-- >>> inits (fromList [1,2,3]) +-- [[],[1],[1,2],[1,2,3]] +inits :: Seq a -> Seq (Seq a) +inits t0 = snoc (U.evalSState (forwards (Tr.traverse f t0)) t0) t0 + where + f _ = Backwards $ U.sState $ \t -> case unsnoc t of + Nothing -> U.S2 t t -- impossible + Just (t',_) -> U.S2 t' t' + +-- Note [Tails implementation] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- tails :: Seq a -> Seq (Seq a) +-- +-- There are many ways to implement tails (and inits), with different +-- : +-- +-- 1. (Generate or imap) with drop : +-- 2. Send down a stack and rebuild : +-- 3. (Unfold, replicateA or traverse) with uncons : +-- +-- We do 3 for because it is faster in benchmarks. It cannot be done lazily, +-- unlike 1 and 2, but that is fine because Seq is value-strict. + + +-- | \(O \left(\frac{n}{c} \log c \right)\). Split a sequence into chunks of the +-- given length @c@. If @c <= 0@, 'empty' is returned. +-- +-- ==== __Examples__ +-- +-- >>> chunksOf 3 (fromList [1..10]) +-- [[1,2,3],[4,5,6],[7,8,9],[10]] +-- >>> chunksOf 10 (fromList "hello") +-- ["hello"] +-- >>> chunksOf (-1) (singleton 7) +-- [] + +-- See Note [chunksOf complexity] +chunksOf :: Int -> Seq a -> Seq (Seq a) +chunksOf !c t@(Tree x xs) + | c <= 0 = Empty + | c == 1 = fmap singleton t + | length t <= c = singleton t + | otherwise = case chunksOf_ c 1 xs of + U.S3 l m r -> case r of + T.Tip -> Tree (Tree x l) m + _ -> Tree (Tree x l) (T.snoc m (fromTree r)) +chunksOf _ Empty = Empty + +-- Preconditions: +-- 1. c > 1 +-- 2. at least one chunk boundary passes through the tree +chunksOf_ :: Int -> Int -> Tree a -> U.S3 (Tree a) (Tree (Seq a)) (Tree a) +chunksOf_ !_ !_ Tip = error "Seq.chunksOf_: precondition violated" +chunksOf_ c off (Bin sz x l r) = case (lHasSplit, rHasSplit) of + (False, False) -> + -- Here exactly one of (lend==c) and (roff==0) is true. + -- If both are true, precondition 1 was violated. + -- If both are false, precondition 2 was violated. + -- We check roff==0 and assume the other is the complement. + case (off==0, roff==0, rend==c) of + (False, True , False) -> U.S3 (T.snoc l x) T.Tip r + (False, True , True ) -> U.S3 (T.snoc l x) (t1 (fromTree r)) T.Tip + (False, False, False) -> U.S3 l T.Tip (T.cons x r) + (False, False, True ) -> U.S3 l (t1 (Tree x r)) T.Tip + (True , False, False) -> U.S3 T.Tip (t1 (fromTree l)) (T.cons x r) + (True , False, True ) -> U.S3 T.Tip (t2 (fromTree l) (Tree x r)) T.Tip + (True , True , False) -> U.S3 T.Tip (t1 (fromTree (T.snoc l x))) r + (True , True , True ) -> + U.S3 T.Tip (t2 (fromTree (T.snoc l x)) (fromTree r)) T.Tip + (False, True) -> case chunksOf_ c roff r of + U.S3 rl rm rr -> case (off==0, lend==c) of + (False, False) -> U.S3 (T.link x l rl) rm rr + (False, True ) -> U.S3 l (T.cons (Tree x rl) rm) rr + (True , False) -> U.S3 T.Tip (T.cons (fromTree (T.link x l rl)) rm) rr + (True , True ) -> + U.S3 T.Tip (T.cons (fromTree l) (T.cons (Tree x rl) rm)) rr + (True, False) -> case chunksOf_ c off l of + U.S3 ll lm lr -> case (roff==0, rend==c) of + (False, False) -> U.S3 ll lm (T.link x lr r) + (False, True ) -> U.S3 ll (T.snoc lm (fromTree (T.link x lr r))) T.Tip + (True , False) -> U.S3 ll (T.snoc lm (fromTree (T.snoc lr x))) r + (True , True ) -> + U.S3 ll + (T.snoc (T.snoc lm (fromTree (T.snoc lr x))) (fromTree r)) + T.Tip + (True, True) -> case (chunksOf_ c off l, chunksOf_ c roff r) of + (U.S3 ll lm lr, U.S3 rl rm rr) -> + U.S3 ll (T.link (fromTree (T.link x lr rl)) lm rm) rr + where + szl = T.size l + szr = sz - szl - 1 + lend = off + szl + roff = (lend + 1) `rem` c + rend = roff + szr + lHasSplit = lend > c + rHasSplit = rend > c + t1 y = T.Bin 1 y T.Tip T.Tip + t2 y1 y2 = T.Bin 2 y1 T.Tip (T.Bin 1 y2 T.Tip T.Tip) + +-- Note [chunksOf complexity] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- The tree of size n is partitioned into ceil(n/c) chunks, each of size at +-- most c. Each chunk is a contiguous subsequence of the original tree, and +-- such chunks are balanced in O(log c) time. The finalizing step of such a +-- chunk is a link of some left and right trees with a root. Since they can be +-- of any size, bounded by the total c, this is O(log c). The left tree here +-- is the result of multiple links of (root, right child) caused by the split +-- at a chunk boundary, again with total size bounded by c. This takes +-- O(log c). For an explanation see the description of the finishing step in +-- Note [fromList complexity]. The same applies to the right tree. Hence, each +-- chunk is balanced in O(log c) and to balance all the chunks we need +-- O((n/c) log c). +-- +-- Now the result tree has size ceil(n/c), which needs to be balanced. This is +-- done by linking the recursive results from the left and right children, l and +-- r. The results are triples of +-- (left incomplete chunk, complete chunks, right incomplete chunk). +-- The number of complete chunks from the left child, say l', is at least +-- lmin=ceil((lsz-2(c-1))/c) and at most lmax=floor(lsz/c). It is likewise for +-- the right child, which returns rmin<=r'<=rmax chunks. Balancing the linked +-- tree takes O(|log(l'sz) - log(r'sz)|) +-- = O(max(log lmax - log rmin, log rmax - log lmin)) +-- = O(max(log(lmax/rmin), log(rmax/lmin))) +-- = O(max(log(lsz/rsz), log(rsz/lsz)) ; lmax is Θ(lsz/c), rmax is Θ(rsz/c) +-- = O(1) ; lsz<=3*rsz && rsz<=3*lsz by balance +-- So all the balancing work here is done in O(n/c). +-- +-- The total is dominated by balancing all the chunks, giving us O((n/c) log c). + +-------------- +-- Filtering +-------------- + +-- | \(O(n)\). Keep elements that satisfy a predicate. +-- +-- ==== __Examples__ +-- +-- >>> filter even (fromList [1..10]) +-- [2,4,6,8,10] +filter :: (a -> Bool) -> Seq a -> Seq a +filter = + (coerce :: ((a -> Identity Bool) -> Seq a -> Identity (Seq a)) + -> (a -> Bool) -> Seq a -> Seq a) + filterA +{-# INLINE filter #-} + +-- | \(O(n)\). Keep the @Just@s in a sequence. +-- +-- ==== __Examples__ +-- +-- >>> catMaybes (fromList [Just 1, Nothing, Nothing, Just 10, Just 100]) +-- [1,10,100] +catMaybes :: Seq (Maybe a) -> Seq a +catMaybes t = mapMaybe id t + +-- | \(O(n)\). Map over elements and collect the @Just@s. +mapMaybe :: (a -> Maybe b) -> Seq a -> Seq b +mapMaybe = + (coerce :: ((a -> Identity (Maybe b)) -> Seq a -> Identity (Seq b)) + -> (a -> Maybe b) -> Seq a -> Seq b) + mapMaybeA +{-# INLINE mapMaybe #-} + +-- | \(O(n)\). Map over elements and split the @Left@s and @Right@s. +-- +-- ==== __Examples__ +-- +-- >>> mapEither (\x -> if odd x then Left x else Right x) (fromList [1..10]) +-- ([1,3,5,7,9],[2,4,6,8,10]) +mapEither :: (a -> Either b c) -> Seq a -> (Seq b, Seq c) +mapEither = + (coerce :: ((a -> Identity (Either b c)) -> Seq a -> Identity (Seq b, Seq c)) + -> (a -> Either b c) -> Seq a -> (Seq b, Seq c)) + mapEitherA +{-# INLINE mapEither #-} + +-- | \(O(n)\). Keep elements that satisfy an applicative predicate. +filterA :: Applicative f => (a -> f Bool) -> Seq a -> f (Seq a) +filterA f = mapMaybeA (\x -> fmap (\b -> if b then Just x else Nothing) (f x)) +{-# INLINE filterA #-} + +-- | \(O(n)\). Traverse over elements and collect the @Just@s. +mapMaybeA :: Applicative f => (a -> f (Maybe b)) -> Seq a -> f (Seq b) +mapMaybeA f = \case + Tree x xs -> Ap.liftA2 (maybe fromTree Tree) (f x) (T.mapMaybeA f xs) + Empty -> pure Empty +{-# INLINE mapMaybeA #-} + +-- | \(O(n)\). Traverse over elements and split the @Left@s and @Right@s. +mapEitherA + :: Applicative f => (a -> f (Either b c)) -> Seq a -> f (Seq b, Seq c) +mapEitherA f = \case + Tree x xs -> (\g -> Ap.liftA2 g (f x) (T.mapEitherA f xs)) $ \mx z -> + case mx of + Left x' -> unS2 $ bimap (Tree x') fromTree z + Right x' -> unS2 $ bimap fromTree (Tree x') z + Empty -> pure (Empty, Empty) + where + unS2 (U.S2 x y) = (x, y) +{-# INLINE mapEitherA #-} + +-- | \(O(i + \log n)\). The longest prefix of elements that satisfy a predicate. +-- \(i\) is the length of the prefix. +-- +-- ==== __Examples__ +-- +-- >>> takeWhile even (fromList [2,4,6,1,3,2,4]) +-- [2,4,6] +takeWhile :: (a -> Bool) -> Seq a -> Seq a +takeWhile p t = IFo.ifoldr (\i x z -> if p x then z else take i t) t t +{-# INLINE takeWhile #-} + +-- | \(O(i + \log n)\). The remainder after removing the longest prefix of +-- elements that satisfy a predicate. +-- \(i\) is the length of the prefix. +-- +-- ==== __Examples__ +-- +-- >>> dropWhile even (fromList [2,4,6,1,3,2,4]) +-- [1,3,2,4] +dropWhile :: (a -> Bool) -> Seq a -> Seq a +dropWhile p t = IFo.ifoldr (\i x z -> if p x then z else drop i t) Empty t +{-# INLINE dropWhile #-} + +-- | \(O(i + \log n)\). The longest prefix of elements that satisfy a predicate, +-- together with the remainder of the sequence. +-- \(i\) is the length of the prefix. +-- +-- @span p xs == ('takeWhile' p xs, 'dropWhile' p xs)@ +-- +-- ==== __Examples__ +-- +-- >>> span even (fromList [2,4,6,1,3,2,4]) +-- ([2,4,6],[1,3,2,4]) +span :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +span p t = IFo.ifoldr (\i x z -> if p x then z else splitAt i t) (t, Empty) t +{-# INLINE span #-} + +-- | \(O(i + \log n)\). The longest prefix of elements that /do not/ satisfy a +-- predicate, together with the remainder of the sequence. \(i\) is the length +-- of the prefix. +-- +-- @break p == 'span' (not . p)@ +-- +-- ==== __Examples__ +-- +-- >>> break odd (fromList [2,4,6,1,3,2,4]) +-- ([2,4,6],[1,3,2,4]) +break :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +break p = span (not . p) +{-# INLINE break #-} + +-- | \(O(i + \log n)\). The longest suffix of elements that satisfy a predicate. +-- \(i\) is the length of the suffix. +takeWhileEnd :: (a -> Bool) -> Seq a -> Seq a +takeWhileEnd p t = IFo.ifoldl (\i z x -> if p x then z else drop (i+1) t) t t +{-# INLINE takeWhileEnd #-} + +-- | \(O(i + \log n)\). The remainder after removing the longest suffix of +-- elements that satisfy a predicate. +-- \(i\) is the length of the suffix. +dropWhileEnd :: (a -> Bool) -> Seq a -> Seq a +dropWhileEnd p t = + IFo.ifoldl (\i z x -> if p x then z else take (i+1) t) Empty t +{-# INLINE dropWhileEnd #-} + +-- | \(O(i + \log n)\). The longest suffix of elements that satisfy a predicate, +-- together with the remainder of the sequence. +-- \(i\) is the length of the suffix. +-- +-- @spanEnd p xs == ('dropWhileEnd' p xs, 'takeWhileEnd' p xs)@ +spanEnd :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +spanEnd p t = + IFo.ifoldl (\i z x -> if p x then z else splitAt (i+1) t) (Empty, t) t +{-# INLINE spanEnd #-} + +-- | \(O(i + \log n)\). The longest suffix of elements that /do not/ satisfy a +-- predicate, together with the remainder of the sequence. +-- \(i\) is the length of the suffix. +-- +-- @breakEnd p == 'spanEnd' (not . p)@ +breakEnd :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +breakEnd p = spanEnd (not . p) +{-# INLINE breakEnd #-} + +-------------- +-- Transform +-------------- + +-- | \(O(n)\). Reverse a sequence. +-- +-- ==== __Examples__ +-- +-- >>> reverse (fromList [1,2,3,4,5]) +-- [5,4,3,2,1] +reverse :: Seq a -> Seq a +reverse (Tree x xs) = case T.uncons (rev xs) of + U.SNothing -> Tree x Tip + U.SJust (U.S2 x' xs') -> Tree x' (T.snoc xs' x) + where + rev T.Tip = T.Tip + rev (T.Bin sz y l r) = T.Bin sz y (rev r) (rev l) +reverse Empty = Empty + +-- | \(O(n)\). Intersperse an element between the elements of a sequence. +-- +-- ==== __Examples__ +-- +-- >>> intersperse '.' (fromList "HELLO") +-- "H.E.L.L.O" +intersperse :: a -> Seq a -> Seq a +intersperse y (Tree x xs) = case T.unsnoc (go xs) of + U.SNothing -> error "Seq.intersperse: impossible" + U.SJust (U.S2 xs' _) -> Tree x xs' + where + go T.Tip = T.Bin 1 y T.Tip T.Tip + go (T.Bin sz z l r) = T.Bin (sz*2+1) z (go l) (go r) + -- No need to balance, x <= 3y => 2x+1 <= 3(2y+1) +intersperse _ Empty = Empty + +-- | \(O(n)\). Like 'Data.Foldable.foldl'' but keeps all intermediate values. +-- +-- ==== __Examples__ +-- +-- >>> scanl (+) 0 (fromList [1..5]) +-- [0,1,3,6,10,15] +scanl :: (b -> a -> b) -> b -> Seq a -> Seq b +scanl f !z0 = + cons z0 . + flip U.evalSState z0 . + Tr.traverse (\x -> U.sState (\z -> let z' = f z x in U.S2 z' z')) +{-# INLINE scanl #-} + +-- Note [SState for scans] +-- ~~~~~~~~~~~~~~~~~~~~~~ +-- SState is better than Trans.State.Strict. +-- For example, for scanl (+) (0 :: Int), the accumulator Int is unboxed with +-- SState but not with Trans.State.Strict. + +-- | \(O(n)\). Like 'Data.Foldable.foldr'' but keeps all intermediate values. +-- +-- ==== __Examples__ +-- +-- >>> scanr (+) 0 (fromList [1..5]) +-- [15,14,12,9,5,0] +scanr :: (a -> b -> b) -> b -> Seq a -> Seq b +scanr f !z0 = + flip snoc z0 . + flip U.evalSState z0 . + forwards . + Tr.traverse + (\x -> Backwards (U.sState (\z -> let z' = f x z in U.S2 z' z'))) +{-# INLINE scanr #-} +-- See Note [SState for scans] + +-- | \(O(n \log n)\). Sort a sequence. The sort is stable. +-- +-- ==== __Examples__ +-- +-- >>> sort (fromList [4,2,3,5,1]) +-- [1,2,3,4,5] +sort :: Ord a => Seq a -> Seq a +sort = sortBy compare +{-# INLINABLE sort #-} + +-- | \(O(n \log n)\). Sort a sequence using a comparison function. The sort is +-- stable. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Ord (Down, comparing) +-- >>> sortBy (comparing Down) (fromList [4,2,3,5,1]) +-- [5,4,3,2,1] +sortBy :: (a -> a -> Ordering) -> Seq a -> Seq a +sortBy cmp xs = IFu.imap (\i _ -> A.indexArray xa i) xs + where + n = length xs + xa = A.createArray n errorElement $ \ma@(A.MutableArray ma#) -> do + IFo.ifoldr (\i x z -> A.writeArray ma i x *> z) (pure ()) xs + Sam.sortArrayBy cmp ma# 0 n +{-# INLINABLE sortBy #-} + +-- Note [Inlinable sortBy] +-- ~~~~~~~~~~~~~~~~~~~~~~~ +-- Don't INLINE sortBy because sortArrayBy is huge. The user can use Exts.inline +-- if they like. + +-------------------- +-- Search and test +-------------------- + +-- | \(O(n)\). The last element satisfying a predicate. +-- +-- To get the first element, use @Data.Foldable.'Data.Foldable.find'@. +findEnd :: (a -> Bool) -> Seq a -> Maybe a +findEnd f = + Monoid.getLast . foldMap (\x -> Monoid.Last (if f x then Just x else Nothing)) +{-# INLINE findEnd #-} + +-- | \(O(n)\). The index of the first element satisfying a predicate. +-- +-- ==== __Examples__ +-- +-- >>> findIndex even (fromList [1..5]) +-- Just 1 +-- >>> findIndex (<0) (fromList [1..5]) +-- Nothing +findIndex :: (a -> Bool) -> Seq a -> Maybe Int +findIndex f = + Monoid.getFirst . + IFo.ifoldMap (\i x -> Monoid.First (if f x then Just i else Nothing)) +{-# INLINE findIndex #-} + +-- | \(O(n)\). The index of the last element satisfying a predicate. +findIndexEnd :: (a -> Bool) -> Seq a -> Maybe Int +findIndexEnd f = + Monoid.getLast . + IFo.ifoldMap (\i x -> Monoid.Last (if f x then Just i else Nothing)) +{-# INLINE findIndexEnd #-} + +-- | \(O(n_1 + n_2)\). Indices in the second sequence where the first sequence +-- begins as a substring. Includes overlapping occurences. +-- +-- ==== __Examples__ +-- +-- >>> infixIndices (fromList "ana") (fromList "banana") +-- [1,3] +-- >>> infixIndices (fromList [0]) (fromList [1,2,3]) +-- [] +-- >>> infixIndices (fromList "") (fromList "abc") +-- [0,1,2,3] +infixIndices :: Eq a => Seq a -> Seq a -> [Int] +infixIndices t1 t2 + | null t1 = [0 .. length t2] + | compareLength t1 t2 == GT = [] + | otherwise = X.build $ \lcons lnil -> + let n1 = length t1 + t1a = infixIndicesMkArray n1 t1 + !(!mt, !mt0) = KMP.build t1a + f !i x k = \ !m -> case KMP.step mt m x of + (b,m') -> + if b + then lcons (i-n1+1) (k m') + else k m' + in IFo.ifoldr f (\ !_ -> lnil) t2 mt0 +{-# INLINE infixIndices #-} -- Inline for fusion + +infixIndicesMkArray :: Int -> Seq a -> A.Array a +infixIndicesMkArray !n !t = A.createArray n errorElement $ \ma -> + IFo.ifoldr (\i x z -> A.writeArray ma i x *> z) (pure ()) t + +-- | \(O(\log n)\). Binary search for an element in a sequence. +-- +-- Given a function @f@ this function returns an arbitrary element @x@, if it +-- exists, such that @f x = EQ@. @f@ must be monotonic on the sequence— +-- specifically @fmap f@ must result in a sequence which has many (possibly +-- zero) @LT@s, followed by many @EQ@s, followed by many @GT@s. +-- +-- ==== __Examples__ +-- +-- >>> binarySearchFind (`compare` 8) (fromList [2,4..10]) +-- Just 8 +-- >>> binarySearchFind (`compare` 3) (fromList [2,4..10]) +-- Nothing +binarySearchFind :: (a -> Ordering) -> Seq a -> Maybe a +binarySearchFind f t = case t of + Empty -> Nothing + Tree x xs -> case f x of + LT -> go xs + EQ -> Just x + GT -> Nothing + where + go Tip = Nothing + go (Bin _ y l r) = case f y of + LT -> go r + EQ -> Just y + GT -> go l +{-# INLINE binarySearchFind #-} + +-- | \(O(\min(n_1,n_2))\). Whether the first sequence is a prefix of the second. +-- +-- ==== __Examples__ +-- +-- >>> fromList "has" `isPrefixOf` fromList "haskell" +-- True +-- >>> fromList "ask" `isPrefixOf` fromList "haskell" +-- False +isPrefixOf :: Eq a => Seq a -> Seq a -> Bool +isPrefixOf t1 t2 = + compareLength t1 t2 /= GT && Stream.isPrefixOf (stream t1) (stream t2) +{-# INLINABLE isPrefixOf #-} + +-- | \(O(\min(n_1,n_2))\). Whether the first sequence is a suffix of the second. +-- +-- ==== __Examples__ +-- +-- >>> fromList "ell" `isSuffixOf` fromList "haskell" +-- True +-- >>> fromList "ask" `isSuffixOf` fromList "haskell" +-- False +isSuffixOf :: Eq a => Seq a -> Seq a -> Bool +isSuffixOf t1 t2 = + compareLength t1 t2 /= GT && Stream.isPrefixOf (streamEnd t1) (streamEnd t2) +{-# INLINABLE isSuffixOf #-} + +-- | \(O(n_1 + n_2)\). Whether the first sequence is a substring of the second. +-- +-- ==== __Examples__ +-- +-- >>> fromList "meow" `isInfixOf` fromList "homeowner" +-- True +-- >>> fromList [2,4] `isInfixOf` fromList [2,3,4] +-- False +isInfixOf :: Eq a => Seq a -> Seq a -> Bool +isInfixOf t1 t2 = not (null (infixIndices t1 t2)) +{-# INLINABLE isInfixOf #-} + +-- | \(O(n_1 + n_2)\). Whether the first sequence is a subsequence of the +-- second. +-- +-- ==== __Examples__ +-- +-- >>> fromList [2,4] `isSubsequenceOf` [2,3,4] +-- True +-- >>> fromList "tab" `isSubsequenceOf` fromList "bat" +-- False +isSubsequenceOf :: Eq a => Seq a -> Seq a -> Bool +isSubsequenceOf t1 t2 = + compareLength t1 t2 /= GT && Stream.isSubsequenceOf (stream t1) (stream t2) +{-# INLINABLE isSubsequenceOf #-} + +-------- +-- Zip +-------- + +-- | \(O(\min(n_1,n_2))\). Zip two sequences. The result is as long as the +-- shorter sequence. +zip :: Seq a -> Seq b -> Seq (a, b) +zip t1 t2 = zipWith (,) t1 t2 + +-- | \(O(\min(n_1,n_2,n_3))\). Zip three sequences. The result is as long as the +-- shortest sequence. +zip3 :: Seq a -> Seq b -> Seq c -> Seq (a, b, c) +zip3 t1 t2 t3 = zipWith3 (,,) t1 t2 t3 + +-- | \(O(\min(n_1,n_2))\). Zip two sequences with a function. The result is +-- as long as the shorter sequence. +-- +-- ==== __Examples__ +-- +-- >>> zipWith (+) (fromList [1,2,3]) (fromList [1,1,1,1,1]) +-- [2,3,4] +zipWith :: (a -> b -> c) -> Seq a -> Seq b -> Seq c +zipWith = + (coerce :: ((a -> b -> Identity c) -> Seq a -> Seq b -> Identity (Seq c)) + -> (a -> b -> c) -> Seq a -> Seq b -> Seq c) + zipWithM +{-# INLINE zipWith #-} + +-- | \(O(\min(n_1,n_2,n_3))\). Zip three sequences with a function. The result +-- is as long as the shortest sequence. +zipWith3 :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d +zipWith3 = + (coerce :: ((a -> b -> c -> Identity d) -> Seq a -> Seq b -> Seq c -> Identity (Seq d)) + -> (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d) + zipWith3M +{-# INLINE zipWith3 #-} + +-- | \(O(\min(n_1,n_2))\). Zip two sequences with a monadic function. The result +-- is as long as the shorter sequence. +zipWithM :: Monad m => (a -> b -> m c) -> Seq a -> Seq b -> m (Seq c) +zipWithM f t1 t2 = zipWithStreamM f t1 (stream t2) +{-# INLINE zipWithM #-} + +-- | \(O(\min(n_1,n_2,n_3))\). Zip three sequences with a monadic function. The +-- result is as long as the shortest sequence. +zipWith3M + :: Monad m => (a -> b -> c -> m d) -> Seq a -> Seq b -> Seq c -> m (Seq d) +zipWith3M f t1 t2 t3 = + zipWithStreamM + (\x (U.S2 y z) -> f x y z) + t1 + (Stream.zipWith U.S2 (stream t2) (stream t3)) +{-# INLINE zipWith3M #-} + +zipWithStreamM :: Monad m => (a -> b -> m c) -> Seq a -> Stream b -> m (Seq c) +zipWithStreamM f t strm = case t of + Empty -> pure Empty + Tree x xs -> case strm of + Stream step s -> case step s of + Done -> pure Empty + Yield y s1 -> + Ap.liftA2 Tree (f x y) (T.zipWithStreamM f xs (Stream step s1)) +{-# INLINE zipWithStreamM #-} + +-- | \(O(n)\). Unzip a sequence of pairs. +unzip :: Seq (a, b) -> (Seq a, Seq b) +unzip t = unzipWith id t + +-- | \(O(n)\). Unzip a sequence of triples. +unzip3 :: Seq (a, b, c) -> (Seq a, Seq b, Seq c) +unzip3 t = unzipWith3 id t + +-- | \(O(n)\). Map over a sequence and unzip the result. +-- +-- ==== __Examples__ +-- +-- >>> unzipWith (\x -> (x-1, x*2)) (fromList [1..5]) +-- ([0,1,2,3,4],[2,4,6,8,10]) +unzipWith :: (a -> (b, c)) -> Seq a -> (Seq b, Seq c) +unzipWith f t = case t of + Tree x xs -> + case (f x, T.unzipWithA (Identity U.#. f) xs) of + ((x1,x2), Identity (U.S2 xs1 xs2)) -> + let !t1 = Tree x1 xs1 + !t2 = Tree x2 xs2 + in (t1,t2) + Empty -> (Empty, Empty) +{-# INLINE unzipWith #-} + +-- | \(O(n)\). Map over a sequence and unzip the result. +unzipWith3 :: (a -> (b, c, d)) -> Seq a -> (Seq b, Seq c, Seq d) +unzipWith3 f t = case t of + Tree x xs -> + case (f x, T.unzipWith3A (Identity U.#. f) xs) of + ((x1,x2,x3), Identity (U.S3 xs1 xs2 xs3)) -> + let !t1 = Tree x1 xs1 + !t2 = Tree x2 xs2 + !t3 = Tree x3 xs3 + in (t1,t2,t3) + Empty -> (Empty, Empty, Empty) +{-# INLINE unzipWith3 #-} + +-------- +-- Util +-------- + +fromTree :: Tree a -> Seq a +fromTree t = case T.uncons t of + U.SNothing -> Empty + U.SJust (U.S2 x xs) -> Tree x xs +{-# INLINE fromTree #-} + +-- Note [compareLength] +-- ~~~~~~~~~~~~~~~~~~~~ +-- The following functions exist for a bit of efficiency. GHC generates some +-- unnecessary branches for the simple `compare (length x) (length y)`, because +-- it does not know that the size of a Bin is always > the size of a Tip. + +compareLength :: Seq a -> Seq b -> Ordering +compareLength l r = case l of + Tree _ xs -> case r of + Tree _ ys -> compareSize xs ys + Empty -> GT + Empty -> case r of + Tree _ _ -> LT + Empty -> EQ +{-# INLINE compareLength #-} + +compareSize :: Tree a -> Tree b -> Ordering +compareSize l r = case l of + Bin szl _ _ _ -> case r of + Bin szr _ _ _ -> compare szl szr + Tip -> GT + Tip -> case r of + Bin _ _ _ _ -> LT + Tip -> EQ +{-# INLINE compareSize #-} + +---------- +-- Build +---------- + +-- WARNING +-- +-- The functions below are similar but they should not be mixed together! All of +-- them operate on Stack, but what the Stack means is not the same between +-- functions. +-- +-- left-to-right, 1 element at a time: ltrPush, ltrFinish +-- left-to-right, many elements at a time: ltrPushMany, ltrFinish +-- right-to-left, 1 element at a time: rtlPush, rtlFinish + +-- Note [fromList implementation] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- fromList is implemented by keeping a Stack where Seqs at each level have +-- their size as a power of 2. The powers of 2 increase down the stack. New +-- elements are pushed on the stack as Seqs of size 1. They are linked down the +-- stack when they match the next Seq's size. Every such link links two perfect +-- binary trees in O(1) using Bin. For n elements this takes O(n). +-- +-- At the end, the Seqs in the stack are linked together from small to large, +-- balancing as necessary. Linking two Seqs A and B takes +-- O(|log(size A) - log(size B)|). The sizes of the Seqs in the stack are the +-- component powers of 2 of the total size n. +-- Let there be k powers of 2 in n, i.e. n = \sum_{i=1}^k p_i. +-- Then the total cost of the links is +-- O(\sum_{i=2}^k (\log 2^{p_i} - \log (\sum_{j=1}^{i-1} 2^{p_j}))) +-- = O(\sum_{i=2}^k (\log 2^{p_i} - \log 2^{p_{i-1}})) +-- = O(\sum_{i=2}^k (p_i - p_{i-1})) +-- = O(p_k - p_1) +-- = O(\log n) + +-- Note [concatMap implementation] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- The concatMap implementation is not unlike the fromList implementation. +-- Since arbitrary sized Seqs have to be concatenated, the Stack will not have +-- trees of sizes as perfect powers of 2. Instead, an invariant is maintained +-- that 2*size(stack!!0) <= size(stack!!1). This keeps the depth of the stack +-- bounded by O(log N), where N is the total size of Seqs in the stack. +-- +-- If a new Seq is to be added and it is small enough to not violate the +-- invariant, it is simply pushed on the stack. If it would violate the +-- invariant, Seqs on the stack are linked to restore it. The idea here is +-- to merge Seqs of similar sizes as much as possible, since we want to minimize +-- the cost of linking which is O(|log(size A) - log(size B)|). The exact +-- strategy used for this is "2-merge", which has been described and +-- analyzed by Sam Buss and Alexander Knop in "Strategies for Stable Merge +-- Sorting" (https://arxiv.org/abs/1801.04641) for use in mergesort. +-- A merge strategy for mergesort hasn't been blindly adopted here. I arrived +-- at this strategy in an attempt to adopt fromList's implementation before I +-- was aware of the above paper. The incentives to merge similar sizes are very +-- clear here, more so compared to mergesort. +-- +-- Finding a good complexity bound for this algorithm is a little tricky. +-- Given that we merge m Seqs of sizes n_i with total size N, I estimate the +-- complexity to be +-- O(log N - log_{n_m} + \sum_{i=1}^m \max(1, log_{n_i} - log_{n_{i-1})). +-- The log_{n_i} - log_{n_{i-1}} term is an upper bound on the cost of restoring +-- invariants when adding Seqs that turn out to be too big. +-- The log N - log_{n_m} is the final linking cost. +-- A special case of this is if all n_i are equal, say s. The complexity becomes +-- O(log sm - log s + m) = O(m). +-- The worst case occurs when Seqs of size 1 are interleaved into a sequence of +-- Seqs. The complexity becomes O(\sum_{i=1}^m log_{n_i}). + +ltrPush :: Stack a -> a -> Stack a +ltrPush stk y = case stk of + Push x Tip stk' -> ltrPushLoop stk' x 1 (T.Bin 1 y Tip Tip) + _ -> Push y Tip stk + +ltrPushLoop :: Stack a -> a -> Int -> Tree a -> Stack a +ltrPushLoop stk y !ysz ys = case stk of + Push x xs@(Bin xsz _ _ _) stk' + | xsz == ysz -> ltrPushLoop stk' x sz (Bin sz y xs ys) + where + sz = xsz+xsz+1 + _ -> Push y ys stk + +rtlPush :: a -> Stack a -> Stack a +rtlPush x = \case + Push y Tip stk' -> rtlPushLoop x 1 (T.Bin 1 y Tip Tip) stk' + stk -> Push x Tip stk + +rtlPushLoop :: a -> Int -> Tree a -> Stack a -> Stack a +rtlPushLoop x !xsz xs = \case + Push y ys@(Bin ysz _ _ _) stk' + | xsz == ysz -> rtlPushLoop x sz (Bin sz y xs ys) stk' + where + sz = xsz+xsz+1 + stk -> Push x xs stk + +ltrPushMany :: Stack a -> a -> Tree a -> Stack a +ltrPushMany stk y ys = case stk of + Push x xs stk' + | ysz > xsz `div` 2 -> ltrPushManyLoop stk' x xsz xs y ysz ys + | otherwise -> Push y ys stk + where + xsz = 1 + T.size xs + ysz = 1 + T.size ys + Nil -> Push y ys Nil + +ltrPushManyLoop + :: Stack a -> a -> Int -> Tree a -> a -> Int -> Tree a -> Stack a +ltrPushManyLoop stk y !ysz ys z !zsz zs = case stk of + Push x xs@(Bin xsz1 _ _ _) stk' + | xsz < zsz + -> ltrPushManyLoop stk' x (xsz + ysz) (T.link y xs ys) z zsz zs + | yzsz > xsz `div` 2 + -> ltrPushManyLoop stk' x xsz xs y yzsz (T.link z ys zs) + | otherwise + -> Push y (T.link z ys zs) stk + where + xsz = 1+xsz1 + yzsz = ysz+zsz + _ -> Push y (T.link z ys zs) stk + +ltrFinish :: Stack a -> Seq a +ltrFinish = wrapUpStack + Empty + U.S2 + (\(U.S2 y ys) x xs -> U.S2 x (T.link y xs ys)) + (\(U.S2 y ys) -> Tree y ys) + +rtlFinish :: Stack a -> Seq a +rtlFinish = wrapUpStack + Empty + U.S2 + (\(U.S2 x xs) y ys -> U.S2 x (T.link y xs ys)) + (\(U.S2 x xs) -> Tree x xs) + +----------- +-- Stream +----------- + +-- Note [Streams] +-- ~~~~~~~~~~~~~~~~ +-- Streams are used here for two reasons. +-- +-- 1. It is better to implement lazy folds (foldr, foldl, etc) using Streams +-- rather than tree traversals. This is because they form loops which GHC +-- can optimize better on fusing with a consumer. For instance, the "cps +-- sum foldr" benchmark takes ~85% more time if foldr is implemented as a +-- recursive tree traversal. However, such an implementation is a little +-- faster for non-fusion use cases. For instance, the "foldr short-circuit" +-- benchmark takes ~30% less time. This behavior can be obtained when +-- desirable using foldMap with Endo. +-- 2. Streams can fuse for zip-like operations, so we use it to implement such +-- functions. These are decently fast, and we are saved from having to write +-- messy multi-tree traversals. Note that fold/build cannot fuse zips. +-- `zip = fromList (List.zip (toList t) (toList t))`, for instance, takes +-- ~40% more time compared to the stream-based zip. + +stream :: Seq a -> Stream a +stream !t = Stream step s + where + s = case t of + Tree x xs -> Push x xs Nil + Empty -> Nil + step = \case + Push x xs stk -> let !stk' = down xs stk in Yield x stk' + Nil -> Done + {-# INLINE [0] step #-} +{-# INLINE stream #-} + +streamEnd :: Seq a -> Stream a +streamEnd !t = Stream step s + where + s = case t of + Tree x xs -> Push x xs Nil + Empty -> Nil + step = \case + Push x xs stk -> case rDown x xs stk of + U.S2 y stk' -> Yield y stk' + Nil -> Done + {-# INLINE [0] step #-} +{-# INLINE streamEnd #-} + +down :: Tree a -> Stack a -> Stack a +down (Bin _ x l r) stk = down l (Push x r stk) +down Tip stk = stk + +rDown :: a -> Tree a -> Stack a -> U.S2 a (Stack a) +rDown !y (Bin _ x l r) stk = rDown x r (Push y l stk) +rDown y Tip stk = U.S2 y stk + +---------- +-- Stack +---------- + +-- This is used in various places. What it stores depends on the specific use +-- case. +data Stack a = Push !a !(Tree a) !(Stack a) | Nil + +wrapUpStack + :: c -- empty + -> (a -> Tree a -> b) -- initial + -> (b -> a -> Tree a -> b) -- fold fun + -> (b -> c) -- finish + -> Stack a + -> c +wrapUpStack z0 f0 f fin = go + where + go Nil = z0 + go (Push x xs stk) = go1 (f0 x xs) stk + go1 !z Nil = fin z + go1 z (Push x xs stk) = go1 (f z x xs) stk +{-# INLINE wrapUpStack #-} + +------------ +-- Testing +------------ + +valid :: Seq a -> Bool +valid = \case + Tree _ xs -> T.valid xs + Empty -> True + +debugShowsPrec :: Show a => Int -> Seq a -> ShowS +debugShowsPrec p = \case + Tree x xs -> + showParen (p > 10) $ + showString "Tree " . + showsPrec 11 x . + showString " " . + T.debugShowsPrec 11 xs + Empty -> showString "Empty" + +---------- +-- Error +---------- + +errorElement :: a +errorElement = error "Seq: errorElement" diff --git a/src/Data/Seqn/Internal/Stream.hs b/src/Data/Seqn/Internal/Stream.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/Stream.hs @@ -0,0 +1,111 @@ +{-# 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 + (==) = eq1 + {-# INLINE (==) #-} + +instance Eq1 Stream where + liftEq f (Stream step1 s10) (Stream step2 s20) = go s10 s20 + where + go s1 s2 = case step1 s1 of + Yield x1 s1' -> case step2 s2 of + Yield x2 s2' -> f x1 x2 && go s1' s2' + Done -> False + Done -> case step2 s2 of + Yield _ _ -> False + Done -> True + {-# INLINE liftEq #-} + +instance Ord a => Ord (Stream a) where + compare = compare1 + {-# INLINE compare #-} + +instance Ord1 Stream where + liftCompare f (Stream step1 s10) (Stream step2 s20) = go s10 s20 + where + go s1 s2 = case step1 s1 of + Yield x1 s1' -> case step2 s2 of + Yield x2 s2' -> f x1 x2 <> go s1' s2' + Done -> GT + Done -> case step2 s2 of + Yield _ _ -> LT + Done -> EQ + {-# INLINE liftCompare #-} + +foldr :: (a -> b -> b) -> b -> Stream a -> b +foldr f z (Stream step s0) = go s0 + where + go s = case step s of + Yield x s' -> f x (go s') + Done -> z +{-# INLINE foldr #-} + +ifoldr :: (Int -> a -> b -> b) -> b -> Int -> (Int -> Int) -> Stream a -> b +ifoldr f z i0 istep (Stream step s0) = go i0 s0 + where + go !i s = case step s of + Yield x s' -> f i x (go (istep i) s') + Done -> z +{-# INLINE ifoldr #-} + +isPrefixOf :: Eq a => Stream a -> Stream a -> Bool +isPrefixOf (Stream step1 s10) (Stream step2 s20) = go s10 s20 + where + go s1 s2 = case step1 s1 of + Yield x1 s1' -> case step2 s2 of + Yield x2 s2' -> x1 == x2 && go s1' s2' + Done -> False + Done -> True +{-# INLINE isPrefixOf #-} + +isSubsequenceOf :: Eq a => Stream a -> Stream a -> Bool +isSubsequenceOf (Stream step1 s10) (Stream step2 s20) = go1 s10 s20 + where + go1 s1 s2 = case step1 s1 of + Yield x s1' -> go2 x s1' s2 + Done -> True + go2 !x s1' s2 = case step2 s2 of + Yield y s2' + | x == y -> go1 s1' s2' + | otherwise -> go2 x s1' s2' + Done -> False +{-# INLINE isSubsequenceOf #-} + +zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c +zipWith f (Stream step1 s10) (Stream step2 s20) = Stream step (U.S2 s10 s20) + where + step (U.S2 s1 s2) = case step1 s1 of + Yield x1 s1' -> case step2 s2 of + Yield x2 s2' -> Yield (f x1 x2) (U.S2 s1' s2') + Done -> Done + Done -> Done + {-# INLINE [0] step #-} +{-# INLINE zipWith #-} diff --git a/src/Data/Seqn/Internal/Tree.hs b/src/Data/Seqn/Internal/Tree.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/Tree.hs @@ -0,0 +1,615 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE LambdaCase #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- This is an internal module. You probably don't need to import this. Use +-- "Data.Seqn.Seq" instead. +-- +-- = WARNING +-- +-- Definitions in this module allow violating invariants that would otherwise be +-- guaranteed by "Data.Seqn.Seq". Use at your own risk! +-- +module Data.Seqn.Internal.Tree + ( + -- * Tree + Tree(..) + + -- * Basic + , size + , bin + + -- * Folds + , foldl' + , ifoldl' + , foldr' + , ifoldr' + , traverse + , itraverse + + -- * Construct + , generateA + + -- * Index + , adjustF + , insertAt + , deleteAt + + -- * Slice + , cons + , snoc + , uncons + , unsnoc + , splitAtF + + -- * Transform + , mapMaybeA + , mapEitherA + + -- * Zip and unzip + , zipWithStreamM + , unzipWithA + , unzipWith3A + + -- * Tree helpers + , fold + , foldSimple + , link + , glue + , merge + , balanceL + , balanceR + + -- * Testing + , valid + , debugShowsPrec + ) where + +import Prelude hiding (concatMap, break, drop, dropWhile, filter, foldl', lookup, map, replicate, reverse, scanl, scanr, span, splitAt, take, takeWhile, traverse, unzip, unzip3, zip, zip3, zipWith, zipWith3) +import qualified Control.Applicative as Ap +import Control.DeepSeq (NFData(..), NFData1(..)) +import Data.Bifunctor (Bifunctor(..)) + +import qualified Data.Seqn.Internal.Util as U +import Data.Seqn.Internal.Stream (Stream(..), Step(..)) + +data Tree a + = Bin {-# UNPACK #-} !Int !a !(Tree a) !(Tree a) + | Tip + +-------------- +-- Instances +-------------- + +instance NFData a => NFData (Tree a) where + rnf = \case + Bin _ x l r -> rnf x `seq` rnf l `seq` rnf r + Tip -> () + {-# INLINABLE rnf #-} + +instance NFData1 Tree where + liftRnf f = go + where + go (Bin _ x l r) = f x `seq` go l `seq` go r + go Tip = () + {-# INLINE liftRnf #-} + +------------------- +-- Basic Tree ops +------------------- + +singleton :: a -> Tree a +singleton x = Bin 1 x Tip Tip +{-# INLINE singleton #-} + +size :: Tree a -> Int +size (Bin n _ _ _) = n +size Tip = 0 +{-# INLINE size #-} + +-- O(1). Link two trees with a value in between. Precondition: The trees are +-- balanced wrt each other. +bin :: a -> Tree a -> Tree a -> Tree a +bin x l r = Bin (size l + size r + 1) x l r +{-# INLINE bin #-} + +---------- +-- Folds +---------- + +-- Note [Folds] +-- ~~~~~~~~~~~~ +-- Certain functions, such as folds, are implemented recursively on non-empty +-- trees, i.e. `go :: Int -> a -> Tree a -> Tree a -> b` instead of +-- `go :: Tree a -> b`. This is simply because benchmarks show this to be +-- faster. + +foldl' :: (b -> a -> b) -> b -> Tree a -> b +foldl' f !z0 = \case + Bin _ x l r -> go z0 x l r + Tip -> z0 + where + go !z !x l r = case l of + Bin _ lx ll lr -> case r of + Bin _ rx rl rr -> + let !z' = go z lx ll lr + in go (f z' x) rx rl rr + Tip -> + let !z' = go z lx ll lr + in f z' x + Tip -> case r of + Bin _ rx rl rr -> go (f z x) rx rl rr + Tip -> f z x +{-# INLINE foldl' #-} + +ifoldl' :: (Int -> b -> a -> b) -> b -> Int -> Tree a -> b +ifoldl' f !z0 !i0 = \case + Bin _ x l r -> go z0 i0 x l r + Tip -> z0 + where + go !z !i !x l r = case l of + Bin lsz lx ll lr -> case r of + Bin _ rx rl rr -> + let !z' = go z i lx ll lr + in go (f (i+lsz) z' x) (i+lsz+1) rx rl rr + Tip -> + let !z' = go z i lx ll lr + in f (i+lsz) z' x + Tip -> case r of + Bin _ rx rl rr -> go (f i z x) (i+1) rx rl rr + Tip -> f i z x +{-# INLINE ifoldl' #-} + +foldr' :: (a -> b -> b) -> b -> Tree a -> b +foldr' f !z0 = \case + Bin _ x l r -> go z0 x l r + Tip -> z0 + where + go !z !x l r = case l of + Bin _ lx ll lr -> case r of + Bin _ rx rl rr -> + let !z' = go z rx rl rr + in go (f x z') lx ll lr + Tip -> go (f x z) lx ll lr + Tip -> case r of + Bin _ rx rl rr -> f x $! go z rx rl rr + Tip -> f x z +{-# INLINE foldr' #-} + +ifoldr' :: (Int -> a -> b -> b) -> b -> Int -> Tree a -> b +ifoldr' f !z0 !i0 = \case + Bin _ x l r -> go z0 i0 x l r + Tip -> z0 + where + go !z !i !x l r = case l of + Bin _ lx ll lr -> case r of + Bin rsz rx rl rr -> + let !z' = go z i rx rl rr + in go (f (i-rsz) x z') (i-rsz-1) lx ll lr + Tip -> go (f i x z) (i-1) lx ll lr + Tip -> case r of + Bin rsz rx rl rr -> f (i-rsz) x $! go z i rx rl rr + Tip -> f i x z +{-# INLINE ifoldr' #-} + +fold + :: b + -> (Int -> a -> b -> b -> b) + -> (Int -> a -> b -> b) + -> (Int -> a -> b -> b) + -> (a -> b) + -> Tree a + -> b +fold tip glr gl gr g = \case + Bin sz x l r -> go sz x l r + Tip -> tip + where + go !sz !x l r = case l of + Bin lsz lx ll lr -> case r of + Bin rsz rx rl rr -> glr sz x (go lsz lx ll lr) (go rsz rx rl rr) + Tip -> gl sz x (go lsz lx ll lr) + Tip -> case r of + Bin rsz rx rl rr -> gr sz x (go rsz rx rl rr) + Tip -> g x +{-# INLINE fold #-} + +foldSimple :: b -> (Int -> a -> b -> b -> b) -> Tree a -> b +foldSimple tip f = fold tip f gl gr g + where + gl !sz x ml = f sz x ml tip + {-# INLINE gl #-} + gr !sz x mr = f sz x tip mr + {-# INLINE gr #-} + g x = f 1 x tip tip + {-# INLINE g #-} +{-# INLINE foldSimple #-} + +traverse :: Applicative f => (a -> f b) -> Tree a -> f (Tree b) +traverse f = fold (pure Tip) glr gl gr g + where + glr !sz x ml mr = liftA3R' (flip (Bin sz)) ml (f x) mr + -- See Note [Traverse liftA3R'] + {-# INLINE glr #-} + gl !sz x ml = Ap.liftA2 (\l' x' -> Bin sz x' l' Tip) ml (f x) + {-# INLINE gl #-} + gr !sz x mr = Ap.liftA2 (\x' r' -> Bin sz x' Tip r') (f x) mr + {-# INLINE gr #-} + g x = fmap singleton (f x) + {-# INLINE g #-} +{-# INLINE traverse #-} + +itraverse :: Applicative f => (Int -> a -> f b) -> Int -> Tree a -> f (Tree b) +itraverse f !i0 = \case + Bin sz x l r -> go i0 sz x l r + Tip -> pure Tip + where + go !i !sz x l r = case l of + Bin lsz lx ll lr -> case r of + Bin rsz rx rl rr -> + liftA3R' + (flip (Bin sz)) + (go i lsz lx ll lr) + (f (i+lsz) x) + (go (i+lsz+1) rsz rx rl rr) + -- See Note [Traverse liftA3R'] + Tip -> + Ap.liftA2 + (\l' x' -> Bin sz x' l' Tip) + (go i lsz lx ll lr) + (f (i+lsz) x) + Tip -> case r of + Bin rsz rx rl rr -> + Ap.liftA2 (\x' r' -> Bin sz x' Tip r') (f i x) (go (i+1) rsz rx rl rr) + Tip -> + fmap (\x' -> Bin sz x' Tip Tip) (f i x) + -- See Note [Traverse] +{-# INLINE itraverse #-} + +-- Note [Traverse liftA3R'] +-- ~~~~~~~~~~~~~~~~~~~~~~~~ +-- We want to associate to the right because we define foldMap using traverse +-- and ifoldMap using itraverse. It is more appropriate to be right-associative +-- for <>. + +-- Right associative and strict +liftA3R' :: Applicative f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d +liftA3R' f mx my mz = + Ap.liftA2 + (\x (U.S2 y z) -> f x y z) + mx + (Ap.liftA2 U.S2 my mz) +{-# INLINE liftA3R' #-} + +-------------- +-- Construct +-------------- + +generateA :: Applicative f => (Int -> f a) -> Int -> Int -> f (Tree a) +generateA f = go + where + go !i n + | n <= 0 = pure Tip + | otherwise = + Ap.liftA3 + (flip (Bin n)) + (go i lsz) + (f (i+lsz)) + (go (i+lsz+1) (n-lsz-1)) + where + lsz = (n-1) `div` 2 +{-# INLINE generateA #-} + +---------- +-- Index +---------- + +-- Precondition: 0 <= i < size xs +adjustF :: Functor f => (a -> f a) -> Int -> Tree a -> f (Tree a) +adjustF f = go + where + go !i = \case + Bin sz x l r -> case compare i szl of + LT -> fmap (\l' -> Bin sz x l' r) (go i l) + EQ -> fmap (\x' -> Bin sz x' l r) (f x) + GT -> fmap (Bin sz x l) (go (i-szl-1) r) + where + szl = size l + Tip -> errorOutOfBounds "Tree.adjustF" +{-# INLINE adjustF #-} + +-- Inserts at ends if not in bounds +insertAt :: Int -> a -> Tree a -> Tree a +insertAt !i x (Bin _ y l r) + | i <= szl = balanceL y (insertAt i x l) r + | otherwise = balanceR y l (insertAt (i-szl-1) x r) + where + szl = size l +insertAt _ x Tip = singleton x + +-- Precondition: 0 <= i < size xs +deleteAt :: Int -> Tree a -> Tree a +deleteAt !i (Bin _ x l r) = case compare i szl of + LT -> balanceR x (deleteAt i l) r + EQ -> glue l r + GT -> balanceL x l (deleteAt (i-szl-1) r) + where + szl = size l +deleteAt _ Tip = errorOutOfBounds "Tree.deleteAt" + +---------- +-- Slice +---------- + +cons :: a -> Tree a -> Tree a +cons x Tip = singleton x +cons x (Bin _ y l r) = balanceL y (cons x l) r + +snoc :: Tree a -> a -> Tree a +snoc Tip x = singleton x +snoc (Bin _ y l r) x = balanceR y l (snoc r x) + +uncons :: Tree a -> U.SMaybe (U.S2 a (Tree a)) +uncons (Bin _ x l r) = U.SJust (unconsSure x l r) +uncons Tip = U.SNothing +{-# INLINE uncons #-} + +unconsSure :: a -> Tree a -> Tree a -> U.S2 a (Tree a) +unconsSure x (Bin _ lx ll lr) r = case unconsSure lx ll lr of + U.S2 y l' -> U.S2 y (balanceR x l' r) +unconsSure x Tip r = U.S2 x r + +unsnoc :: Tree a -> U.SMaybe (U.S2 (Tree a) a) +unsnoc (Bin _ x l r) = U.SJust $ unsnocSure x l r +unsnoc Tip = U.SNothing +{-# INLINE unsnoc #-} + +unsnocSure :: a -> Tree a -> Tree a -> U.S2 (Tree a) a +unsnocSure x l (Bin _ rx rl rr) = case unsnocSure rx rl rr of + U.S2 r' y -> U.S2 (balanceL x l r') y +unsnocSure x l Tip = U.S2 l x + +-- Precondition: 0 <= i < size xs +splitAtF + :: U.Biapplicative f + => Int -> Tree a -> f (Tree a) (U.S2 a (Tree a)) +splitAtF = go + where + go !i (Bin _ x l r) = case compare i szl of + LT -> second (second (\lr -> link x lr r)) (go i l) + EQ -> U.bipure l (U.S2 x r) + GT -> first (link x l) (go (i-szl-1) r) + where + szl = size l + go _ Tip = errorOutOfBounds "Tree.splitAtF" +{-# INLINE splitAtF #-} + +-------------- +-- Transform +-------------- + +mapMaybeA :: Applicative f => (a -> f (Maybe b)) -> Tree a -> f (Tree b) +mapMaybeA f = foldSimple tip g + where + tip = pure Tip + {-# INLINE tip #-} + g _ x ml mr = (\h -> Ap.liftA3 h ml (f x) mr) $ \l my r -> + case my of + Nothing -> merge l r + Just y -> link y l r + {-# INLINE g #-} +{-# INLINE mapMaybeA #-} + +mapEitherA + :: Applicative f + => (a -> f (Either b c)) -> Tree a -> f (U.S2 (Tree b) (Tree c)) +mapEitherA f = foldSimple tip g + where + tip = pure (U.bipure Tip Tip) + {-# INLINE tip #-} + g _ x ml mr = (\h -> Ap.liftA3 h ml (f x) mr) $ \l my r -> + case my of + Left y -> U.biliftA2 (link y) merge l r + Right y -> U.biliftA2 merge (link y) l r + {-# INLINE g #-} +{-# INLINE mapEitherA #-} + +------------------ +-- Zip and unzip +------------------ + +zipWithStreamM :: Monad m => (a -> b -> m c) -> Tree a -> Stream b -> m (Tree c) +zipWithStreamM f t (Stream step s) = U.evalSStateT (foldSimple tip g t) s + where + tip = pure Tip + {-# INLINE tip #-} + g _ x ml mr = U.SStateT $ \s2 -> do + U.S2 s3 l <- U.runSStateT ml s2 + case step s3 of + Done -> pure $ U.S2 s3 l + Yield y s4 -> do + z <- f x y + U.S2 s5 r <- U.runSStateT mr s4 + pure $! U.S2 s5 (link z l r) + {-# INLINE g #-} +{-# INLINE zipWithStreamM #-} + +unzipWithA + :: Applicative f => (a -> f (b, c)) -> Tree a -> f (U.S2 (Tree b) (Tree c)) +unzipWithA f = foldSimple tip g + where + tip = pure (U.S2 Tip Tip) + {-# INLINE tip #-} + g !sz x ml mr = (\h -> Ap.liftA3 h ml (f x) mr) $ + \(U.S2 l1 l2) (x1,x2) (U.S2 r1 r2) -> + U.S2 (Bin sz x1 l1 r1) (Bin sz x2 l2 r2) + {-# INLINE g #-} +{-# INLINE unzipWithA #-} + +unzipWith3A + :: Applicative f + => (a -> f (b, c, d)) + -> Tree a + -> f (U.S3 (Tree b) (Tree c) (Tree d)) +unzipWith3A f = foldSimple tip g + where + tip = pure (U.S3 Tip Tip Tip) + {-# INLINE tip #-} + g !sz x ml mr = (\h -> Ap.liftA3 h ml (f x) mr) $ + \(U.S3 l1 l2 l3) (x1,x2,x3) (U.S3 r1 r2 r3) -> + U.S3 (Bin sz x1 l1 r1) (Bin sz x2 l2 r2) (Bin sz x3 l3 r3) + {-# INLINE g #-} +{-# INLINE unzipWith3A #-} + +----------- +-- Errors +----------- + +errorOutOfBounds :: String -> a +errorOutOfBounds name = error (name ++ ": out of bounds") + +------------ +-- Balance +------------ + +-- O(|log n1 - log n2|). Link two trees with a value in between. +link :: a -> Tree a -> Tree a -> Tree a +link !x Tip r = cons x r +link x l Tip = snoc l x +link x l@(Bin ls lx ll lr) r@(Bin rs rx rl rr) + | delta*ls < rs = balanceL rx (linkL x ls l rl) rr + | delta*rs < ls = balanceR lx ll (linkR x lr rs r) + | otherwise = Bin (1+ls+rs) x l r +{-# INLINE link #-} + +linkL :: a -> Int -> Tree a -> Tree a -> Tree a +linkL !x !ls !l r = case r of + Bin rs rx rl rr + | delta*ls < rs -> balanceL rx (linkL x ls l rl) rr + | otherwise -> Bin (1+ls+rs) x l r + Tip -> error "Tree.linkL: impossible" + +linkR :: a -> Tree a -> Int -> Tree a -> Tree a +linkR !x l !rs !r = case l of + Bin ls lx ll lr + | delta*rs < ls -> balanceR lx ll (linkR x lr rs r) + | otherwise -> Bin (1+ls+rs) x l r + Tip -> error "Tree.linkR: impossible" + +-- O(log (n1 + n2)). Link two trees. +merge :: Tree a -> Tree a -> Tree a +merge Tip r = r +merge l Tip = l +merge l@(Bin ls lx ll lr) r@(Bin rs rx rl rr) + | ls < rs = case unsnocSure lx ll lr of U.S2 l' mx -> link mx l' r + | otherwise = case unconsSure rx rl rr of U.S2 mx r' -> link mx l r' +{-# INLINE merge #-} + +-- O(log (n1 + n2)). Link two trees. Precondition: The trees must be balanced +-- wrt each other. +glue :: Tree a -> Tree a -> Tree a +glue Tip r = r +glue l Tip = l +glue l@(Bin ls lx ll lr) r@(Bin rs rx rl rr) + | ls > rs = case unsnocSure lx ll lr of U.S2 l' m -> balanceR m l' r + | otherwise = case unconsSure rx rl rr of U.S2 m r' -> balanceL m l r' +{-# INLINE glue #-} + +-- Note [Balance] +-- ~~~~~~~~~~~~~~ +-- The balancing code here is largely influenced by the implementation of +-- for the Set type in containers: https://hackage.haskell.org/package/containers +-- The linked papers in Data.Seqn.Seq describe the structure in greater detail. +-- +-- To summarize: +-- * A tree is balanced if size(left child) < delta*size(right child) and vice +-- versa, which a special case for Tips. See `balanceOk` in `valid`, which is +-- used to check that balance holds in tests. +-- * Rebalancing involves rotations. The rotation can be single or double. The +-- constant `ratio` determines whether a double rotation is performed. + +delta, ratio :: Int +delta = 3 +ratio = 2 + +-- O(1). Restores balance with at most one right rotation. Precondition: One +-- right rotation must be enough to restore balance. This is the case when the +-- left tree might have been inserted to or the right tree deleted from. +balanceL :: a -> Tree a -> Tree a -> Tree a +balanceL !x l r = case r of + Tip -> case l of + Tip -> Bin 1 x Tip Tip + Bin _ lx ll lr -> case lr of + Tip -> case ll of + Tip -> Bin 2 x l Tip + Bin _ _ _ _ -> Bin 3 lx ll (Bin 1 x Tip Tip) + Bin _ lrx _ _ -> case ll of + Tip -> Bin 3 lrx (Bin 1 lx Tip Tip) (Bin 1 x Tip Tip) + Bin _ _ _ _ -> Bin 4 lx ll (Bin 2 x lr Tip) + Bin rs _ _ _ -> case l of + Tip -> Bin (1+rs) x Tip r + Bin ls lx ll lr + | ls > delta*rs -> case (ll, lr) of + (Bin lls _ _ _, Bin lrs lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls+rs) lx ll (Bin (1+rs+lrs) x lr r) + | otherwise -> Bin (1+ls+rs) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+rs+size lrr) x lrr r) + _ -> error "Tree.balanceL: impossible" + | otherwise -> Bin (1+ls+rs) x l r +{-# NOINLINE balanceL #-} + +-- O(1). Restores balance with at most one left rotation. Precondition: One left +-- rotation must be enough to restore balance. This is the case when the right +-- tree might have been inserted to or the left tree deleted from. +balanceR :: a -> Tree a -> Tree a -> Tree a +balanceR !x l r = case l of + Tip -> case r of + Tip -> Bin 1 x Tip Tip + Bin _ rx rl rr -> case rl of + Tip -> case rr of + Tip -> Bin 2 x Tip r + Bin _ _ _ _ -> Bin 3 rx (Bin 1 x Tip Tip) rr + Bin _ rlx _ _ -> case rr of + Tip -> Bin 3 rlx (Bin 1 x Tip Tip) (Bin 1 rx Tip Tip) + Bin _ _ _ _ -> Bin 4 rx (Bin 2 x Tip rl) rr + Bin ls _ _ _ -> case r of + Tip -> Bin (1+ls) x l Tip + Bin rs rx rl rr + | rs > delta*ls -> case (rl, rr) of + (Bin rls rlx rll rlr, Bin rrs _ _ _) + | rls < ratio*rrs -> Bin (1+ls+rs) rx (Bin (1+ls+rls) x l rl) rr + | otherwise -> Bin (1+ls+rs) rlx (Bin (1+ls+size rll) x l rll) (Bin (1+rrs+size rlr) rx rlr rr) + _ -> error "Tree.balanceR: impossible" + | otherwise -> Bin (1+ls+rs) x l r +{-# NOINLINE balanceR #-} + +------------ +-- Testing +------------ + +valid :: Tree a -> Bool +valid s = balanceOk s && sizeOk s + where + balanceOk = \case + Bin _ _ l r -> ok && balanceOk l && balanceOk r + where + ok = size l + size r <= 1 || + (size l <= delta * size r && size r <= delta * size l) + Tip -> True + + sizeOk = \case + Bin sz _ l r -> sizeOk l && sizeOk r && size l + size r + 1 == sz + Tip -> True + +debugShowsPrec :: Show a => Int -> Tree a -> ShowS +debugShowsPrec p = \case + Bin sz x l r -> + showParen (p > 10) $ + showString "Bin " . + shows sz . + showString " " . + showsPrec 11 x . + showString " " . + debugShowsPrec 11 l . + showString " " . + debugShowsPrec 11 r + Tip -> showString "Tip" diff --git a/src/Data/Seqn/Internal/Util.hs b/src/Data/Seqn/Internal/Util.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Internal/Util.hs @@ -0,0 +1,105 @@ +-- | +-- This is an internal module. You probably don't need to import this. Use +-- "Data.Seqn.Seq", "Data.Seqn.MSeq", or "Data.Seqn.PQueue" instead. +-- +-- The only reason to use this module is to use the constructs defined here with +-- other internal modules. +-- +module Data.Seqn.Internal.Util + ( Biapplicative(..) + , S2(..) + , S3(..) + , SMaybe(..) + , Tagged(..) + , SStateT(..) + , evalSStateT + , SState + , sState + , evalSState + , (#.) + ) where + +import qualified Control.Applicative -- for before liftA2 in Prelude +import Data.Bifunctor (Bifunctor(..)) +import Data.Coerce (Coercible, coerce) +import Data.Functor.Const (Const(..)) +import Data.Functor.Identity (Identity(..)) + +class Bifunctor p => Biapplicative p where + bipure :: a -> b -> p a b + biliftA2 :: (a -> b -> c) -> (d -> e -> f) -> p a d -> p b e -> p c f + +instance Biapplicative Const where + bipure x _ = coerce x + biliftA2 f _ = coerce f + +data S2 a b = S2 !a !b + +instance Functor (S2 a) where + fmap f (S2 x y) = S2 x (f y) + +instance Bifunctor S2 where + bimap f g (S2 x y) = S2 (f x) (g y) + +instance Biapplicative S2 where + bipure = S2 + biliftA2 f g (S2 x1 y1) (S2 x2 y2) = S2 (f x1 x2) (g y1 y2) + +data S3 a b c = S3 !a !b !c + +data SMaybe a + = SNothing + | SJust !a + +newtype Tagged a b = Tagged { unTagged :: b } + +instance Functor (Tagged a) where + fmap = coerce + +instance Bifunctor Tagged where + bimap _ = coerce + +instance Biapplicative Tagged where + bipure _ = coerce + biliftA2 _ = coerce + +-- Strict in the state, value, and bind. +newtype SStateT s m a = SStateT { runSStateT :: s -> m (S2 s a) } + +evalSStateT :: Functor m => SStateT s m a -> s -> m a +evalSStateT m s = fmap (\(S2 _ x) -> x) (runSStateT m s) +{-# INLINE evalSStateT #-} + +type SState s = SStateT s Identity + +sState :: (s -> S2 s a) -> SState s a +sState = coerce + +evalSState :: SState s a -> s -> a +evalSState m s = case runSStateT m s of Identity (S2 _ x) -> x +{-# INLINE evalSState #-} + +instance Functor m => Functor (SStateT s m) where + fmap f m = SStateT $ \s -> (fmap . fmap) f (runSStateT m s) + {-# INLINE fmap #-} + +instance Monad m => Applicative (SStateT s m) where + pure x = SStateT $ \s -> pure $ S2 s x + {-# INLINE pure #-} + + liftA2 f m1 m2 = SStateT $ \s -> do + S2 s1 x <- runSStateT m1 s + S2 s2 y <- runSStateT m2 s1 + pure $ S2 s2 (f x y) + {-# INLINE liftA2 #-} + +-- Borrow a trick from base in case GHC has trouble optimizing certain +-- function coercions. +-- +-- See Note [Function coercion] in +-- https://gitlab.haskell.org/ghc/ghc/-/blob/8d67f247c3e4ca3810712654e1becbf927405f6b/libraries/ghc-internal/src/GHC/Internal/Data/Functor/Utils.hs#L134-160 +(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c) +(#.) _ = coerce +{-# INLINE (#.) #-} + +infixr 9 #. diff --git a/src/Data/Seqn/MSeq.hs b/src/Data/Seqn/MSeq.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/MSeq.hs @@ -0,0 +1,304 @@ +-- | +-- = Finite measured sequences +-- +-- A value of type @MSeq a@ is a sequence with elements of type @a@. +-- An @MSeq@ is +-- +-- * Spine-strict, hence finite. @MSeq@ cannot represent infinite sequences. +-- * Value-strict. It is guaranteed that if a @MSeq@ is in +-- [weak head normal form](https://wiki.haskell.org/Weak_head_normal_form) +-- (WHNF), every element of the @Seq@ is also in WHNF. +-- +-- An @MSeq@ provides quick access to the combined \"measure\" of all elements +-- of the sequence. Please see the [Tutorial](#g:tutorial) at the end of this +-- page for an explanation. +-- +-- It is recommended to import this module qualified to avoid name clashes. +-- +-- @ +-- import Data.Seqn.MSeq (Measured, MSeq) +-- import qualified Data.Seqn.MSeq as MSeq +-- @ +-- +-- === Warning +-- +-- The length of a @MSeq@ must not exceed @(maxBound \`div\` 3) :: Int@. If this +-- length is exceeded, the behavior of a @MSeq@ is undefined. This value is very +-- large in practice, greater than \(7 \cdot 10^8\) on 32-bit systems and +-- \(3 \cdot 10^{18}\) on 64-bit systems. +-- +-- === Implementation +-- +-- @MSeq@ is implemented as a +-- [weight-balanced binary tree](https://en.wikipedia.org/wiki/Weight-balanced_tree). +-- This structure is described by +-- +-- * J. Nievergelt and E. M. Reingold, +-- /\"Binary search trees of bounded balance\"/, +-- SIAM Journal of Computing 2(1), 1973, +-- https://doi.org/10.1137/0202005 +-- +-- * Stephen Adams, +-- /\"Efficient sets—a balancing act\"/, +-- Journal of Functional Programming 3(4), 553-561, 1993, +-- https://doi.org/10.1017/S0956796800000885 +-- +-- * Yoichi Hirai and Kazuhiko Yamamoto, +-- /\"Balancing weight-balanced trees\"/, +-- Journal of Functional Programming 21(3), 287-307, 2011, +-- https://doi.org/10.1017/S0956796811000104 +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- /\"Parallel Ordered Sets Using Join\"/, 2016, +-- https://doi.org/10.48550/arXiv.1602.02120 +-- +module Data.Seqn.MSeq + ( + -- * MSeq + S.MSeq + , M.Measured(..) + + -- * Measured queries + , S.summaryMay + , S.summary + , S.binarySearchPrefix + , S.binarySearchSuffix + + -- * Construct + , S.empty + , S.singleton + , S.fromList + , S.fromRevList + , S.replicate + , S.replicateA + , S.generate + , S.generateA + , S.unfoldr + , S.unfoldl + , S.unfoldrM + , S.unfoldlM + , S.concatMap + , S.mfix + + -- * Convert + , S.toRevList + + -- * Index + , S.lookup + , S.index + , (S.!?) + , (S.!) + , S.update + , S.adjust + , S.insertAt + , S.deleteAt + + -- * Slice + , S.cons + , S.snoc + , S.uncons + , S.unsnoc + , S.take + , S.drop + , S.slice + , S.splitAt + , S.takeEnd + , S.dropEnd + , S.splitAtEnd + + -- * Filter + , S.filter + , S.mapMaybe + , S.mapEither + , S.filterA + , S.mapMaybeA + , S.mapEitherA + , S.takeWhile + , S.dropWhile + , S.span + , S.break + , S.takeWhileEnd + , S.dropWhileEnd + , S.spanEnd + , S.breakEnd + + -- * Transform + , S.map + , S.liftA2 + , S.traverse + , S.imap + , S.itraverse + , S.reverse + , S.intersperse + , S.scanl + , S.scanr + , S.sort + , S.sortBy + + -- * Search and test + , S.findEnd + , S.findIndex + , S.findIndexEnd + , S.infixIndices + , S.binarySearchFind + , S.isPrefixOf + , S.isSuffixOf + , S.isInfixOf + , S.isSubsequenceOf + + -- * Zip and unzip + , S.zipWith + , S.zipWith3 + , S.zipWithM + , S.zipWith3M + , S.unzipWith + , S.unzipWith3 + + -- * Force + , S.liftRnf2 + + -- * Tutorial #tutorial# + -- $tutorial + ) where + +import qualified Data.Seqn.Internal.MTree as M +import qualified Data.Seqn.Internal.MSeq as S + +-- $tutorial +-- +-- @MSeq@, like @Seq@, is a sequence which supports operations like @lookup@, +-- @splitAt@, @(<>)@, @foldr@, and more. What makes it different is that it +-- maintains \"measure\"s of the elements in it. Every element in the sequence +-- has an associated measure. The type of this measure must have a @Semigroup@ +-- instance. An @MSeq@ allows accessing the combined measure of all its elements +-- in \(O(1)\) time. +-- +-- == Example: Sum +-- +-- @ +-- data Task = Task +-- !Text -- ^ Name +-- !Word -- ^ Cost +-- deriving Show +-- @ +-- +-- Consider that we need to maintain a sequence of tasks, where each task has +-- some cost. Tasks will be added and removed over time. At various points, the +-- total cost of all the tasks in the sequence must be computed. +-- +-- We may use a @Seq@ to store the task, and calculate the sum when required in +-- \(O(n)\). This is reasonable if such events are rare, but a poor strategy +-- if the sum has to be calculated frequently. In the latter case, we could use +-- an @MSeq@. +-- +-- We begin with some imports. +-- +-- @ +-- import Data.Seqn.MSeq (Measured, MSeq) +-- import qualified Data.Seqn.MSeq as MSeq +-- @ +-- +-- Next, we define the t'Data.Seqn.MSeq.Measured' instance for @Task@. +-- +-- @ +-- {-# LANGUAGE TypeFamilies #-} +-- import "Data.Monoid" (Sum(..)) +-- +-- instance Measured Task where +-- type Measure Task = Sum Word +-- measure (Task _ cost) = Sum cost +-- @ +-- +-- >>> let tasks = MSeq.fromList [Task "A" 50, Task "B" 30, Task "C" 60] +-- >>> tasks +-- [Task "A" 50,Task "B" 30,Task "C" 60] +-- +-- We now have access to the combined measure of the @MSeq@, called the +-- 'Data.Seqn.MSeq.summary', in \(O(1)\). +-- +-- >>> MSeq.summary tasks +-- Sum {getSum = 140} +-- +-- If we modify the task list, the summary will change accordingly. +-- +-- >>> let tasks' = MSeq.deleteAt 2 $ MSeq.cons (Task "D" 100) tasks +-- >>> tasks' +-- [Task "D" 100,Task "A" 50,Task "C" 60] +-- >>> MSeq.summary tasks' +-- Sum {getSum = 210} +-- +-- == Example 2: Max +-- +-- Consider that we now need the maximum cost instead of the sum, or both +-- sum and max. We need only change the @Measured@ instance to use another +-- @Semigroup@ that fits the requirement. +-- +-- @ +-- data SumMax = SumMax +-- { sum_ :: !Word +-- , max_ :: !Word +-- } deriving Show +-- +-- instance Semigroup SumMax where +-- SumMax sum1 max1 <> SumMax sum2 max2 = +-- SumMax (sum1+sum2) (max max1 max2) +-- +-- instance Measured Task where +-- type Measure Task = SumMax +-- measure (Task _ cost) = SumMax cost cost +-- @ +-- +-- We can see that it works as expected. +-- +-- >>> let tasks = MSeq.fromList [Task "A" 50, Task "B" 30, Task "C" 60] +-- >>> MSeq.summaryMay tasks +-- Just (SumMax {sum_ = 140, max_ = 60}) +-- +-- Note that we used 'Data.Seqn.MSeq.summaryMay' instead of @summary@, since we +-- did not define a monoid instance for @SumMax@. +-- +-- Aside: For the above scenario you may have considered using @(Sum Word, +-- t'Data.Monoid.Max' Word)@ as the measure, since the @Semigroup@ instance for +-- it is already defined. While that would work, it would be inefficient because +-- @(a,b)@ and its @(<>)@ implementation are lazy in @a@ and @b@. +-- +-- == Example 3: Binary search +-- +-- Consider that there are events where we unlock the ability to process tasks +-- with a total cost \(c\). To handle such events, we need to split out the +-- maximum number of tasks from the beginning of the sequence such that their +-- total does not exceed \(c\), and send them for processing. +-- +-- We can do this efficiently with an @MSeq@. The prefix sums of costs, which +-- is a component of our measure, forms a monotonic non-decreasing sequence. +-- We can take advantage of this and use binary search to find the point where +-- the sequence should be split. +-- +-- @ +-- splitAtMost :: Word -> MSeq Task -> Maybe (MSeq Task, MSeq Task) +-- splitAtMost c tasks = +-- case MSeq.'Data.Seqn.MSeq.binarySearchPrefix' (\\(SumMax c' _) -> c' > c) tasks of +-- (Nothing, _) -> Nothing -- c is too small for even the first task +-- (Just i, _) -> Just $! MSeq.'Data.Seqn.MSeq.splitAt' (i+1) tasks +-- @ +-- +-- >>> let tasks = MSeq.fromList [Task "A" 50, Task "B" 30, Task "C" 60] +-- >>> splitAtMost 100 tasks +-- Just ([Task "A" 50,Task "B" 30],[Task "C" 60]) +-- >>> splitAtMost 10 tasks +-- Nothing +-- +-- Note that the running time of @splitAtMost@ is simply \(O(\log n)\), and not +-- dependent on how many tasks are split out. +-- +-- == More information +-- +-- More uses of measured sequences can be found in the paper on finger trees: +-- +-- * Ralf Hinze and Ross Paterson, +-- /\"Finger trees: a simple general-purpose data structure\"/, +-- Journal of Functional Programming 16(2), 197-217, 2006, +-- https://doi.org/10.1017/S0956796805005769 +-- +-- One such use, priority queues, is implemented in this package and can be +-- found in the module "Data.Seqn.PQueue". diff --git a/src/Data/Seqn/PQueue.hs b/src/Data/Seqn/PQueue.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/PQueue.hs @@ -0,0 +1,43 @@ +-- | +-- = Priority queues +-- +-- @PQueue@ is a minimum priority queue implemented using an +-- t'Data.Seqn.MSeq.MSeq'. +-- +-- * It is spine-strict, and can contain only a finite number of elements. +-- * It is value-strict. It is guaranteed that if a @PQueue@ is in +-- [weak head normal form](https://wiki.haskell.org/Weak_head_normal_form) +-- (WHNF), every element of the @PQueue@ is also in WHNF. +-- * It maintains insertion order. If two elements compare equal, the one +-- which was inserted first will be removed first. Elements can also be +-- folded over in insertion order. +-- * It is a mergeable priority queue. Two queues can be concatenated +-- efficiently in logarithmic time. +-- +-- It is recommended to import this module qualified to avoid name clashes. +-- +-- @ +-- import Data.Seqn.PQueue (PQueue) +-- import qualified Data.Seqn.PQueue as PQueue +-- @ +-- +module Data.Seqn.PQueue + ( + -- * PQueue + P.PQueue + , P.empty + , P.singleton + , P.fromList + , P.concatMap + , P.insert + , P.min + , P.minView + , P.toSortedList + + -- * Entry + , P.Entry(..) + , P.entryPrio + , P.entryValue + ) where + +import qualified Data.Seqn.Internal.PQueue as P diff --git a/src/Data/Seqn/Seq.hs b/src/Data/Seqn/Seq.hs new file mode 100644 --- /dev/null +++ b/src/Data/Seqn/Seq.hs @@ -0,0 +1,149 @@ +-- | +-- = Finite sequences +-- +-- A value of type @Seq a@ is a sequence with elements of type @a@. +-- A @Seq@ is +-- +-- * Spine-strict, hence finite. @Seq@ cannot represent infinite sequences. +-- * Value-strict. It is guaranteed that if a @Seq@ is in +-- [weak head normal form](https://wiki.haskell.org/Weak_head_normal_form) +-- (WHNF), every element of the @Seq@ is also in WHNF. +-- +-- It is recommended to import this module qualified to avoid name clashes. +-- +-- @ +-- import Data.Seqn.Seq (Seq) +-- import qualified Data.Seqn.Seq as Seq +-- @ +-- +-- === Warning +-- +-- The length of a @Seq@ must not exceed @(maxBound \`div\` 3) :: Int@. If this +-- length is exceeded, the behavior of a @Seq@ is undefined. This value is very +-- large in practice, greater than \(7 \cdot 10^8\) on 32-bit systems and +-- \(3 \cdot 10^{18}\) on 64-bit systems. +-- +-- === Implementation +-- +-- @Seq@ is implemented as a +-- [weight-balanced binary tree](https://en.wikipedia.org/wiki/Weight-balanced_tree). +-- This structure is described by +-- +-- * J. Nievergelt and E. M. Reingold, +-- /\"Binary search trees of bounded balance\"/, +-- SIAM Journal of Computing 2(1), 1973, +-- https://doi.org/10.1137/0202005 +-- +-- * Stephen Adams, +-- /\"Efficient sets—a balancing act\"/, +-- Journal of Functional Programming 3(4), 553-561, 1993, +-- https://doi.org/10.1017/S0956796800000885 +-- +-- * Yoichi Hirai and Kazuhiko Yamamoto, +-- /\"Balancing weight-balanced trees\"/, +-- Journal of Functional Programming 21(3), 287-307, 2011, +-- https://doi.org/10.1017/S0956796811000104 +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- /\"Parallel Ordered Sets Using Join\"/, 2016, +-- https://doi.org/10.48550/arXiv.1602.02120 +-- +module Data.Seqn.Seq + ( + -- * Seq + S.Seq + + -- * Construct + , S.empty + , S.singleton + , S.fromList + , S.fromRevList + , S.replicate + , S.replicateA + , S.generate + , S.generateA + , S.unfoldr + , S.unfoldl + , S.unfoldrM + , S.unfoldlM + , S.concatMap + + -- * Convert + , S.toRevList + + -- * Index + , S.lookup + , S.index + , (S.!?) + , (S.!) + , S.update + , S.adjust + , S.insertAt + , S.deleteAt + + -- * Slice + , S.cons + , S.snoc + , S.uncons + , S.unsnoc + , S.take + , S.drop + , S.slice + , S.splitAt + , S.takeEnd + , S.dropEnd + , S.splitAtEnd + , S.tails + , S.inits + , S.chunksOf + + -- * Filter + , S.filter + , S.catMaybes + , S.mapMaybe + , S.mapEither + , S.filterA + , S.mapMaybeA + , S.mapEitherA + , S.takeWhile + , S.dropWhile + , S.span + , S.break + , S.takeWhileEnd + , S.dropWhileEnd + , S.spanEnd + , S.breakEnd + + -- * Transform + , S.reverse + , S.intersperse + , S.scanl + , S.scanr + , S.sort + , S.sortBy + + -- * Search and test + , S.findEnd + , S.findIndex + , S.findIndexEnd + , S.infixIndices + , S.binarySearchFind + , S.isPrefixOf + , S.isSuffixOf + , S.isInfixOf + , S.isSubsequenceOf + + -- * Zip and unzip + , S.zip + , S.zip3 + , S.zipWith + , S.zipWith3 + , S.zipWithM + , S.zipWith3M + , S.unzip + , S.unzip3 + , S.unzipWith + , S.unzipWith3 + ) where + +import qualified Data.Seqn.Internal.Seq as S diff --git a/test/ListExtra.hs b/test/ListExtra.hs new file mode 100644 --- /dev/null +++ b/test/ListExtra.hs @@ -0,0 +1,51 @@ +module ListExtra where + +import qualified Data.List as L + +unfoldlL :: (a -> Maybe (a, b)) -> a -> [b] +unfoldlL f = go [] + where + go xs z = case f z of + Nothing -> xs + Just (z',x) -> go (x:xs) z' + +-- In base since 4.19 +lookupL :: Int -> [a] -> Maybe a +lookupL i xs = + foldr (\x k j -> if j == 0 then Just x else k (j-1)) (const Nothing) xs i + +-- In base since 4.19 +unsnocL :: [a] -> Maybe ([a], a) +unsnocL = + foldr (\x -> Just . maybe ([], x) (\(~(a, b)) -> (x : a, b))) Nothing + +takeEndL :: Int -> [a] -> [a] +takeEndL n xs = foldr (\_ k -> k . tail) id (L.drop n xs) xs + +dropEndL :: Int -> [a] -> [a] +dropEndL n xs = L.zipWith const xs (L.drop n xs) + +adjustL :: (a -> a) -> Int -> [a] -> [a] +adjustL f i xs + | i < 0 = xs + | (l,x:r) <- L.splitAt i xs = l ++ f x : r + | otherwise = xs + +insertAtL :: Int -> a -> [a] -> [a] +insertAtL i x xs = L.take i xs ++ [x] ++ L.drop i xs + +deleteAtL :: Int -> [a] -> [a] +deleteAtL i xs = L.take i xs ++ L.drop (i+1) xs + +spanEndL :: (a -> Bool) -> [a] -> ([a], [a]) +spanEndL f xs = case L.span f (L.reverse xs) of + (ys,zs) -> (L.reverse zs, L.reverse ys) + +infixIndicesL :: Eq a => [a] -> [a] -> [Int] +infixIndicesL xs ys = + [i | (i,ys') <- L.zip [0..] (L.tails ys), xs `L.isPrefixOf` ys'] + +chunksOfL :: Int -> [a] -> [[a]] +chunksOfL c xs = case L.splitAt c xs of + ([], _) -> [] + (ys, zs) -> ys : chunksOfL c zs diff --git a/test/ListLikeTests.hs b/test/ListLikeTests.hs new file mode 100644 --- /dev/null +++ b/test/ListLikeTests.hs @@ -0,0 +1,561 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ImpredicativeTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +module ListLikeTests where + +import Prelude hiding ((<>), break, concatMap, drop, dropWhile, filter, liftA2, lookup, map, read, replicate, reverse, show, span, splitAt, take, takeWhile) +import qualified Control.Applicative as Ap +import Control.Monad.Trans.State (runState, state) +import Data.Bifunctor (Bifunctor(..)) +import qualified Data.Either as Either +import qualified Data.Foldable as F +import qualified Data.Foldable.WithIndex as IFo +import qualified Data.Functor.WithIndex as IFu +import qualified Data.List as L +import qualified Data.Maybe as Maybe +import qualified Data.Monoid as Monoid +import Data.Ord (comparing) +import qualified Data.Semigroup as Semigroup +import Test.Tasty +import Test.Tasty.QuickCheck hiding (generate) +import qualified Text.Read as Read +import qualified Text.Show as Show + +import ListExtra +import TestUtil (MaybeBottom(..), ListLike(..), isBottom, eval) + +-- This is defined for convenience. Not all tests require all of these +-- constraints. +type Common t = + ( ListLike t + , Eq t + , Show t + , Arbitrary t + , Eq (E t) + , Show (E t) + , Arbitrary (E t) + ) +type FuncE t = + ( CoArbitrary (E t) + , Function (E t) + ) + +listLike :: forall t. Common t => TestTree +listLike = testGroup "ListLike" + [ testProperty "toList . fromList == id" $ \xs -> + toL (fromL xs :: t) === xs + , testProperty "fromList . toList == id" $ \(xs :: t) -> + fromL (toL xs) === xs + ] + +-------------- +-- Construct +-------------- + +empty :: ListLike t => t -> TestTree +empty e = testProperty "empty" $ null (toL e) + +singleton :: Common t => (E t -> t) -> TestTree +singleton f = testProperty "singleton" $ \x -> toL (f x) === [x] + +fromRevList :: Common t => ([E t] -> t) -> TestTree +fromRevList f = testProperty "fromRevList" $ \xs -> toL (f xs) === L.reverse xs + +replicate :: Common t => (Int -> E t -> t) -> TestTree +replicate f = testProperty "replicate" $ \n x -> toL (f n x) === L.replicate n x + +replicateA :: Common t => (forall f. Applicative f => Int -> f (E t) -> f t) -> TestTree +replicateA f = testProperty "replicateA" $ \xs -> + let n = length xs in + runState (fmap toL (f n (state (\(y:ys) -> (y,ys))))) xs === (xs, []) + +generate :: Common t => (Int -> (Int -> E t) -> t) -> TestTree +generate f = testProperty "generate" $ \n fn -> + toL (f n (applyFun fn)) === L.map (applyFun fn) [0..n-1] + +generateA :: Common t => (forall f. Applicative f => Int -> (Int -> f (E t)) -> f t) -> TestTree +generateA f = testProperty "generateA" $ \xs -> + let n = length xs in + fmap toL (f n (\i -> ([i], xs !! i))) === ([0..n-1], xs) + +unfoldr :: Common t => (([a] -> Maybe (a, [a])) -> [E t] -> t) -> TestTree +unfoldr f = testProperty "unfoldr" $ \xs -> toL (f L.uncons xs) === xs + +unfoldl :: Common t => (([a] -> Maybe ([a], a)) -> [E t] -> t) -> TestTree +unfoldl f = testProperty "unfoldl" $ \xs -> toL (f unsnocL xs) === xs + +unfoldrM :: Common t => (forall b m. Monad m => (b -> m (Maybe (E t, b))) -> b -> m t) -> TestTree +unfoldrM f = testProperty "unfoldrM" $ \xs -> + fmap toL (f (\ys -> ([ys], L.uncons ys)) xs) === (L.tails xs, xs) + +unfoldlM :: Common t => (forall b m. Monad m => (b -> m (Maybe (b, E t))) -> b -> m t) -> TestTree +unfoldlM f = testProperty "unfoldlM" $ \xs -> + fmap toL (f (\ys -> ([ys], unsnocL ys)) xs) === + (L.reverse (L.inits xs), xs) + +(<>) :: Common t => (t -> t -> t) -> TestTree +(<>) f = testProperty "<>" $ \xs ys -> + toL (f xs ys) === toL xs Semigroup.<> toL ys + +stimes :: Common t => (Int -> t -> t) -> TestTree +stimes f = testProperty "stimes" $ \(n :: Int) xs -> + toL (f n xs) === Semigroup.stimes (max 0 n) (toL xs) + +mconcat :: Common t => ([t] -> t) -> TestTree +mconcat f = testProperty "mconcat" $ \xss -> + toL (f xss) === Monoid.mconcat (fmap toL xss) + +concatMap :: Common t => (forall f a. Foldable f => (a -> t) -> f a -> t) -> TestTree +concatMap f = testProperty "concatMap" $ \fn (xs :: [Int]) -> + toL (f (applyFun fn) xs) === L.concatMap (toL . applyFun fn) xs + +read :: forall t. (Common t, Read t, Read (E t)) => TestTree +read = testProperty "read" $ \s -> case s of + Left str -> + fmap toL (Read.readMaybe str :: Maybe t) === Read.readMaybe str + Right (xs :: [E t]) -> + fmap toL (Read.readMaybe (Show.show xs) :: Maybe t) === Just xs + +------------ +-- Convert +------------ + +toRevList :: Common t => (t -> [E t]) -> TestTree +toRevList f = testProperty "toRevList" $ \xs -> f xs === L.reverse (toL xs) + +show :: forall t. Common t => TestTree +show = testProperty "show" $ \(xs :: t) -> Show.show xs === Show.show (toL xs) + +--------- +-- Fold +--------- + +-- Note: For the foldr/foldl family, strictness is checked in addition to +-- correctness. For the foldMap family strictness is /not checked/ because +-- these usually follow the internal structure of the Foldable. This results in +-- different strictness in different places compared to a list. + +foldMap :: Common t => (forall m. Monoid m => (E t -> m) -> t -> m) -> TestTree +foldMap f = testProperty "foldMap" $ \xs -> f (:[]) xs === toL xs + +foldMap' :: Common t => (forall m. Monoid m => (E t -> m) -> t -> m) -> TestTree +foldMap' f = testProperty "foldMap'" $ \xs -> f (:[]) xs === toL xs + +foldr :: (Common t, FuncE t) => (forall b. (E t -> b -> b) -> b -> t -> b) -> TestTree +foldr f = testProperty "foldr" $ \xs fn (z0 :: Int) -> + let g x z = case applyFun fn x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (L.foldr g z0) xs + +foldl' :: (Common t, FuncE t) => (forall b. (b -> E t -> b) -> b -> t -> b) -> TestTree +foldl' f = testProperty "foldl'" $ \xs fn (z0 :: Int) -> + let g z x = case applyFun fn x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (L.foldl' g z0) xs + +foldr' :: (Common t, FuncE t) => (forall b. (E t -> b -> b) -> b -> t -> b) -> TestTree +foldr' f = testProperty "foldr'" $ \xs fn (z0 :: Int) -> + let g x z = case applyFun fn x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (F.foldr' g z0) xs + +foldl :: (Common t, FuncE t) => (forall b. (b -> E t -> b) -> b -> t -> b) -> TestTree +foldl f = testProperty "foldl" $ \xs fn (z0 :: Int) -> + let g z x = case applyFun fn x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (L.foldl g z0) xs + +ifoldMap :: Common t => (forall m. Monoid m => (Int -> E t -> m) -> t -> m) -> TestTree +ifoldMap f = testProperty "ifoldMap" $ \xs -> + f (\i x -> [(i,x)]) xs === L.zip [0..] (toL xs) + +ifoldMap' :: Common t => (forall m. Monoid m => (Int -> E t -> m) -> t -> m) -> TestTree +ifoldMap' f = testProperty "ifoldMap'" $ \xs -> + f (\i x -> [(i,x)]) xs === L.zip [0..] (toL xs) + +ifoldr :: (Common t, FuncE t) => (forall b. (Int -> E t -> b -> b) -> b -> t -> b) -> TestTree +ifoldr f = testProperty "ifoldr" $ \xs fn (z0 :: Int) -> + let g i x z = case applyFun2 fn i x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (IFo.ifoldr g z0) xs + +ifoldl :: (Common t, FuncE t) => (forall b. (Int -> b -> E t -> b) -> b -> t -> b) -> TestTree +ifoldl f = testProperty "ifoldl" $ \xs fn (z0 :: Int) -> + let g i z x = case applyFun2 fn i x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (IFo.ifoldl g z0) xs + +ifoldr' :: (Common t, FuncE t) => (forall b. (Int -> E t -> b -> b) -> b -> t -> b) -> TestTree +ifoldr' f = testProperty "ifoldr'" $ \xs fn (z0 :: Int) -> + let g i x z = case applyFun2 fn i x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (IFo.ifoldr' g z0) xs + +ifoldl' :: (Common t, FuncE t) => (forall b. (Int -> b -> E t -> b) -> b -> t -> b) -> TestTree +ifoldl' f = testProperty "ifoldl" $ \xs fn (z0 :: Int) -> + let g i z x = case applyFun2 fn i x of + Left h -> applyFun h z + Right Bottom -> error "bottom" + Right (NotBottom z') -> z' + in testFold (f g z0) (IFo.ifoldl' g z0) xs + +testFold + :: (Common t, Eq a, Show a) => (t -> a) -> ([E t] -> a) -> t -> Property +testFold f1 f2 xs = ioProperty $ do + r1 <- eval (f1 xs) + r2 <- eval (f2 (toL xs)) + pure $ + classify (isBottom r2) "bottom" $ + r1 === r2 + +---------- +-- Index +---------- + +lookup :: Common t => (Int -> t -> Maybe (E t)) -> TestTree +lookup f = testProperty "lookup" $ \i xs -> f i xs === lookupL i (toL xs) + +index :: Common t => (Int -> t -> E t) -> TestTree +index f = testProperty "index" $ \i xs -> + 0 <= i && i < length (toL xs) ==> + f i xs === toL xs !! i + +update :: Common t => (Int -> E t -> t -> t) -> TestTree +update f = testProperty "update" $ \i x xs -> + toL (f i x xs) === adjustL (\_ -> x) i (toL xs) + +adjust :: (Common t, FuncE t) => ((E t -> E t) -> Int -> t -> t) -> TestTree +adjust f = testProperty "adjust" $ \fn i xs -> + toL (f (applyFun fn) i xs) === adjustL (applyFun fn) i (toL xs) + +insertAt :: Common t => (Int -> E t -> t -> t) -> TestTree +insertAt f = testProperty "insertAt" $ \i x xs -> + toL (f i x xs) === insertAtL i x (toL xs) + +deleteAt :: Common t => (Int -> t -> t) -> TestTree +deleteAt f = testProperty "deleteAt" $ \i xs -> + toL (f i xs) === deleteAtL i (toL xs) + +---------- +-- Slice +---------- + +cons :: Common t => (E t -> t -> t) -> TestTree +cons f = testProperty "cons" $ \x xs -> toL (f x xs) === x : toL xs + +snoc :: Common t => (t -> E t -> t) -> TestTree +snoc f = testProperty "snoc" $ \xs x -> toL (f xs x) === toL xs ++ [x] + +uncons :: Common t => (t -> Maybe (E t, t)) -> TestTree +uncons f = testProperty "uncons" $ \xs -> + fmap (fmap toL) (f xs) === L.uncons (toL xs) + +unsnoc :: Common t => (t -> Maybe (t, E t)) -> TestTree +unsnoc f = testProperty "unsnoc" $ \xs -> + fmap (first toL) (f xs) === unsnocL (toL xs) + +take :: Common t => (Int -> t -> t) -> TestTree +take f = testProperty "take" $ \n xs -> toL (f n xs) === L.take n (toL xs) + +drop :: Common t => (Int -> t -> t) -> TestTree +drop f = testProperty "drop" $ \n xs -> toL (f n xs) === L.drop n (toL xs) + +slice :: Common t => ((Int, Int) -> t -> t) -> TestTree +slice f = testProperty "slice" $ \(i,j) xs -> + toL (f (i,j) xs) === L.drop i (L.take (j+1) (toL xs)) + +splitAt :: Common t => (Int -> t -> (t,t)) -> TestTree +splitAt f = testProperty "splitAt" $ \n xs -> + bimap toL toL (f n xs) === L.splitAt n (toL xs) + +takeEnd :: Common t => (Int -> t -> t) -> TestTree +takeEnd f = testProperty "takeEnd" $ \n xs -> + toL (f n xs) === takeEndL n (toL xs) + +dropEnd :: Common t => (Int -> t -> t) -> TestTree +dropEnd f = testProperty "dropEnd" $ \n xs -> + toL (f n xs) === dropEndL n (toL xs) + +splitAtEnd :: Common t => (Int -> t -> (t,t)) -> TestTree +splitAtEnd f = testProperty "splitAtEnd" $ \n xs -> + bimap toL toL (f n xs) === + (dropEndL n (toL xs), takeEndL n (toL xs)) + +tails :: (Common t, ListLike t2, t ~ E t2) => (t -> t2) -> TestTree +tails f = testProperty "tails" $ \xs -> + L.map toL (toL (f xs)) === L.tails (toL xs) + +inits :: (Common t, ListLike t2, t ~ E t2) => (t -> t2) -> TestTree +inits f = testProperty "inits" $ \xs -> + L.map toL (toL (f xs)) === L.inits (toL xs) + +chunksOf :: (Common t, ListLike t2, t ~ E t2) => (Int -> t -> t2) -> TestTree +chunksOf f = testProperty "chunksOf" $ \c xs -> + L.map toL (toL (f c xs)) === chunksOfL c (toL xs) + +----------- +-- Filter +----------- + +filter :: (Common t, FuncE t) => ((E t -> Bool) -> t -> t) -> TestTree +filter f = testProperty "filter" $ \fn xs -> + toL (f (applyFun fn) xs) === L.filter (applyFun fn) (toL xs) + +mapMaybe :: (Common t1, FuncE t1, Common t2) => ((E t1 -> Maybe (E t2)) -> t1 -> t2) -> TestTree +mapMaybe f = testProperty "mapMaybe" $ \fn xs -> + toL (f (applyFun fn) xs) === + Maybe.mapMaybe (applyFun fn) (toL xs) + +mapEither :: (Common t1, FuncE t1, Common t2, Common t3) => ((E t1 -> Either (E t2) (E t3)) -> t1 -> (t2, t3)) -> TestTree +mapEither f = testProperty "mapEither" $ \fn xs -> + bimap toL toL (f (applyFun fn) xs) === + Either.partitionEithers (fmap (applyFun fn) (toL xs)) + +filterM :: (Common t, FuncE t) => (forall m. Monad m => (E t -> m Bool) -> t -> m t) -> TestTree +filterM f = testProperty "filter" $ \fn xs -> + fmap toL (f (\x -> ([x], applyFun fn x)) xs) === + (toL xs, L.filter (applyFun fn) (toL xs)) + +mapMaybeM :: (Common t1, FuncE t1, Common t2) => (forall m. Monad m => (E t1 -> m (Maybe (E t2))) -> t1 -> m t2) -> TestTree +mapMaybeM f = testProperty "mapMaybeM" $ \fn xs -> + fmap toL (f (\x -> ([x], applyFun fn x)) xs) === + (toL xs, Maybe.mapMaybe (applyFun fn) (toL xs)) + +mapEitherM :: (Common t1, FuncE t1, Common t2, Common t3) => (forall m. Monad m => (E t1 -> m (Either (E t2) (E t3))) -> t1 -> m (t2, t3)) -> TestTree +mapEitherM f = testProperty "mapEitherM" $ \fn xs -> + fmap (bimap toL toL) (f (\x -> ([x], applyFun fn x)) xs) === + (toL xs, Either.partitionEithers (fmap (applyFun fn) (toL xs))) + +takeWhile :: (Common t, FuncE t) => ((E t -> Bool) -> t -> t) -> TestTree +takeWhile f = testProperty "takeWhile" $ \fn xs -> + toL (f (applyFun fn) xs) === L.takeWhile (applyFun fn) (toL xs) + +dropWhile :: (Common t, FuncE t) => ((E t -> Bool) -> t -> t) -> TestTree +dropWhile f = testProperty "dropWhile" $ \fn xs -> + toL (f (applyFun fn) xs) === L.dropWhile (applyFun fn) (toL xs) + +span :: (Common t, FuncE t) => ((E t -> Bool) -> t -> (t,t)) -> TestTree +span f = testProperty "span" $ \fn xs -> + bimap toL toL (f (applyFun fn) xs) === L.span (applyFun fn) (toL xs) + +break :: (Common t, FuncE t) => ((E t -> Bool) -> t -> (t,t)) -> TestTree +break f = testProperty "break" $ \fn xs -> + bimap toL toL (f (applyFun fn) xs) === L.break (applyFun fn) (toL xs) + +takeWhileEnd :: (Common t, FuncE t) => ((E t -> Bool) -> t -> t) -> TestTree +takeWhileEnd f = testProperty "takeWhileEnd" $ \fn xs -> + toL (f (applyFun fn) xs) === + snd (spanEndL (applyFun fn) (toL xs)) + +dropWhileEnd :: (Common t, FuncE t) => ((E t -> Bool) -> t -> t) -> TestTree +dropWhileEnd f = testProperty "dropWhileEnd" $ \fn xs -> + toL (f (applyFun fn) xs) === + fst (spanEndL (applyFun fn) (toL xs)) + +spanEnd :: (Common t, FuncE t) => ((E t -> Bool) -> t -> (t,t)) -> TestTree +spanEnd f = testProperty "spanEnd" $ \fn xs -> + bimap toL toL (f (applyFun fn) xs) === + spanEndL (applyFun fn) (toL xs) + +breakEnd :: (Common t, FuncE t) => ((E t -> Bool) -> t -> (t,t)) -> TestTree +breakEnd f = testProperty "breakEnd" $ \fn xs -> + bimap toL toL (f (applyFun fn) xs) === + spanEndL (not . applyFun fn) (toL xs) + +-------------- +-- Transform +-------------- + +map :: (Common t1, FuncE t1, Common t2) => ((E t1 -> E t2) -> t1 -> t2) -> TestTree +map f = testProperty "map" $ \fn xs -> + toL (f (applyFun fn) xs) === L.map (applyFun fn) (toL xs) + +liftA2 :: (Common t1, FuncE t1, Common t2, FuncE t2, Common t3) => ((E t1 -> E t2 -> E t3) -> t1 -> t2 -> t3) -> TestTree +liftA2 f = testProperty "liftA2" $ + \fn xs ys -> + toL (f (applyFun2 fn) xs ys) === + Ap.liftA2 (applyFun2 fn) (toL xs) (toL ys) + +(<*) :: (Common t1, Common t2) => (t1 -> t2 -> t1) -> TestTree +(<*) f = testProperty "<*" $ \xs ys -> + toL (f xs ys) === (toL xs Ap.<* toL ys) + +(*>) :: (Common t1, Common t2) => (t1 -> t2 -> t2) -> TestTree +(*>) f = testProperty "*>" $ \xs ys -> + toL (f xs ys) === (toL xs Ap.*> toL ys) + +bind :: (Common t1, FuncE t1, Common t2) => (t1 -> (E t1 -> t2) -> t2) -> TestTree +bind f = testProperty ">>=" $ \xs fn -> + toL (f xs (applyFun fn)) === (toL xs >>= toL . applyFun fn) + +traverse :: (Common t1, FuncE t1, Common t2) => (forall f. Applicative f => (E t1 -> f (E t2)) -> t1 -> f t2) -> TestTree +traverse f = testProperty "traverse" $ \fn xs -> + fmap toL (f (\x -> ([x], applyFun fn x)) xs) === + (toL xs, fmap (applyFun fn) (toL xs)) + +imap :: (Common t1, FuncE t1, Common t2) => ((Int -> E t1 -> E t2) -> t1 -> t2) -> TestTree +imap f = testProperty "imap" $ \fn xs -> + toL (f (applyFun2 fn) xs) === IFu.imap (applyFun2 fn) (toL xs) + +itraverse :: (Common t1, FuncE t1, Common t2) => (forall f. Applicative f => (Int -> E t1 -> f (E t2)) -> t1 -> f t2) -> TestTree +itraverse f = testProperty "itraverse" $ \fn xs -> + let ixs = L.zip [0..] (toL xs) in + fmap toL (f (\i x -> ([(i,x)], applyFun2 fn i x)) xs) === + (ixs, L.map (uncurry (applyFun2 fn)) ixs) + +reverse :: Common t => (t -> t) -> TestTree +reverse f = testProperty "reverse" $ \xs -> + toL (f xs) === L.reverse (toL xs) + +intersperse :: Common t => (E t -> t -> t) -> TestTree +intersperse f = testProperty "intersperse" $ \x xs -> + toL (f x xs) === L.intersperse x (toL xs) + +scanl :: (Common t1, FuncE t1, Common t2, FuncE t2) => ((E t2 -> E t1 -> E t2) -> E t2 -> t1 -> t2) -> TestTree +scanl f = testProperty "scanl" $ \fn z xs -> + toL (f (applyFun2 fn) z xs) === L.scanl' (applyFun2 fn) z (toL xs) + +scanr :: (Common t1, FuncE t1, Common t2, FuncE t2) => ((E t1 -> E t2 -> E t2) -> E t2 -> t1 -> t2) -> TestTree +scanr f = testProperty "scanr" $ \fn z xs -> + toL (f (applyFun2 fn) z xs) === L.scanr (applyFun2 fn) z (toL xs) + +sort :: (Common t, Ord (E t)) => (t -> t) -> TestTree +sort f = testProperty "sort" $ \xs -> + toL (f xs) === L.sort (toL xs) + +sortBy :: forall t. (Common t, FuncE t) => ((E t -> E t -> Ordering) -> t -> t) -> TestTree +sortBy f = testProperty "sortBy" $ \(fn :: Fun (E t) Int) xs -> + toL (f (comparing (applyFun fn)) xs) === + L.sortBy (comparing (applyFun fn)) (toL xs) + +-------------------- +-- Search and test +-------------------- + +eq :: Common t => (t -> t -> Bool) -> TestTree +eq f = testProperty "==" $ \xs ys -> + let tru = toL xs == toL ys in + classify (tru && not (null (toL xs))) "non-trivial yes" $ + f xs ys === tru + +cmp :: (Common t, Ord (E t)) => (t -> t -> Ordering) -> TestTree +cmp f = testProperty "compare" $ \xs ys -> + let res = compare (toL xs) (toL ys) in + collect res $ + f xs ys === res + +findEnd :: (Common t, FuncE t) => ((E t -> Bool) -> t -> Maybe (E t)) -> TestTree +findEnd f = testProperty "findEnd" $ \fn xs -> + f (applyFun fn) xs === L.find (applyFun fn) (L.reverse (toL xs)) + +findIndex :: (Common t, FuncE t) => ((E t -> Bool) -> t -> Maybe Int) -> TestTree +findIndex f = testProperty "findIndex" $ \fn xs -> + f (applyFun fn) xs === L.findIndex (applyFun fn) (toL xs) + +findIndexEnd :: (Common t, FuncE t) => ((E t -> Bool) -> t -> Maybe Int) -> TestTree +findIndexEnd f = testProperty "findIndexEnd" $ \fn xs -> + f (applyFun fn) xs === + fmap snd (unsnocL (L.findIndices (applyFun fn) (toL xs))) + +infixIndices :: Common t => (t -> t -> [Int]) -> TestTree +infixIndices f = testProperty "infixIndices" $ \xs ys -> + let res = infixIndicesL (toL xs) (toL ys) in + classify (not (null res) && not (null (toL xs))) "non-trivial yes" $ + f xs ys === res + +binarySearchFind :: (Common t, Ord (E t)) => ((E t -> Ordering) -> t -> Maybe (E t)) -> TestTree +binarySearchFind f = testProperty "binarySearchFind" $ \(Sorted xs) x -> + f (`compare` x) (fromL xs) === L.find (x==) xs + +isPrefixOf :: Common t => (t -> t -> Bool) -> TestTree +isPrefixOf f = testProperty "isPrefixOf" $ \xs ys -> + let tru = toL xs `L.isPrefixOf` toL ys in + classify (tru && not (null (toL xs))) "non-trivial yes" $ + f xs ys === tru + +isSuffixOf :: Common t => (t -> t -> Bool) -> TestTree +isSuffixOf f = testProperty "isSuffixOf" $ \xs ys -> + let tru = toL xs `L.isSuffixOf` toL ys in + classify (tru && not (null (toL xs))) "non-trivial yes" $ + f xs ys === tru + +isInfixOf :: Common t => (t -> t -> Bool) -> TestTree +isInfixOf f = testProperty "isInfixOf" $ \xs ys -> + let tru = toL xs `L.isInfixOf` toL ys in + classify (tru && not (null (toL xs))) "non-trivial yes" $ + f xs ys === tru + +isSubsequenceOf :: Common t => (t -> t -> Bool) -> TestTree +isSubsequenceOf f = testProperty "isSubsequenceOf" $ \xs ys -> + let tru = toL xs `L.isSubsequenceOf` toL ys in + classify (tru && not (null (toL xs))) "non-trivial yes" $ + f xs ys === tru + +------------------ +-- Zip and unzip +------------------ + +zip :: (Common t1, Common t2, Common t3, E t3 ~ (E t1, E t2)) => (t1 -> t2 -> t3) -> TestTree +zip f = testProperty "zip" $ \xs ys -> toL (f xs ys) === L.zip (toL xs) (toL ys) + +zip3 :: (Common t1, Common t2, Common t3, Common t4, E t4 ~ (E t1, E t2, E t3)) => (t1 -> t2 -> t3 -> t4) -> TestTree +zip3 f = testProperty "zip3" $ \xs ys zs -> + toL (f xs ys zs) === L.zip3 (toL xs) (toL ys) (toL zs) + +zipWith :: (Common t1, FuncE t1, Common t2, FuncE t2, Common t3) => ((E t1 -> E t2 -> E t3) -> t1 -> t2 -> t3) -> TestTree +zipWith f = testProperty "zipWith" $ \fn xs ys -> + toL (f (applyFun2 fn) xs ys) === + L.zipWith (applyFun2 fn) (toL xs) (toL ys) + +zipWith3 :: (Common t1, FuncE t1, Common t2, FuncE t2, Common t3, FuncE t3, Common t4) => ((E t1 -> E t2 -> E t3 -> E t4) -> t1 -> t2 -> t3 -> t4) -> TestTree +zipWith3 f = testProperty "zipWith3" $ \fn xs ys zs -> + toL (f (applyFun3 fn) xs ys zs) === + L.zipWith3 (applyFun3 fn) (toL xs) (toL ys) (toL zs) + +zipWithM :: (Common t1, FuncE t1, Common t2, FuncE t2, Common t3) => (forall m. Monad m => (E t1 -> E t2 -> m (E t3)) -> t1 -> t2 -> m t3) -> TestTree +zipWithM f = testProperty "zipWith" $ \fn xs ys -> + fmap toL (f (\x y -> ([(x,y)], applyFun2 fn x y)) xs ys) === + (L.zip (toL xs) (toL ys), L.zipWith (applyFun2 fn) (toL xs) (toL ys)) + +zipWith3M :: (Common t1, FuncE t1, Common t2, FuncE t2, Common t3, FuncE t3, Common t4) => (forall m. Monad m => (E t1 -> E t2 -> E t3 -> m (E t4)) -> t1 -> t2 -> t3 -> m t4) -> TestTree +zipWith3M f = testProperty "zipWith" $ \fn xs ys zs -> + fmap toL (f (\x y z -> ([(x,y,z)], applyFun3 fn x y z)) xs ys zs) === + ( L.zip3 (toL xs) (toL ys) (toL zs) + , L.zipWith3 (applyFun3 fn) (toL xs) (toL ys) (toL zs) + ) + +unzip :: (Common t1, Common t2, Common t3, E t1 ~ (E t2, E t3)) => (t1 -> (t2, t3)) -> TestTree +unzip f = testProperty "unzip" $ \xs -> + bimap toL toL (f xs) === L.unzip (toL xs) + +unzip3 :: (Common t1, Common t2, Common t3, Common t4, E t1 ~ (E t2, E t3, E t4)) => (t1 -> (t2, t3, t4)) -> TestTree +unzip3 f = testProperty "unzip3" $ \xs -> + (\(a,b,c) -> (toL a, toL b, toL c)) (f xs) === L.unzip3 (toL xs) + +unzipWith :: (Common t1, FuncE t1, Common t2, Common t3) => ((E t1 -> (E t2, E t3)) -> t1 -> (t2, t3)) -> TestTree +unzipWith f = testProperty "unzipWith" $ \fn xs -> + bimap toL toL (f (applyFun fn) xs) === + L.unzip (fmap (applyFun fn) (toL xs)) + +unzipWith3 :: (Common t1, FuncE t1, Common t2, Common t3, Common t4) => ((E t1 -> (E t2, E t3, E t4)) -> t1 -> (t2, t3, t4)) -> TestTree +unzipWith3 f = testProperty "unzipWith3" $ \fn xs -> + (\(as,bs,cs) -> (toL as, toL bs, toL cs)) (f (applyFun fn) xs) === + L.unzip3 (fmap (applyFun fn) (toL xs)) diff --git a/test/MSeq.hs b/test/MSeq.hs new file mode 100644 --- /dev/null +++ b/test/MSeq.hs @@ -0,0 +1,396 @@ +{-# OPTIONS_GHC -Wno-orphans #-} -- Arbitrary instances +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +module MSeq (mseqTests) where + +import Prelude hiding (break, concatMap, drop, dropWhile, filter, liftA2, lookup, map, replicate, reverse, splitAt, scanl, scanr, span, take, takeWhile, traverse, zipWith, zipWith3) +import Data.Coerce (coerce) +import qualified Control.Applicative as Ap +import qualified Data.Foldable as F +import qualified Data.Foldable.WithIndex as IFo +import Data.Functor.Identity (Identity(..)) +import qualified Data.List as L +import Data.Monoid (Sum(..)) +import Data.Proxy (Proxy(..)) +import Data.Semigroup (stimes) +import Test.Tasty +import Test.Tasty.QuickCheck hiding (generate) +import qualified Test.QuickCheck.Classes.Base as QLaws + +import Data.Seqn.MSeq +import qualified Data.Seqn.Internal.MSeq as MSeqInternal +import qualified Data.Seqn.Internal.MTree as MTreeInternal +import ListExtra (unsnocL) +import qualified ListLikeTests as LL +import TestUtil ((.:), ListLike(..), Sqrt1(..), tastyLaws) + +mseqTests :: TestTree +mseqTests = testGroup "Data.Seqn.MSeq" + [ testGroup "properties" + [ + LL.listLike @(MSeq A) + + -- Measured queries + , testProperty "summaryMay" $ \(xs :: MSeq A) -> + summaryMay xs === foldMap (Just . measure) (F.toList xs) + , testProperty "summary" $ \(xs :: MSeq D) -> + summary xs === foldMap measure (F.toList xs) + , testProperty "binarySearchPrefix" $ \(xs :: MSeq S) y -> + let p = (>=y) . getSum + xs' = F.toList xs + iws = + [ (i, foldMap measure (L.take (i+1) xs')) + | i <- [0 .. length xs - 1] + ] + lastFalse = snd <$> unsnocL [i | (i,w) <- iws, not (p w)] + firstTrue = fst <$> L.uncons [i | (i,w) <- iws, p w] + in binarySearchPrefix p xs === (lastFalse, firstTrue) + , testProperty "binarySearchSuffix" $ \(xs :: MSeq S) y -> + let p = (>=y) . getSum + xs' = F.toList xs + iws = + [ (i, foldMap measure (L.drop i xs')) + | i <- [0 .. length xs - 1] + ] + lastTrue = snd <$> unsnocL [i | (i,w) <- iws, p w] + firstFalse = fst <$> L.uncons [i | (i,w) <- iws, not (p w)] + in binarySearchSuffix p xs === (lastTrue, firstFalse) + + -- Construct + , LL.fromRevList @(MSeq A) fromRevList + , LL.empty @(MSeq A) empty + , LL.singleton @(MSeq A) singleton + , LL.replicate @(MSeq A) replicate + , LL.replicateA @(MSeq A) replicateA + , LL.generate @(MSeq A) generate + , LL.generateA @(MSeq A) generateA + , LL.unfoldr @(MSeq A) unfoldr + , LL.unfoldl @(MSeq A) unfoldl + , LL.unfoldrM @(MSeq A) unfoldrM + , LL.unfoldlM @(MSeq A) unfoldlM + , (LL.<>) @(MSeq A) (<>) + , LL.stimes @(MSeq A) stimes + , LL.mconcat @(MSeq A) mconcat + , LL.concatMap @(MSeq A) concatMap + , testProperty "mfix" $ \n -> + F.toList (mkMfix n) === + fmap (LI . L.replicate 10 . fromIntegral) [0 .. n-1] + , LL.read @(MSeq A) + + -- Convert + , LL.toRevList @(MSeq A) toRevList + , LL.show @(MSeq A) + + -- Fold + , LL.foldMap @(MSeq A) foldMap + , LL.foldMap' @(MSeq A) F.foldMap' + , LL.foldr @(MSeq A) foldr + , LL.foldl @(MSeq A) F.foldl + , LL.foldl' @(MSeq A) F.foldl' + , LL.foldr' @(MSeq A) F.foldr' + , LL.ifoldMap @(MSeq A) IFo.ifoldMap + , LL.ifoldr @(MSeq A) IFo.ifoldr + , LL.ifoldl @(MSeq A) IFo.ifoldl + , LL.ifoldr' @(MSeq A) IFo.ifoldr' + , LL.ifoldl' @(MSeq A) IFo.ifoldl' + + -- Index + , LL.lookup @(MSeq A) lookup + , LL.index @(MSeq A) index + , LL.update @(MSeq A) update + , LL.adjust @(MSeq A) adjust + , LL.insertAt @(MSeq A) insertAt + , LL.deleteAt @(MSeq A) deleteAt + + -- Slice + , LL.cons @(MSeq A) cons + , LL.snoc @(MSeq A) snoc + , LL.uncons @(MSeq A) uncons + , LL.unsnoc @(MSeq A) unsnoc + , LL.take @(MSeq A) take + , LL.drop @(MSeq A) drop + , LL.slice @(MSeq A) slice + , LL.splitAt @(MSeq A) splitAt + , LL.takeEnd @(MSeq A) takeEnd + , LL.dropEnd @(MSeq A) dropEnd + , LL.splitAtEnd @(MSeq A) splitAtEnd + + -- Filter + , LL.filter @(MSeq A) filter + , LL.mapMaybe @(MSeq A) @(MSeq B) mapMaybe + , LL.mapEither @(MSeq A) @(MSeq B) @(MSeq C) mapEither + , LL.filterM @(MSeq A) filterA + , LL.mapMaybeM @(MSeq A) @(MSeq B) mapMaybeA + , LL.mapEitherM @(MSeq A) @(MSeq B) @(MSeq C) mapEitherA + , LL.takeWhile @(MSeq A) takeWhile + , LL.dropWhile @(MSeq A) dropWhile + , LL.span @(MSeq A) span + , LL.break @(MSeq A) break + , LL.takeWhileEnd @(MSeq A) takeWhileEnd + , LL.dropWhileEnd @(MSeq A) dropWhileEnd + , LL.spanEnd @(MSeq A) spanEnd + , LL.breakEnd @(MSeq A) breakEnd + + -- Transform + , LL.map @(MSeq A) @(MSeq B) map + , LL.liftA2 @(Sqrt1 MSeq A) @(Sqrt1 MSeq B) @(Sqrt1 MSeq C) (coerce liftA2) + , LL.traverse @(MSeq A) @(MSeq B) traverse + , LL.imap @(MSeq A) @(MSeq B) imap + , LL.itraverse @(MSeq A) @(MSeq B) itraverse + , LL.reverse @(MSeq A) reverse + , LL.intersperse @(MSeq A) intersperse + , LL.scanl @(MSeq A) @(MSeq B) scanl + , LL.scanr @(MSeq A) @(MSeq B) scanr + , LL.sort @(MSeq A) sort + , LL.sortBy @(MSeq A) sortBy + + -- Search and test + , LL.eq @(MSeq A) (==) + , LL.cmp @(MSeq A) compare + , LL.findEnd @(MSeq A) findEnd + , LL.findIndex @(MSeq A) findIndex + , LL.findIndexEnd @(MSeq A) findIndexEnd + , LL.infixIndices @(MSeq A) infixIndices + , LL.binarySearchFind @(MSeq A) binarySearchFind + , LL.isPrefixOf @(MSeq A) isPrefixOf + , LL.isSuffixOf @(MSeq A) isSuffixOf + , LL.isInfixOf @(MSeq A) isInfixOf + , LL.isSubsequenceOf @(MSeq A) isSubsequenceOf + + -- Zip and unzip + , LL.zipWith @(MSeq A) @(MSeq B) @(MSeq C) zipWith + , LL.zipWith3 @(MSeq A) @(MSeq B) @(MSeq C) @(MSeq A) zipWith3 + , LL.zipWithM @(MSeq A) @(MSeq B) @(MSeq C) zipWithM + , LL.zipWith3M @(MSeq A) @(MSeq B) @(MSeq C) @(MSeq A) zipWith3M + , LL.unzipWith @(MSeq A) @(MSeq B) @(MSeq C) unzipWith + , LL.unzipWith3 @(MSeq A) @(MSeq B) @(MSeq C) @(MSeq A) unzipWith3 + ] + + , testGroup "valid" + [ + -- Arbitrary + testProperty "arbitrary" $ id @(MSeq A) + + -- Construct + , testProperty "fromList" $ fromList @A + , testProperty "fromRevList" $ fromRevList @A + , testProperty "replicate" $ replicate @A + , testProperty "generate" $ \n -> generate n . applyFun @Int @A + , testProperty "unfoldr" $ unfoldr (L.uncons @A) + , testProperty "unfoldl" $ unfoldl (unsnocL @A) + , testProperty "<>" $ (<>) @(MSeq A) + , testProperty "stimes" $ stimes @(MSeq A) @Int + , testProperty "mconcat" $ mconcat @(MSeq A) + , testProperty "concatMap []" $ concatMap @B @[] . applyFun @A + , testProperty "mfix" mkMfix + + -- Index + , testProperty "adjust" $ adjust . applyFun @A + , testProperty "update" $ update @A + , testProperty "insertAt" $ insertAt @A + , testProperty "deleteAt" $ deleteAt @A + + -- Slice + , testProperty "cons" $ cons @A + , testProperty "snoc" $ snoc @A + , testProperty "uncons" $ fmap snd . uncons @A + , testProperty "unsnoc" $ fmap fst . unsnoc @A + , testProperty "take" $ take @A + , testProperty "drop" $ drop @A + , testProperty "slice" $ slice @A + , testProperty "splitAt" $ splitAt @A + , testProperty "takeEnd" $ takeEnd @A + , testProperty "dropEnd" $ dropEnd @A + , testProperty "splitAtEnd" $ splitAtEnd @A + + -- Transform + , testProperty "fmap" $ map . applyFun @A @B + , testProperty "liftA2" $ \(Sqrt1 xs) (Sqrt1 ys) (f :: Fun (A,B) C) -> + liftA2 (applyFun2 f) xs ys + , testProperty "traverse" $ runIdentity .: traverse . applyFun @A @(Identity B) + , testProperty "imap" $ imap . applyFun2 @_ @A @B + , testProperty "itraverse" $ runIdentity .: itraverse . applyFun2 @_ @A @(Identity B) + , testProperty "reverse" $ reverse @A + , testProperty "intersperse" $ intersperse @A + , testProperty "scanl" $ scanl . applyFun2 @A @B + , testProperty "scanr" $ scanr . applyFun2 @A @B + , testProperty "sort" $ sort @A + , testProperty "sortBy" $ sortBy @A compare + + -- Filter + , testProperty "filter" $ filter . applyFun @A + , testProperty "mapMaybe" $ mapMaybe . applyFun @A @(Maybe B) + , testProperty "mapEither" $ mapEither . applyFun @A @(Either B C) + , testProperty "takeWhile" $ takeWhile . applyFun @A + , testProperty "dropWhile" $ dropWhile . applyFun @A + , testProperty "span" $ span . applyFun @A + , testProperty "break" $ break . applyFun @A + , testProperty "takeWhileEnd" $ takeWhileEnd . applyFun @A + , testProperty "dropWhileEnd" $ dropWhileEnd . applyFun @A + , testProperty "spanEnd" $ spanEnd . applyFun @A + , testProperty "breakEnd" $ breakEnd . applyFun @A + + -- Zip and unzip + , testProperty "zipWith" $ zipWith . applyFun2 @A @B @C + , testProperty "zipWith3" $ zipWith3 . applyFun3 @A @B @A @C + , testProperty "unzipWith" $ unzipWith . applyFun @A @(B,C) + , testProperty "unzipWith3" $ unzipWith3 . applyFun @A @(B,C,B) + + -- Random + , testProperty "random transforms" $ \tf (xs :: MSeq A) -> + let xs' = unTransform tf xs + in classify (not (null xs')) "non-empty" xs' + ] + + , testGroup "laws" $ L.map tastyLaws $ + let pa = Proxy @(MSeq A) + porda = Proxy @(MSeq A) + in + [ QLaws.eqLaws pa + , QLaws.ordLaws porda + , QLaws.isListLaws pa + , QLaws.semigroupLaws pa + , QLaws.monoidLaws pa + , QLaws.showLaws pa + , QLaws.showReadLaws pa + ] + ] + +instance (Arbitrary a, Measured a) => Arbitrary (MSeq a) where + arbitrary = oneof + [ fromList <$> arbitrary + , fromRevList <$> arbitrary + , randomStructure + ] + where + randomStructure = sized $ \n -> do + n' <- choose (0,n) + if n' == 0 + then pure empty + else Ap.liftA2 MSeqInternal.MTree arbitrary (go (n'-1)) + where + go 0 = pure MTreeInternal.MTip + go n = do + ln <- choose (0,n-1) + Ap.liftA3 MTreeInternal.link arbitrary (go ln) (go (n-1-ln)) + + shrink = fmap fromList . shrink . F.toList + +newtype Transform a = Transform { unTransform :: MSeq a -> MSeq a } + +instance Show (Transform a) where + show _ = "Transform" + +instance (Measured a, Arbitrary a, CoArbitrary a) => Arbitrary (Transform a) where + arbitrary = Transform . foldr (.) id <$> listOf tf + where + tf = oneof + [ (fst .) . splitAt <$> arbitrary + , (snd .) . splitAt <$> arbitrary + , (<>) <$> arbitrary + , flip (<>) <$> arbitrary + , insertAt <$> arbitrary <*> arbitrary + , deleteAt <$> arbitrary + , map <$> arbitrary + , fmap cons arbitrary + , fmap (flip snoc) arbitrary + , pure (maybe empty snd . uncons) + , pure (maybe empty fst . unsnoc) + , filter <$> arbitrary + , mapMaybe <$> arbitrary + ] + +instance (Show a, Measured a, Eq (Measure a), Show (Measure a)) => + Testable (MSeq a) where + property t = + counterexample ("Invalid: " ++ MSeqInternal.debugShowsPrec 0 t "") $ + MSeqInternal.valid t + +instance (Testable a, Testable b, a ~ MSeq x, b ~ MSeq y) => Testable (a, b) where + property (a, b) = property a .&&. property b + +instance + (Testable a, Testable b, Testable c, a ~ MSeq x, b ~ MSeq y, c ~ MSeq z) => + Testable (a, b, c) where + property (a, b, c) = property a .&&. property b .&&. property c + +instance (Testable a, a ~ MSeq x) => Testable [a] where + property = conjoin + +instance Measured a => ListLike (MSeq a) where + type E (MSeq a) = a + fromL = fromList + toL = F.toList + +-- A non-commutative semigroup +-- From https://math.stackexchange.com/a/2893712 +newtype Bato = Bato Integer deriving (Eq, Show) + +instance Semigroup Bato where + Bato x <> Bato y = Bato (abs x * y) + +newtype A = A Integer + deriving newtype (Eq, Ord, Read, Show, Arbitrary, CoArbitrary) + +instance Function A where + function = functionMap (\(A x) -> x) A + +instance Measured A where + type Measure A = Bato + measure (A x) = Bato x + +newtype B = B Integer + deriving newtype (Eq, Ord, Show, Arbitrary, CoArbitrary) + +instance Function B where + function = functionMap (\(B x) -> x) B + +instance Measured B where + type Measure B = Bato + measure (B x) = Bato x + +newtype C = C Integer + deriving newtype (Eq, Ord, Show, Arbitrary, CoArbitrary) + +instance Function C where + function = functionMap (\(C x) -> x) C + +instance Measured C where + type Measure C = Bato + measure (C x) = Bato x + +newtype D = D Integer + deriving newtype (Eq, Ord, Show, Arbitrary, CoArbitrary) + +instance Measured D where + type Measure D = Maybe Bato + measure (D x) = Just (Bato x) + +data LI = LI [Integer] deriving (Eq, Show) + +takeL :: Int -> LI -> LI +takeL n (LI xs) = LI (L.take n xs) + +instance Measured LI where + type Measure LI = Bato + measure (LI xs) = Bato (head xs) + +-- map (replicate 10) [0..n-1] +mkMfix :: Int -> MSeq LI +mkMfix n = + map (takeL 10) + (mfix (\ ~(LI is) -> generate n (\i -> LI (fromIntegral i : is)))) + +newtype S = S Word + deriving newtype (Eq, Ord, Show, Arbitrary) + +instance Measured S where + type Measure S = Sum Word + measure (S x) = Sum x diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,13 @@ +import Test.Tasty +import Test.Tasty.QuickCheck + +import Seq (seqTests) +import MSeq (mseqTests) +import PQueue (pQueueTests) + +main :: IO () +main = defaultMain $ localOption (QuickCheckTests 2000) $ testGroup "All" + [ seqTests + , mseqTests + , pQueueTests + ] diff --git a/test/PQueue.hs b/test/PQueue.hs new file mode 100644 --- /dev/null +++ b/test/PQueue.hs @@ -0,0 +1,98 @@ +{-# OPTIONS_GHC -Wno-orphans #-} -- Arbitrary instances +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +module PQueue (pQueueTests) where + +import Prelude hiding (concatMap, min) +import qualified Control.Applicative as Ap +import qualified Data.Foldable as F +import Data.Foldable (toList) +import qualified Data.Foldable.WithIndex as IFo +import qualified Data.List as L +import Data.Proxy (Proxy(..)) +import qualified Test.QuickCheck.Classes.Base as QLaws +import Test.QuickCheck.Poly (A, OrdA) +import Test.Tasty +import Test.Tasty.QuickCheck + +import Data.Seqn.PQueue +import TestUtil (ListLike(..), tastyLaws) +import qualified ListLikeTests as LL + +pQueueTests :: TestTree +pQueueTests = testGroup "Data.Seqn.PQueue" + [ testGroup "properties" + [ + LL.listLike @(PQueue OrdA) + + , LL.empty @(PQueue OrdA) empty + , LL.singleton @(PQueue OrdA) singleton + , LL.snoc @(PQueue OrdA) (flip insert) + , LL.concatMap @(PQueue OrdA) concatMap + , LL.read @(PQueue Int) + + , testProperty "min" $ \(q :: PQueue EA) -> + min q === minL (toList q) + , testProperty "minView" $ \(q :: PQueue EA) -> + (fmap . fmap) toList (minView q) === minViewL (toList q) + , testProperty "toSortedList" $ \(q :: PQueue EA) -> + toSortedList q === L.sort (toList q) + + , LL.foldMap @(PQueue OrdA) foldMap + , LL.foldMap' @(PQueue OrdA) F.foldMap' + , LL.foldr @(PQueue OrdA) foldr + , LL.foldl' @(PQueue OrdA) F.foldl' + , LL.foldl @(PQueue OrdA) foldl + , LL.foldr' @(PQueue OrdA) F.foldr' + , LL.ifoldMap @(PQueue OrdA) IFo.ifoldMap + , LL.ifoldMap' @(PQueue OrdA) IFo.ifoldMap' + , LL.ifoldr @(PQueue OrdA) IFo.ifoldr + , LL.ifoldr' @(PQueue OrdA) IFo.ifoldr' + , LL.ifoldl @(PQueue OrdA) IFo.ifoldl + , LL.ifoldl' @(PQueue OrdA) IFo.ifoldl' + , LL.show @(PQueue OrdA) + ] + + , testGroup "laws" $ map tastyLaws $ + let pe = Proxy @(PQueue EA) + pint = Proxy @(PQueue Int) + in + [ QLaws.eqLaws pe + , QLaws.ordLaws pe + , QLaws.isListLaws pe + , QLaws.semigroupLaws pe + , QLaws.monoidLaws pe + , QLaws.showLaws pe + , QLaws.showReadLaws pint + ] + ] + +-- Use Entry instead of just OrdA to test queue stability. +type EA = Entry OrdA A + +instance (Arbitrary k, Arbitrary a) => Arbitrary (Entry k a) where + arbitrary = Ap.liftA2 Entry arbitrary arbitrary + shrink (Entry k x) = uncurry Entry <$> shrink (k,x) + +instance (Ord a, Arbitrary a) => Arbitrary (PQueue a) where + arbitrary = fromList <$> arbitrary + shrink = fmap fromList . shrink . toList + +instance (CoArbitrary k, CoArbitrary a) => CoArbitrary (Entry k a) where + coarbitrary (Entry k x) = coarbitrary (k,x) + +instance (Function k, Function a) => Function (Entry k a) where + function = functionMap (\(Entry k x) -> (k,x)) (uncurry Entry) + +instance Ord a => ListLike (PQueue a) where + type E (PQueue a) = a + fromL = fromList + toL = F.toList + +minL :: Ord a => [a] -> Maybe a +minL [] = Nothing +minL xs = Just (minimum xs) + +minViewL :: Ord a => [a] -> Maybe (a, [a]) +minViewL xs = fmap (\x -> (x, xs L.\\ [x])) (minL xs) diff --git a/test/Seq.hs b/test/Seq.hs new file mode 100644 --- /dev/null +++ b/test/Seq.hs @@ -0,0 +1,342 @@ +{-# OPTIONS_GHC -Wno-orphans #-} -- Arbitrary instances +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +module Seq (seqTests) where + +import Prelude hiding (unzip3, zip3, unzip, zip, break, concatMap, drop, dropWhile, filter, lookup, replicate, reverse, splitAt, scanl, scanr, span, take, takeWhile, zipWith, zipWith3) +import qualified Control.Applicative as Ap +import Data.Coerce (coerce) +import Control.Monad.Fix (MonadFix(..)) +import qualified Data.Foldable as F +import Data.Foldable (toList) +import qualified Data.Foldable.WithIndex as IFo +import Data.Functor.Identity (Identity(..)) +import qualified Data.Functor.WithIndex as IFu +import qualified Data.List as L +import Data.Proxy (Proxy(..)) +import Data.Semigroup (stimes) +import qualified Data.Traversable.WithIndex as ITr +import Test.Tasty +import Test.Tasty.QuickCheck hiding (generate) +import qualified Test.QuickCheck.Classes.Base as QLaws +import Test.QuickCheck.Poly (A, B, C, OrdA) + +import Data.Seqn.Seq +import qualified Data.Seqn.Internal.Seq as SeqInternal +import qualified Data.Seqn.Internal.Tree as TreeInternal +import ListExtra (unsnocL) +import qualified ListLikeTests as LL +import TestUtil (Sqrt1(..), tastyLaws, (.:), (.:.), ListLike(..)) + +seqTests :: TestTree +seqTests = testGroup "Data.Seqn.Seq" + [ testGroup "properties" + [ + LL.listLike @(Seq A) + + -- Construct + , LL.empty @(Seq A) empty + , LL.singleton @(Seq A) singleton + , LL.replicate @(Seq A) replicate + , LL.replicateA @(Seq A) replicateA + , LL.generate @(Seq A) generate + , LL.generateA @(Seq A) generateA + , LL.unfoldr @(Seq A) unfoldr + , LL.unfoldl @(Seq A) unfoldl + , LL.unfoldrM @(Seq A) unfoldrM + , LL.unfoldlM @(Seq A) unfoldlM + , (LL.<>) @(Seq A) (<>) + , LL.stimes @(Seq A) stimes + , LL.mconcat @(Seq A) mconcat + , LL.concatMap @(Seq A) concatMap + , testProperty "mfix" $ \n -> + toList (mkMfix n) === fmap (Solo . L.replicate 10) [0 .. n-1] + , LL.read @(Seq Int) + + -- Convert + , LL.toRevList @(Seq A) toRevList + , LL.show @(Seq A) + + -- Fold + , LL.foldMap @(Seq A) foldMap + , LL.foldMap' @(Seq A) F.foldMap' + , LL.foldr @(Seq A) foldr + , LL.foldl @(Seq A) F.foldl + , LL.foldl' @(Seq A) F.foldl' + , LL.foldr' @(Seq A) F.foldr' + , LL.ifoldMap @(Seq A) IFo.ifoldMap + , LL.ifoldMap' @(Seq A) IFo.ifoldMap' + , LL.ifoldr @(Seq A) IFo.ifoldr + , LL.ifoldl @(Seq A) IFo.ifoldl + , LL.ifoldr' @(Seq A) IFo.ifoldr' + , LL.ifoldl' @(Seq A) IFo.ifoldl' + + -- Index + , LL.lookup @(Seq A) lookup + , LL.index @(Seq A) index + , LL.update @(Seq A) update + , LL.adjust @(Seq A) adjust + , LL.insertAt @(Seq A) insertAt + , LL.deleteAt @(Seq A) deleteAt + + -- Slice + , LL.cons @(Seq A) cons + , LL.snoc @(Seq A) snoc + , LL.uncons @(Seq A) uncons + , LL.unsnoc @(Seq A) unsnoc + , LL.take @(Seq A) take + , LL.drop @(Seq A) drop + , LL.slice @(Seq A) slice + , LL.splitAt @(Seq A) splitAt + , LL.takeEnd @(Seq A) takeEnd + , LL.dropEnd @(Seq A) dropEnd + , LL.splitAtEnd @(Seq A) splitAtEnd + , LL.tails @(Seq A) tails + , LL.inits @(Seq A) inits + , LL.chunksOf @(Seq A) chunksOf + + -- Filter + , LL.filter @(Seq A) filter + , LL.mapMaybe @(Seq A) @(Seq B) mapMaybe + , LL.mapEither @(Seq A) @(Seq B) @(Seq C) mapEither + , LL.filterM @(Seq A) filterA + , LL.mapMaybeM @(Seq A) @(Seq B) mapMaybeA + , LL.mapEitherM @(Seq A) @(Seq B) @(Seq C) mapEitherA + , LL.takeWhile @(Seq A) takeWhile + , LL.dropWhile @(Seq A) dropWhile + , LL.span @(Seq A) span + , LL.break @(Seq A) break + , LL.takeWhileEnd @(Seq A) takeWhileEnd + , LL.dropWhileEnd @(Seq A) dropWhileEnd + , LL.spanEnd @(Seq A) spanEnd + , LL.breakEnd @(Seq A) breakEnd + + -- Transform + , LL.map @(Seq A) @(Seq B) fmap + , LL.liftA2 @(Sqrt1 Seq A) @(Sqrt1 Seq B) @(Sqrt1 Seq C) (coerce (Ap.liftA2 @Seq @A @B @C)) + , (LL.<*) @(Seq A) @(Seq B) (<*) + , (LL.*>) @(Seq A) @(Seq B) (*>) + , LL.bind @(Seq A) @(Seq B) (>>=) + , LL.traverse @(Seq A) @(Seq B) traverse + , LL.imap @(Seq A) @(Seq B) IFu.imap + , LL.itraverse @(Seq A) @(Seq B) ITr.itraverse + , LL.reverse @(Seq A) reverse + , LL.intersperse @(Seq A) intersperse + , LL.scanl @(Seq A) @(Seq B) scanl + , LL.scanr @(Seq A) @(Seq B) scanr + , LL.sort @(Seq OrdA) sort + , LL.sortBy @(Seq A) sortBy + + -- Search and test + , LL.eq @(Seq A) (==) + , LL.cmp @(Seq OrdA) compare + , LL.findEnd @(Seq A) findEnd + , LL.findIndex @(Seq A) findIndex + , LL.findIndexEnd @(Seq A) findIndexEnd + , LL.infixIndices @(Seq A) infixIndices + , LL.binarySearchFind @(Seq OrdA) binarySearchFind + , LL.isPrefixOf @(Seq A) isPrefixOf + , LL.isSuffixOf @(Seq A) isSuffixOf + , LL.isInfixOf @(Seq A) isInfixOf + , LL.isSubsequenceOf @(Seq A) isSubsequenceOf + + -- Zip and unzip + , LL.zip @(Seq A) @(Seq B) zip + , LL.zip3 @(Seq A) @(Seq B) @(Seq C) zip3 + , LL.zipWith @(Seq A) @(Seq B) @(Seq C) zipWith + , LL.zipWith3 @(Seq A) @(Seq B) @(Seq C) @(Seq A) zipWith3 + , LL.zipWithM @(Seq A) @(Seq B) @(Seq C) zipWithM + , LL.zipWith3M @(Seq A) @(Seq B) @(Seq C) @(Seq A) zipWith3M + , LL.unzip @(Seq (A, B)) unzip + , LL.unzip3 @(Seq (A, B, C)) unzip3 + , LL.unzipWith @(Seq A) @(Seq B) @(Seq C) unzipWith + , LL.unzipWith3 @(Seq A) @(Seq B) @(Seq C) @(Seq A) unzipWith3 + ] + + , testGroup "valid" + [ + -- Arbitrary + testProperty "arbitrary" $ id @(Seq A) + + -- Construct + , testProperty "fromList" $ fromList @A + , testProperty "fromRevList" $ fromRevList @A + , testProperty "replicate" $ replicate @A + , testProperty "generate" $ \n -> generate n . applyFun @Int @A + , testProperty "unfoldr" $ unfoldr (L.uncons @A) + , testProperty "unfoldl" $ unfoldl (unsnocL @A) + , testProperty "<>" $ (<>) @(Seq A) + , testProperty "stimes" $ stimes @(Seq A) @Int + , testProperty "mconcat" $ mconcat @(Seq A) + , testProperty "concatMap" $ concatMap @[] . applyFun @A @(Seq B) + , testProperty "mfix" mkMfix + + -- Index + , testProperty "adjust" $ adjust . applyFun @A + , testProperty "update" $ update @A + , testProperty "insertAt" $ insertAt @A + , testProperty "deleteAt" $ deleteAt @A + + -- Slice + , testProperty "cons" $ cons @A + , testProperty "snoc" $ snoc @A + , testProperty "uncons" $ fmap snd . uncons @A + , testProperty "unsnoc" $ fmap fst . unsnoc @A + , testProperty "take" $ take @A + , testProperty "drop" $ drop @A + , testProperty "slice" $ slice @A + , testProperty "splitAt" $ splitAt @A + , testProperty "takeEnd" $ takeEnd @A + , testProperty "dropEnd" $ dropEnd @A + , testProperty "splitAtEnd" $ splitAtEnd @A + , testProperty "tails" $ tails @A + , testProperty "tails each" $ toList . tails @A + , testProperty "inits" $ inits @A + , testProperty "inits each" $ toList . inits @A + , testProperty "chunksOf" $ chunksOf @A + , testProperty "chunksOf each" $ toList .: chunksOf @A + + -- Transform + , testProperty "fmap" $ fmap @Seq . applyFun @A @B + , testProperty "liftA2" $ unSqrt1 .:. Ap.liftA2 @(Sqrt1 Seq) . applyFun2 @A @B @C + , testProperty "<*" $ (<*) @Seq @A @B + , testProperty "*>" $ (*>) @Seq @A @B + , testProperty ">>=" $ flip ((>>=) @Seq) . applyFun @A @(Seq B) + , testProperty "traverse" $ runIdentity .: traverse @Seq . applyFun @A @(Identity B) + , testProperty "imap" $ IFu.imap @_ @Seq . applyFun2 @_ @A @B + , testProperty "itraverse" $ runIdentity .: ITr.itraverse @_ @Seq . applyFun2 @_ @A @(Identity B) + , testProperty "reverse" $ reverse @A + , testProperty "intersperse" $ intersperse @A + , testProperty "scanl" $ scanl . applyFun2 @A @B + , testProperty "scanr" $ scanr . applyFun2 @A @B + , testProperty "sort" $ sort @OrdA + , testProperty "sortBy" $ sortBy @OrdA compare + + -- Filter + , testProperty "filter" $ filter . applyFun @A + , testProperty "mapMaybe" $ mapMaybe . applyFun @A @(Maybe B) + , testProperty "mapEither" $ mapEither . applyFun @A @(Either B C) + , testProperty "takeWhile" $ takeWhile . applyFun @A + , testProperty "dropWhile" $ dropWhile . applyFun @A + , testProperty "span" $ span . applyFun @A + , testProperty "break" $ break . applyFun @A + , testProperty "takeWhileEnd" $ takeWhileEnd . applyFun @A + , testProperty "dropWhileEnd" $ dropWhileEnd . applyFun @A + , testProperty "spanEnd" $ spanEnd . applyFun @A + , testProperty "breakEnd" $ breakEnd . applyFun @A + + -- Zip + , testProperty "zipWith" $ zipWith . applyFun2 @A @B @C + , testProperty "zipWith3" $ zipWith3 . applyFun3 @A @B @A @C + , testProperty "unzipWith" $ unzipWith . applyFun @A @(B,C) + , testProperty "unzipWith3" $ unzipWith3 . applyFun @A @(B,C,B) + + -- Random + , testProperty "random transforms" $ \tf (xs :: Seq A) -> + let xs' = unTransform tf xs + in classify (not (null xs')) "non-empty" xs' + ] + + , testGroup "laws" $ map tastyLaws $ + let p = Proxy @Seq + sqrtp = Proxy @(Sqrt1 Seq) + pa = Proxy @(Seq A) + porda = Proxy @(Seq OrdA) + pint = Proxy @(Seq Int) + in + [ QLaws.eqLaws pa + , QLaws.ordLaws porda + , QLaws.isListLaws pa + , QLaws.semigroupLaws pa + , QLaws.monoidLaws pa + , QLaws.showLaws pa + , QLaws.showReadLaws pint + , QLaws.foldableLaws p + , QLaws.traversableLaws p + , QLaws.functorLaws p + , QLaws.applicativeLaws sqrtp + , QLaws.alternativeLaws p + , QLaws.monadLaws sqrtp + , QLaws.monadPlusLaws p + , QLaws.monadZipLaws p + ] + ] + +instance Arbitrary a => Arbitrary (Seq a) where + arbitrary = oneof + [ fromList <$> arbitrary + , fromRevList <$> arbitrary + , randomStructure + ] + where + randomStructure = sized $ \n -> do + n' <- choose (0,n) + if n' == 0 + then pure empty + else Ap.liftA2 SeqInternal.Tree arbitrary (go (n'-1)) + where + go 0 = pure TreeInternal.Tip + go n = do + ln <- choose (0,n-1) + Ap.liftA3 TreeInternal.link arbitrary (go ln) (go (n-1-ln)) + + shrink = fmap fromList . shrink . F.toList + +newtype Transform a = Transform { unTransform :: Seq a -> Seq a } + +instance Show (Transform a) where + show _ = "Transform" + +instance (Arbitrary a, CoArbitrary a) => Arbitrary (Transform a) where + arbitrary = Transform . foldr (.) id <$> listOf tf + where + tf = oneof + [ (fst .) . splitAt <$> arbitrary + , (snd .) . splitAt <$> arbitrary + , (<>) <$> arbitrary + , flip (<>) <$> arbitrary + , insertAt <$> arbitrary <*> arbitrary + , deleteAt <$> arbitrary + , fmap <$> arbitrary + , fmap cons arbitrary + , fmap (flip snoc) arbitrary + , pure (maybe empty snd . uncons) + , pure (maybe empty fst . unsnoc) + , filter <$> arbitrary + , mapMaybe <$> arbitrary + ] + +instance Show a => Testable (Seq a) where + property t = + counterexample ("Invalid: " ++ SeqInternal.debugShowsPrec 0 t "") $ + SeqInternal.valid t + +instance (Testable a, Testable b, a ~ Seq x, b ~ Seq y) => Testable (a, b) where + property (a, b) = property a .&&. property b + +instance + (Testable a, Testable b, Testable c, a ~ Seq x, b ~ Seq y, c ~ Seq z) => + Testable (a, b, c) where + property (a, b, c) = property a .&&. property b .&&. property c + +instance (Testable a, a ~ Seq x) => Testable [a] where + property = conjoin + +instance ListLike (Seq a) where + type E (Seq a) = a + fromL = fromList + toL = toList + +-- map (replicate 10) [0..n-1] +mkMfix :: Int -> Seq (Solo [Int]) +mkMfix n = + fmap (fmap (L.take 10)) + (mfix (\ ~(Solo is) -> generate n (\i -> Solo (i : is)))) + +-- In Data.Tuple since 4.15 +data Solo a = Solo a + deriving (Eq, Show, Functor) diff --git a/test/TestUtil.hs b/test/TestUtil.hs new file mode 100644 --- /dev/null +++ b/test/TestUtil.hs @@ -0,0 +1,70 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +module TestUtil + ( ListLike(..) + , Sqrt1(..) + , tastyLaws + , MaybeBottom(..) + , isBottom + , eval + , (.:) + , (.:.) + ) where + +import Control.Exception (try, evaluate, ErrorCall) +import Test.Tasty +import Test.Tasty.QuickCheck +import Test.QuickCheck.Classes.Base (Laws(..)) + +class ListLike t where + type E t + fromL :: [E t] -> t + toL :: t -> [E t] + +newtype Sqrt1 f a = Sqrt1 { unSqrt1 :: f a } + deriving stock (Eq, Ord, Show) + deriving newtype (Functor, Applicative, Monad) + +instance Arbitrary (f a) => Arbitrary (Sqrt1 f a) where + arbitrary = sized $ \n -> Sqrt1 <$> resize (isqrt n) arbitrary + where + isqrt n = round (sqrt (fromIntegral n) :: Double) + shrink = fmap Sqrt1 . shrink . unSqrt1 + +instance (ListLike (f a), E (f a) ~ a) => ListLike (Sqrt1 f a) where + type E (Sqrt1 f a) = a + fromL = Sqrt1 . fromL + toL = toL . unSqrt1 + +data MaybeBottom a = Bottom | NotBottom a deriving (Eq, Show) + +instance Arbitrary a => Arbitrary (MaybeBottom a) where + arbitrary = frequency [(1, pure Bottom), (5, NotBottom <$> arbitrary)] + shrink = \case + Bottom -> [] + NotBottom x -> Bottom : fmap NotBottom (shrink x) + +isBottom :: MaybeBottom a -> Bool +isBottom = \case + Bottom -> True + NotBottom _ -> False + +eval :: a -> IO (MaybeBottom a) +eval = fmap (either @ErrorCall (const Bottom) NotBottom) . try . evaluate + +tastyLaws :: Laws -> TestTree +tastyLaws (Laws class_ tests) = + testGroup class_ (map (uncurry testProperty) tests) + +infixr 8 .: +(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d +(.:) = (.) . (.) + +infixr 8 .:. +(.:.) :: (d -> e) -> (a -> b -> c -> d) -> a -> b -> c -> e +(.:.) = (.) . (.:)