Add sorting functions

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`

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

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

src/Data/RRBVector/Internal.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -653,18 +652,26 @@
 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) []
 
@@ -716,7 +723,7 @@
     -- the type signature is necessary to compile
     mergeRebalance :: forall a. Shift -> A.Array (Tree a) -> A.Array (Tree a) -> A.Array (Tree a) -> A.Array (Tree a)
     mergeRebalance !sh !left !center !right
         | sh == blockShift = mergeRebalance' (\(Leaf arr) -> arr) Leaf
         | otherwise = mergeRebalance' treeToArray (computeSizes (down sh))
       where
         mergeRebalance' :: (Tree a -> A.Array t) -> (A.Array t -> Tree a) -> A.Array (Tree a)

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'
+    , foldrMap1, ifoldrMap1Step
     , empty, singleton, from2, wrap
     , replicate, replicateSnoc
     , index, head, last
@@ -107,6 +107,24 @@ 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
 
+-- helper function for implementing foldToMaybeTree
+foldrMap1 :: (a -> b) -> (b -> b -> b) -> Array a -> b
+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 (f x) (go (i + 1))
+    in go start
+
+-- helper function for implementing foldToMaybeWithIndexTree
+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"
 

src/Data/RRBVector/Internal/Sorting.hs

@@ -0,0 +1,126 @@
+module Data.RRBVector.Internal.Sorting
+    ( sort
+    , sortBy
+    , sortOn
+    , unstableSort
+    , unstableSortBy
+    , unstableSortOn
+    ) where
+
+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 (foldrMap1, 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 (<+>) arr
+    foldTree (Unbalanced arr _) = foldrMap1 foldTree (<+>) arr
+    foldTree (Leaf arr) = foldrMap1 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'

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