diff --git a/CHANGELOG.md b/CHANGELOG.md index 01d77d2..96fc490 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +# 0.2.2.0 - upcoming + +* Add `sort`, `sortBy`, `sortOn`, `unstableSort`, `unstableSortBy`, `unstableSortOn` ([#22](https://github.com/konsumlamm/rrb-vector/pull/22)) + # 0.2.1.0 - December 2023 * Add `findIndexL`, `findIndexR`, `findIndicesL`, `findIndicesR` diff --git a/rrb-vector.cabal b/rrb-vector.cabal index 961b35a..6aabfc0 100644 --- a/rrb-vector.cabal +++ b/rrb-vector.cabal @@ -47,9 +47,16 @@ library Data.RRBVector.Internal.Array Data.RRBVector.Internal.Buffer Data.RRBVector.Internal.IntRef - build-depends: base >= 4.11 && < 5, deepseq >= 1.4.3 && < 1.6, indexed-traversable ^>= 0.1, primitive >= 0.7 && < 0.10 + Data.RRBVector.Internal.Sorting + build-depends: + base >= 4.11 && < 5, + containers >= 0.5.11 && < 0.8, + deepseq >= 1.4.3 && < 1.6, + indexed-traversable ^>= 0.1, + primitive >= 0.7 && < 0.10 ghc-options: -O2 -Wall -Wno-name-shadowing -Werror=missing-methods -Werror=missing-fields default-language: Haskell2010 + default-extensions: BangPatterns test-suite test hs-source-dirs: test diff --git a/src/Data/RRBVector.hs b/src/Data/RRBVector.hs index 2a42dde..b04b8d4 100644 --- a/src/Data/RRBVector.hs +++ b/src/Data/RRBVector.hs @@ -44,6 +44,9 @@ module Data.RRBVector , map, map', reverse -- * Zipping and unzipping , zip, zipWith, unzip, unzipWith + -- * Sorting + , sort, sortBy, sortOn + , unstableSort, unstableSortBy, unstableSortOn ) where import Prelude hiding (replicate, lookup, take, drop, splitAt, map, reverse, zip, zipWith, unzip) @@ -53,3 +56,4 @@ import Data.Functor.WithIndex import Data.Traversable.WithIndex import Data.RRBVector.Internal +import Data.RRBVector.Internal.Sorting diff --git a/src/Data/RRBVector/Internal.hs b/src/Data/RRBVector/Internal.hs index b67dd51..7d89b0b 100644 --- a/src/Data/RRBVector/Internal.hs +++ b/src/Data/RRBVector/Internal.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -653,18 +652,26 @@ deleteAt :: Int -> Vector a -> Vector a deleteAt i v = let (left, right) = splitAt (i + 1) v in take i left >< right -- | \(O(n)\). Find the first index from the left that satisfies the predicate. +-- +-- @since 0.2.1.0 findIndexL :: (a -> Bool) -> Vector a -> Maybe Int findIndexL f = ifoldr (\i x acc -> if f x then Just i else acc) Nothing -- | \(O(n)\). Find the first index from the right that satisfies the predicate. +-- +-- @since 0.2.1.0 findIndexR :: (a -> Bool) -> Vector a -> Maybe Int findIndexR f = ifoldl (\i acc x -> if f x then Just i else acc) Nothing -- | \(O(n)\). Find the indices that satisfy the predicate, starting from the left. +-- +-- @since 0.2.1.0 findIndicesL :: (a -> Bool) -> Vector a -> [Int] findIndicesL f = ifoldr (\i x acc -> if f x then i : acc else acc) [] -- | \(O(n)\). Find the indices that satisfy the predicate, starting from the right. +-- +-- @since 0.2.1.0 findIndicesR :: (a -> Bool) -> Vector a -> [Int] findIndicesR f = ifoldl (\i acc x -> if f x then i : acc else acc) [] diff --git a/src/Data/RRBVector/Internal/Array.hs b/src/Data/RRBVector/Internal/Array.hs index 5110014..69b0c2f 100644 --- a/src/Data/RRBVector/Internal/Array.hs +++ b/src/Data/RRBVector/Internal/Array.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE RankNTypes #-} @@ -14,6 +13,7 @@ module Data.RRBVector.Internal.Array ( Array, MutableArray , ifoldrStep, ifoldlStep, ifoldrStep', ifoldlStep' + , ifoldrMap1Step , empty, singleton, from2, wrap , replicate, replicateSnoc , index, head, last @@ -36,6 +36,7 @@ import Control.DeepSeq (NFData(..)) import Control.Monad (when) import Control.Monad.ST import Data.Foldable (Foldable(..)) +import Data.Foldable1 (Foldable1(foldrMap1)) import Data.Primitive.SmallArray import Prelude hiding (replicate, take, drop, splitAt, head, last, map, traverse, read, unzip, (++)) @@ -74,6 +75,14 @@ instance Foldable Array where length (Array _ len _) = len +instance Foldable1 Array where + foldrMap1 f g (Array start len arr) = + let end = start + len + go i + | i == end - 1, (# x #) <- indexSmallArray## arr i = f x + | (# x #) <- indexSmallArray## arr i = g x (go (i + 1)) + in go start + instance (NFData a) => NFData (Array a) where rnf = foldl' (\_ x -> rnf x) () @@ -107,6 +116,14 @@ ifoldlStep' i0 step f z (Array start len arr) = | (# x #) <- indexSmallArray## arr i = go (i + 1) (j + step x) (f j acc x) in go start i0 z +ifoldrMap1Step :: Int -> (a -> Int) -> (Int -> a -> b) -> (b -> b -> b) -> Array a -> b +ifoldrMap1Step i0 step f g (Array start len arr) = + let end = start + len + go !i !j -- i is the index in arr, j is the index for f + | i == end - 1, (# x #) <- indexSmallArray## arr i = f j x + | (# x #) <- indexSmallArray## arr i = g (f j x) (go (i + 1) (j + step x)) + in go start i0 + uninitialized :: a uninitialized = errorWithoutStackTrace "uninitialized" diff --git a/src/Data/RRBVector/Internal/Sorting.hs b/src/Data/RRBVector/Internal/Sorting.hs new file mode 100644 index 0000000..52224eb --- /dev/null +++ b/src/Data/RRBVector/Internal/Sorting.hs @@ -0,0 +1,127 @@ +module Data.RRBVector.Internal.Sorting + ( sort + , sortBy + , sortOn + , unstableSort + , unstableSortBy + , unstableSortOn + ) where + +import Data.Foldable1 +import Data.Sequence.Internal.Sorting hiding + ( buildQ, buildTQ, buildIQ, buildITQ + , foldToMaybeTree, foldToMaybeWithIndexTree + , sort, sortBy, sortOn + , unstableSort, unstableSortBy, unstableSortOn + ) + +import Data.RRBVector.Internal +import Data.RRBVector.Internal.Array (ifoldrMap1Step) + +-- stable sorting + +foldToMaybeWithIndexTree :: (b -> b -> b) -> (Int -> a -> b) -> Int -> Vector a -> Maybe b +foldToMaybeWithIndexTree _ _ !_ Empty = Nothing +foldToMaybeWithIndexTree (<+>) f i (Root _ sh tree) = Just (foldTree i sh tree) + where + foldTree !i !sh (Balanced arr) = ifoldrMap1Step i (treeSize (down sh)) (flip foldTree (down sh)) (<+>) arr + foldTree i sh (Unbalanced arr _) = ifoldrMap1Step i (treeSize (down sh)) (flip foldTree (down sh)) (<+>) arr + foldTree i _ (Leaf arr) = ifoldrMap1Step i (\_ -> 1) f (<+>) arr + +buildIQ :: (a -> a -> Ordering) -> (Int -> a -> IndexedQueue a) -> Vector a -> Maybe (IndexedQueue a) +buildIQ cmp f = foldToMaybeWithIndexTree (mergeIQ cmp) f 0 + +buildITQ :: (b -> b -> Ordering) -> (Int -> a -> IndexedTaggedQueue b a) -> Vector a -> Maybe (IndexedTaggedQueue b a) +buildITQ cmp f = foldToMaybeWithIndexTree (mergeITQ cmp) f 0 + +-- | \(O(n \log n)\). Sort the vector in ascending order. +-- The sort is stable, meaning the order of equal elements is preserved. +-- +-- If stability is not required, `unstableSort` can be slightly faster and uses less memory. +-- +-- @since 0.2.2.0 +sort :: (Ord a) => Vector a -> Vector a +sort = sortBy compare + +-- | \(O(n \log n)\). Sort the vector in ascending order according to the specified comparison function. +-- The sort is stable, meaning the order of equal elements is preserved. +-- +-- If stability is not required, `unstableSortBy` can be slightly faster and uses less memory. +-- +-- @since 0.2.2.0 +sortBy :: (a -> a -> Ordering) -> Vector a -> Vector a +sortBy cmp v = case buildIQ cmp (\i x -> IQ i x IQNil) v of + Nothing -> Empty + Just q -> fromList $ go (length v) q + where + go 0 _ = [] + go n q = let (q', x) = popMinIQ cmp q in x : go (n - 1) q' + +-- | \(O(n \log n)\). Sort the vector in ascending order by comparing the results of applying the key function to each element. +-- The sort is stable, meaning the order of equal elements is preserved. +-- @`sortOn` f@ is equivalent to @`sortBy` (`Data.Ord.comparing` f)@, but only evaluates @f@ once for each element. +-- +-- If stability is not required, `unstableSortOn` can be slightly faster and uses less memory. +-- +-- @since 0.2.2.0 +sortOn :: (Ord b) => (a -> b) -> Vector a -> Vector a +sortOn f v = case buildITQ compare (\i x -> ITQ i (f x) x ITQNil) v of + Nothing -> Empty + Just q -> fromList $ go (length v) q + where + go 0 _ = [] + go n q = let (q', x) = popMinITQ compare q in x : go (n - 1) q' + +-- unstable sorting + +foldToMaybeTree :: (b -> b -> b) -> (a -> b) -> Vector a -> Maybe b +foldToMaybeTree _ _ Empty = Nothing +foldToMaybeTree (<+>) f (Root _ _ tree) = Just (foldTree tree) + where + foldTree (Balanced arr) = foldrMap1 foldTree ((<+>) . foldTree) arr + foldTree (Unbalanced arr _) = foldrMap1 foldTree ((<+>) . foldTree) arr + foldTree (Leaf arr) = foldrMap1 f ((<+>) . f) arr + +buildQ :: (a -> a -> Ordering) -> (a -> Queue a) -> Vector a -> Maybe (Queue a) +buildQ cmp = foldToMaybeTree (mergeQ cmp) + +buildTQ :: (b -> b -> Ordering) -> (a -> TaggedQueue b a) -> Vector a -> Maybe (TaggedQueue b a) +buildTQ cmp = foldToMaybeTree (mergeTQ cmp) + +-- | \(O(n \log n)\). Sort the vector in ascending order. +-- The sort is unstable, meaning the order of equal elements might not be preserved. +-- +-- If stability is required, use `sort` instead. +-- +-- @since 0.2.2.0 +unstableSort :: (Ord a) => Vector a -> Vector a +unstableSort = unstableSortBy compare + +-- | \(O(n \log n)\). Sort the vector in ascending order according to the specified comparison function. +-- The sort is unstable, meaning the order of equal elements might not be preserved. +-- +-- If stability is required, use `sortBy` instead. +-- +-- @since 0.2.2.0 +unstableSortBy :: (a -> a -> Ordering) -> Vector a -> Vector a +unstableSortBy cmp v = case buildQ cmp (\x -> Q x Nil) v of + Nothing -> Empty + Just q -> fromList $ go (length v) q + where + go 0 _ = [] + go n q = let (q', x) = popMinQ cmp q in x : go (n - 1) q' + +-- | \(O(n \log n)\). Sort the vector in ascending order by comparing the results of applying the key function to each element. +-- The sort is stable, meaning the order of equal elements is preserved. +-- @`unstableSortOn` f@ is equivalent to @`unstableSortBy` (`Data.Ord.comparing` f)@, but only evaluates @f@ once for each element. +-- +-- If stability is required, use `sortOn` instead. +-- +-- @since 0.2.2.0 +unstableSortOn :: (Ord b) => (a -> b) -> Vector a -> Vector a +unstableSortOn f v = case buildTQ compare (\x -> TQ (f x) x TQNil) v of + Nothing -> Empty + Just q -> fromList $ go (length v) q + where + go 0 _ = [] + go n q = let (q', x) = popMinTQ compare q in x : go (n - 1) q' diff --git a/test/Properties.hs b/test/Properties.hs index d93fb64..2bcff87 100644 --- a/test/Properties.hs +++ b/test/Properties.hs @@ -10,7 +10,8 @@ module Properties import Control.Applicative (liftA2) #endif import Data.Foldable (Foldable(..)) -import Data.List (uncons) +import Data.List (uncons, sort, sortOn) +import Data.Ord (comparing, Down(..)) import Data.Proxy (Proxy(..)) import Prelude hiding ((==)) -- use @===@ instead @@ -194,6 +195,30 @@ properties = testGroup "properties" [ testProperty "unzips the vector" $ \v -> (\(xs, ys) -> (toList xs, toList ys)) (V.unzip v) === unzip (toList v) , testProperty "valid" $ \v -> let (v1, v2) = V.unzip v in checkValid v1 .&&. checkValid v2 ] + , localOption (QuickCheckMaxSize 1000) $ testGroup "sorting" + [ testGroup "sort" + [ testProperty "sorts the vector" $ \v -> toList (V.sort v) === sort (toList v) + ] + , testGroup "sortBy" + [ testProperty "satisfies `sortBy compare = sort`" $ \v -> V.sortBy compare v === V.sort v + , testProperty "is stable" $ \v -> let cmp _ _ = EQ in V.sortBy cmp v === v + ] + , testGroup "sortOn" + [ testProperty "sorts the vector" $ \v -> toList (V.sortOn Down v) === sortOn Down (toList v) + , testProperty "satisfies `sortOn f = sortBy (comparing f)`" $ \v -> V.sortOn Down v === V.sortBy (comparing Down) v + , testProperty "is stable" $ \v -> let f _ = () in V.sortOn f v === v + ] + , testGroup "unstableSort" + [ testProperty "sorts the vector" $ \v -> toList (V.unstableSort v) === sort (toList v) + ] + , testGroup "unstableSortBy" + [ testProperty "satisfies `unstableSortBy compare = unstableSort`" $ \v -> V.unstableSortBy compare v === V.unstableSort v + ] + , testGroup "unstableSortOn" + [ testProperty "sorts the vector" $ \v -> toList (V.unstableSortOn id v) === sortOn id (toList v) + , testProperty "satisfies `unstableSortOn f = unstableSortBy (comparing f)`" $ \v -> V.unstableSortOn Down v === V.unstableSortBy (comparing Down) v + ] + ] , instances , laws , issues