samsort (empty) → 0.1.0.0
raw patch · 7 files changed
+1047/−0 lines, 7 filesdep +QuickCheckdep +basedep +primitive
Dependencies added: QuickCheck, base, primitive, samsort, tasty, tasty-quickcheck
Files
- CHANGELOG.md +3/−0
- HOWTO.md +299/−0
- LICENSE +30/−0
- README.md +53/−0
- samsort.cabal +65/−0
- src/Data/SamSort.hs +542/−0
- test/Main.hs +55/−0
+ CHANGELOG.md view
@@ -0,0 +1,3 @@+### 0.1.0.0 -- 2024-05-05++* First version.
+ HOWTO.md view
@@ -0,0 +1,299 @@+# Using this library++This is a short guide to demonstrate how this library may be used. Familiarity+with some form of arrays and mutable arrays in Haskell is expected. If your+array type is not present here, it should still be possible to adapt some of the+code below to your use case.++## Sorting boxed elements++This library offers the function++```hs+sortArrayBy+ :: (a -> a -> Ordering) -- ^ comparison+ -> MutableArray# s a+ -> Int -- ^ offset+ -> Int -- ^ length+ -> ST s ()+```++So how does one use this?++* The first parameter is a comparison function that will be used to order the+ elements.+* The second parameter is a [`MutableArray#`](https://hackage.haskell.org/package/base-4.19.0.0/docs/GHC-Exts.html#t:MutableArray-35-)+ with elements of type `a`. This is a primitive array type provided by GHC.+ This array will be sorted in place.+* The third and fourth parameters are `Int`s which demarcate a slice of the+ array. Elements in this slice will be sorted, and other elements will not be+ touched.+* The return type is an `ST` action. If you are not familiar with `ST`, please+ see the documentation for [`Control.Monad.ST`](https://hackage.haskell.org/package/base-4.19.0.0/docs/Control-Monad-ST.html).++Clearly, to use `sortArrayBy`, an important step is to put the elements to be+sorted into a `MutableArray#`. The most convenient way to do this depends on how+the elements are stored prior to sorting.++### Example 1: [`MVector`](https://hackage.haskell.org/package/vector-0.13.1.0/docs/Data-Vector-Mutable.html#t:MVector)++Consider that we need to sort a mutable vector `MVector` from the `vector`+library. This is quite easy, and in fact we do not need to put elements anywhere+because the underlying representation of an `MVector` is a `MutableArray#`! We+only need to get it out of the `MVector`.++```hs+import Control.Monad.Primitive (PrimMonad(..), stToPrim) -- from the package "primitive"+import Data.Primitive.Array (MutableArray(..)) -- also from "primitive"+import Data.Vector.Mutable (MVector(..))++import qualified Data.SamSort as Sam++-- | Sort a mutable vector in place.+sortMV :: (PrimMonad m, Ord a) => MVector (PrimState m) a -> m ()+sortMV = sortMVBy compare++-- | Sort a mutable vector in place using a comparison function.+sortMVBy :: PrimMonad m => (a -> a -> Ordering) -> MVector (PrimState m) a -> m ()+sortMVBy cmp (MVector off len (MutableArray ma)) =+ stToPrim $ Sam.sortArrayBy cmp ma off len+```++### Example 2: [`Vector`](https://hackage.haskell.org/package/vector-0.13.1.0/docs/Data-Vector.html#t:Vector)++Now consider sorting an (immutable) `Vector`, again from the `vector` library.+Since we cannot mutate it, we will return a sorted copy. The most convenient way+here is to thaw to a `MVector` and sort it as we did above.++```hs+import Data.Vector (Vector)+import qualified Data.Vector as V++-- | Sort a vector.+sortV :: Ord a => Vector a -> Vector a+sortV = sortVBy compare++-- | Sort a vector using a comparison function.+sortVBy :: (a -> a -> Ordering) -> Vector a -> Vector a+sortVBy cmp v = V.create $ do+ mv <- V.thaw v+ sortMVBy cmp mv -- from Example 1 above+ pure mv+```++We can test it out in GHCI.++```hs+>>> sortV (V.fromList [5,2,6,3,4,1])+[1,2,3,4,5,6]+>>> import Data.Ord (comparing)+>>> sortVBy (comparing length) (V.fromList ["Lunar","11.3","Candle","Magic"])+["11.3","Lunar","Magic","Candle"]+```++### Example 3: List++Let us now try to sort a list, like [`Data.List.sort`](https://hackage.haskell.org/package/base-4.19.0.0/docs/Data-List.html#v:sort)+does. We will need to move the elements from the list into a `MutableArray#`.++I recommend using the [`primitive`](https://hackage.haskell.org/package/primitive-0.9.0.0/docs/Data-Primitive-Array.html)+library for this task. `primitive` provides boxed wrappers over GHC primitive+types and functions to work with them. While it is possible to do this without+any library, it is easiest to use what is already available. If you are unable+to use `primitive`, you can take a peek at the relevant definitions there and+use them directly.++```hs+import Control.Monad.Primitive (stToPrim)+import qualified Data.Foldable as F+import Data.Primitive.Array (MutableArray(..))+import qualified Data.Primitive.Array as A++import qualified Data.SamSort as Sam++-- | Sort a list.+sortL :: Ord a => [a] -> [a]+sortL = sortLBy compare++-- | Sort a list using a comparison function.+sortLBy :: (a -> a -> Ordering) -> [a] -> [a]+sortLBy cmp xs = F.toList $ A.runArray $ do+ let a = A.arrayFromList xs+ n = A.sizeofArray a+ ma@(MutableArray ma') <- A.thawArray a 0 n+ stToPrim $ Sam.sortArrayBy cmp ma' 0 n+ pure ma+```++In GHCI,+```hs+>>> sortL ["Fall","In","The","Dark"]+["Dark","Fall","In","The"]+>>> import Data.Ord (Down, comparing)+>>> sortLBy (comparing Down) [3.4,8.5,9.1,7.9,3.1,6.2]+[9.1,8.5,7.9,6.2,3.4,3.1]+```++> [!TIP]+>+> Avoid `Data.List`'s `sort` and `sortBy` when a large number of elements need+> to be fully sorted and performance is a concern. Sorting lists is quite+> inefficient. Put the elements in a mutable array and use this (or some other)+> sorting library instead.++## Sorting `Int`s++Converting to a `MutableArray#` and sorting, as explained in the above section,+should cover the majority of use cases. However, sometimes it is not the+best option. For instance, we may be storing `Int`s in an unboxed array for+efficiency. Having to pull them out and box them for sorting does not sound+good.++The second function provided by this library is++```hs+sortIntArrayBy+ :: (Int -> Int -> Ordering) -- ^ comparison+ -> MutableByteArray# s+ -> Int -- ^ offset in Int#s+ -> Int -- ^ length in Int#s+ -> ST s ()+```++As you might have guessed, this sorts an unboxed array of `Int`s. We can use+this whenever we need to sort `Int`s, or even other types that may be cheaply+converted to and from `Int`s (like `Word`).++### Example 1: [Unboxed `MVector`](https://hackage.haskell.org/package/vector-0.13.1.0/docs/Data-Vector-Unboxed-Mutable.html#t:MVector)++Let us sort a mutable unboxed `MVector Int`. Like with the boxed `MVector`,+we do not need to move the elements because the underlying representation is a+`MutableByteArray#`.++```hs+import Control.Monad.Primitive (PrimMonad(..), stToPrim)+import Data.Primitive.ByteArray (MutableByteArray(..))+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed.Mutable as VUM+import qualified Data.Vector.Unboxed.Base as VUB++import qualified Data.SamSort as Sam++sortVUMInt :: PrimMonad m => VUM.MVector (PrimState m) Int -> m ()+sortVUMInt = sortVUMIntBy compare++sortVUMIntBy+ :: PrimMonad m+ => (Int -> Int -> Ordering) -> VUM.MVector (PrimState m) Int -> m ()+sortVUMIntBy cmp mv = case mv of+ VUB.MV_Int (VP.MVector off len (MutableByteArray ma')) ->+ stToPrim $ Sam.sortIntArrayBy cmp ma' off len+```++> [!WARNING]+>+> Do not try changing the element type above to sort any other unboxed vector.+> `MutableByteArray#` is the underlying representation for many unboxed vectors,+> but it would be incorrect to use the above code if the element type is not+> `Int`.++## Sorting by index++We have now covered sorting boxed values, and sorting `Int`s. What about other+types in unboxed arrays?++### Example 1: [Unboxed `Vector`](https://hackage.haskell.org/package/vector-0.13.1.0/docs/Data-Vector-Unboxed.html#t:Vector)++Consider that we need to sort an unboxed vector of some type `a`. The `vector`+library is designed in a way that the underlying representation of an unboxed+vector can be anything depending on the type `a`. We cannot assume anything+about it.++We know that we can index such a vector efficiently. We also know that we can+construct such vectors from an `Int -> a` using the handy `generate` function. We+will use these facts to sort such a vector.++First we will create an `Int` vector, the elements of which will be indices into+the `a` vector. Then we will sort this `Int` vector using a comparison function+that indexes the `a` vector and compares `a`s. Finally, we will construct a+vector with `a`s in the order of the sorted indices.++This technique is general enough that we can sort any flavor of `Vector`+(boxed, `Unboxed`, `Prim`, `Storable`), so let us use `Vector.Generic` to+define the functions.++```hs+import Control.Monad.Primitive (stToPrim)+import Data.Primitive.ByteArray (MutableByteArray(..))+import qualified Data.Vector.Generic as VG+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Primitive.Mutable as VPM++import qualified Data.SamSort as Sam++-- | Sort a vector.+sortByIdxVG :: (Ord a, VG.Vector v a) => v a -> v a+sortByIdxVG = sortByIdxVGBy compare++-- | Sort a vector using a comparison function.+sortByIdxVGBy :: VG.Vector v a => (a -> a -> Ordering) -> v a -> v a+sortByIdxVGBy cmp v = VG.generate n (VG.unsafeIndex v . VP.unsafeIndex ixa)+ where+ n = VG.length v+ cmp' i j = cmp (VG.unsafeIndex v i) (VG.unsafeIndex v j)+ ixa = VP.create $ do+ ixma <- VPM.generate n id+ case ixma of+ VPM.MVector off len (MutableByteArray ma') ->+ stToPrim $ Sam.sortIntArrayBy cmp' ma' off len+ pure ixma+```++In fact, this technique is general enough to be used whenever indices can+be sorted, using any method, and not just with this library!++Sorting by index is more beneficial the larger the elements are in memory,+since moving around index `Int`s is cheaper than moving around the elements+themselves.++We can see that the sort works as expected in GHCI.++```hs+>>> import Data.Ord (comparing)+>>> import qualified Data.Vector.Unboxed as VU+>>> let v = VU.fromList [(6,4),(5,4),(1,2)] :: VU.Vector (Int,Int)+>>> sortByIdxVGBy (comparing snd) v+[(1,2),(6,4),(5,4)]+```++And we can see [in benchmarks](https://github.com/meooow25/samsort/tree/master/compare#4-sort-105-int-int-ints-unboxed)+that sorting by index is indeed more efficient than sorting directly, for+elements of type `(Int, Int, Int)` and sort implementations which support both.++## Sorting unboxed arrays of small elements++So sorting by index is more beneficial the larger the element is, but what+about small elements? Perhaps we need to sort an unboxed array of `Word8`s, or+`Float`s?++Our options as seen above are++* Convert to a boxed array and sort. Lots of avoidable allocations and slow+ comparisons.+* Sort by index. Better, but has avoidable allocations in the form of the+ index array.++Neither are ideal. The most efficient way to sort small elements is to sort the+array of such elements directly. Unfortunately, this library cannot be used to+do this because there are only two functions, one to sort boxed values, and one+to sort `Int`s. If we must use this library, sorting by index is the method of+choice. It is not ideal, but it will not be slow either.++The [`primitive-sort`](https://hackage.haskell.org/package/primitive-sort)+library may also be a good choice for this task. It can sort such small elements+efficiently, though it has some drawbacks (not adaptive, cannot sort a slice,+cannot sort using a comparison function, more dependencies).++[`vector-algorithms`](https://hackage.haskell.org/package/vector-algorithms)+is also able to sort small elements, however it turns out to be+[slower in practice](https://github.com/meooow25/samsort/tree/master/compare#5-sort-105-word8s-unboxed).
+ LICENSE view
@@ -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.
+ README.md view
@@ -0,0 +1,53 @@+# samsort++[](https://hackage.haskell.org/package/samsort)+[](https://github.com/meooow25/samsort/actions/workflows/haskell-ci.yml)++A stable adapative mergesort implementation++## Features++This is a lightweight library offering two high performance sort functions:++* `sortArrayBy` sorts a GHC [`MutableArray#`](https://hackage.haskell.org/package/base-4.19.0.0/docs/GHC-Exts.html#t:MutableArray-35-)+ of boxed elements in place.+* `sortIntArrayBy` sorts a GHC [`MutableByteArray#`](https://hackage.haskell.org/package/base-4.19.0.0/docs/GHC-Exts.html#t:MutableByteArray-35-)+ of `Int#`s in place.++There are no dependencies outside of `base`. This means that this library is+not tied to array abstractions from any particular library. This also means+that you may need to write a few lines of code to get a `MutableArray#` or+`MutableByteArray#` from your data, which can then be sorted. See+[`HOWTO.md`](https://github.com/meooow25/samsort/blob/master/HOWTO.md)+for a guide.++If you need to use this library in an environment where you cannot depend on+other packages, you may simply copy the lone source file+[`src/Data/SamSort.hs`](https://github.com/meooow25/samsort/blob/master/src/Data/SamSort.hs)+to your project.++## Algorithm details++* The sort is a comparison-based $O(n \log n)$ mergesort.+* The sort is stable, i.e. the order of equal elements in the input is+ preserved.+* The sort is adaptive, i.e. the sort identifies and uses ascending and+ descending runs of elements occuring in the input to perform less work. As a+ result, the sort is $O(n)$ for already sorted inputs.+* The sort is the fastest among implementations from other libraries in most+ scenarios. [See the benchmarks](https://github.com/meooow25/samsort/tree/master/compare)+ for details.++## Known issues++Ideally, this library would offer only an algorithm, capable of sorting arrays+of any flavor. To support different arrays we would need to rely on some+abstraction, either from another library (like `vector`), or created here. We+cannot do either of those while also keeping the library as lightweight as it+is now.++## Contributing++Questions, bug reports, documentation improvements, code contributions welcome!+Please [open an issue](https://github.com/meooow25/samsort/issues) as the first+step. Slow performance counts as a bug!
+ samsort.cabal view
@@ -0,0 +1,65 @@+cabal-version: 3.0+name: samsort+version: 0.1.0.0+synopsis: A stable adaptive mergesort implementation+description: A stable adaptive mergesort implementation.+homepage: https://github.com/meooow25/samsort+bug-reports: https://github.com/meooow25/samsort/issues+license: BSD-3-Clause+license-file: LICENSE+author: Soumik Sarkar+maintainer: soumiksarkar.3120@gmail.com+copyright: (c) 2024 Soumik Sarkar+category: Data+build-type: Simple++extra-doc-files:+ CHANGELOG.md+ HOWTO.md+ README.md++tested-with:+ GHC == 8.4.4+ , GHC == 8.6.5+ , 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/samsort.git++common warnings+ ghc-options: -Wall++library+ import: warnings++ exposed-modules:+ Data.SamSort++ build-depends:+ base >= 4.11 && < 5++ hs-source-dirs: src+ default-language: Haskell2010++test-suite samsort-test+ import: warnings++ main-is: Main.hs+ build-depends:+ base+ , samsort+ , primitive+ , QuickCheck+ , tasty+ , tasty-quickcheck++ default-language: Haskell2010+ type: exitcode-stdio-1.0+ hs-source-dirs: test
+ src/Data/SamSort.hs view
@@ -0,0 +1,542 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}++-- |+-- Copyright: (c) 2024 Soumik Sarkar+-- License: BSD-3-Clause+--+-- A stable adaptive mergesort implementation.+--+-- The merging strategy used is "2-merge" as described by+--+-- * Sam Buss, Alexander Knop,+-- /\"Strategies for Stable Merge Sorting\"/,+-- 2018,+-- https://arxiv.org/abs/1801.04641+--+module Data.SamSort+ ( sortArrayBy+ , sortIntArrayBy+ ) where++import Control.Monad (when)+import Data.Bits (finiteBitSize, countLeadingZeros, shiftR)++import GHC.ST (ST(..))+import GHC.Exts+ ( Int#+ , Int(..)+ , MutableArray#+ , MutableByteArray#+ , (*#)+ , copyMutableArray#+ , copyMutableByteArray#+ , newArray#+ , newByteArray#+ , readArray#+ , readIntArray#+ , writeArray#+ , writeIntArray#+ )++-- | \(O(n \log n)\). Sort a slice of a @MutableArray#@ using a comparison+-- function.+--+-- The comparison must form a total order, as required by the 'Ord' laws.+--+-- @offset@ and @length@ must be valid, i.e.+--+-- * @0 <= offset < array size@ .+-- * @0 <= length@ .+-- * @offset + length <= array size@ .+--+-- This function will inline to get the best performance out of statically+-- known comparison functions. To avoid code duplication, create a wrapping+-- definition and reuse it as necessary.+--+sortArrayBy+ :: (a -> a -> Ordering) -- ^ comparison+ -> MutableArray# s a+ -> Int -- ^ offset+ -> Int -- ^ length+ -> ST s ()+sortArrayBy cmp = -- Inline with 1 arg+ \ma# !off !len -> sortArrayBy' cmp (MA ma#) off len+{-# INLINE sortArrayBy #-}++sortArrayBy'+ :: (a -> a -> Ordering)+ -> MA s a+ -> Int+ -> Int+ -> ST s ()+sortArrayBy' _ !_ !_ len | len < 2 = pure ()+sortArrayBy' cmp ma off len = do+ -- See Note [Algorithm overview]++ !swp <- newA (len `shiftR` 1) errorElement+ !stk <- newI (lg len)++ let -- Merge [i1,i2) and [i2,i3)+ -- Precondition: i1 < i2 < i3+ merge !i1 !i2 !i3+ | i2-i1 <= i3-i2 = mergeCopyLeft1 i1 i2 i3+ | otherwise = mergeCopyRight1 i1 i2 i3++ mergeCopyLeft1 !i1 !i2 !i3 = do+ x0 <- readA ma i1 -- See Note [First iteration]+ y <- readA ma i2+ if y `lt` x0+ then mergeCopyLeft2 i1 i2 i3+ else do+ let skip i | i >= i2 = pure ()+ skip i = do+ x <- readA ma i+ if y `lt` x+ then mergeCopyLeft2 i i2 i3+ else skip (i+1)+ skip (i1+1)++ -- Precondition: i1 < i2 < i3, (ma!i2) `lt` (ma!i1)+ mergeCopyLeft2 !i1 !i2 !i3 = do+ copyA ma i1 swp 0 (i2-i1)+ readA ma i2 >>= writeA ma i1+ if i2+1 < i3+ then loop 0 (i2+1) (i1+1)+ else copyA swp 0 ma (i1+1) len1+ where+ !len1 = i2-i1+ loop !h !j !k = do+ x <- readA swp h+ y0 <- readA ma j -- See Note [First iteration]+ let nxt !j1 !k1 = do+ writeA ma k1 x+ when (h+1 < len1) $+ loop (h+1) j1 (k1+1)+ if y0 `lt` x+ then do+ let loop2 j1 !k1 | j1 >= i3 = copyA swp h ma k1 (len1-h)+ loop2 j1 k1 = do+ y <- readA ma j1+ if y `lt` x+ then do+ writeA ma k1 y+ loop2 (j1+1) (k1+1)+ else+ nxt j1 k1+ writeA ma k y0+ loop2 (j+1) (k+1)+ else+ nxt j k++ mergeCopyRight1 !i1 !i2 !i3 = do+ x <- readA ma (i2-1)+ y0 <- readA ma (i3-1) -- See Note [First iteration]+ if y0 `lt` x+ then mergeCopyRight2 i1 i2 i3+ else do+ let skip j | j < i2 = pure ()+ skip j = do+ y <- readA ma j+ if y `lt` x+ then mergeCopyRight2 i1 i2 (j+1)+ else skip (j-1)+ skip (i3-2)++ -- Precondition: i1 < i2 < i3, (ma!(i3-1)) `lt` (ma!(i2-1))+ mergeCopyRight2 !i1 !i2 !i3 = do+ copyA ma i2 swp 0 (i3-i2)+ readA ma (i2-1) >>= writeA ma (i3-1)+ if i2-2 >= i1+ then loop (i2-2) (i3-i2-1) (i3-2)+ else copyA swp 0 ma i1 (i3-i2)+ where+ loop !h !j !k = do+ x0 <- readA ma h -- See Note [First iteration]+ y <- readA swp j+ let nxt !h1 !k1 = do+ writeA ma k1 y+ when (j > 0) $+ loop h1 (j-1) (k1-1)+ if y `lt` x0+ then do+ let loop2 h1 !_ | h1 < i1 = copyA swp 0 ma i1 (j+1)+ loop2 h1 k1 = do+ x <- readA ma h1+ if y `lt` x+ then do+ writeA ma k1 x+ loop2 (h1-1) (k1-1)+ else+ nxt h1 k1+ writeA ma k x0+ loop2 (h-1) (k-1)+ else+ nxt h k++ mergeStrategy merge getRun stk off end++ where+ lt x y = case cmp x y of LT -> True; _ -> False+ {-# INLINE lt #-}+ -- Note: Use lt instead of gt. Why? Because `compare` for types like Int and+ -- Word are defined in a way that needs one `<` op for LT but two (`<`,`==`)+ -- for GT.++ !end = off + len++ getRun = mkGetRun lt (readA ma) (writeA ma) (reverseA ma) end+{-# INLINE sortArrayBy' #-}++-- | \(O(n \log n)\). Sort a slice of a @MutableByteArray#@ interpreted as an+-- array of @Int#@s using a comparison function.+--+-- The comparison must form a total order, as required by the 'Ord' laws.+--+-- @offset@ and @length@ must be valid, i.e.+--+-- * @0 <= offset < array size@ .+-- * @0 <= length@ .+-- * @offset + length <= array size@ .+--+-- This function will inline to get the best performance out of statically+-- known comparison functions. To avoid code duplication, create a wrapping+-- definition and reuse it as necessary.+--+sortIntArrayBy+ :: (Int -> Int -> Ordering) -- ^ comparison+ -> MutableByteArray# s+ -> Int -- ^ offset in @Int#@s+ -> Int -- ^ length in @Int#@s+ -> ST s ()+sortIntArrayBy cmp = -- Inline with 1 arg+ \ma# !off !len -> sortIntArrayBy' cmp (MIA ma#) off len+{-# INLINE sortIntArrayBy #-}++sortIntArrayBy'+ :: (Int -> Int -> Ordering)+ -> MIA s+ -> Int+ -> Int+ -> ST s ()+sortIntArrayBy' _ !_ !_ len | len < 2 = pure ()+sortIntArrayBy' cmp ma off len = do+ -- See Note [Algorithm overview]++ !swp <- newI (len `shiftR` 1)+ !stk <- newI (lg len)++ let -- Merge [i1,i2) and [i2,i3)+ -- Precondition: i1 < i2 < i3+ merge !i1 !i2 !i3+ | i2-i1 <= i3-i2 = mergeCopyLeft1 i1 i2 i3+ | otherwise = mergeCopyRight1 i1 i2 i3++ mergeCopyLeft1 !i1 !i2 !i3 = readI ma i2 >>= skip i1+ where+ skip !i !y = do+ x <- readI ma i+ if y `lt` x+ then mergeCopyLeft2 i i2 i3+ else+ when (i < i2) $+ skip (i+1) y++ -- Precondition: i1 < i2 < i3, (ma!i2) `lt` (ma!i1)+ mergeCopyLeft2 !i1 !i2 !i3 = do+ copyI ma i1 swp 0 (i2-i1)+ readI ma i2 >>= writeI ma i1+ if i2+1 < i3+ then loop 0 (i2+1) (i1+1)+ else copyI swp 0 ma (i1+1) len1+ where+ !len1 = i2-i1+ loop !h !j !k = readI swp h >>= loop2 j k h+ loop2 j1 !k1 !h !_ | j1 >= i3 = copyI swp h ma k1 (len1-h)+ loop2 j1 k1 h x = do+ y <- readI ma j1+ if y `lt` x+ then do+ writeI ma k1 y+ loop2 (j1+1) (k1+1) h x+ else do+ writeI ma k1 x+ when (h+1 < len1) $+ loop (h+1) j1 (k1+1)++ mergeCopyRight1 !i1 !i2 !i3 = readI ma (i2-1) >>= skip (i3-1)+ where+ skip !j !x = do+ y <- readI ma j+ if y `lt` x+ then mergeCopyRight2 i1 i2 (j+1)+ else+ when (j >= i2) $+ skip (j-1) x++ -- Precondition: i1 < i2 < i3, (ma!(i3-1)) `lt` (ma!(i2-1))+ mergeCopyRight2 !i1 !i2 !i3 = do+ copyI ma i2 swp 0 (i3-i2)+ readI ma (i2-1) >>= writeI ma (i3-1)+ if i2-2 >= i1+ then loop (i2-2) (i3-i2-1) (i3-2)+ else copyI swp 0 ma i1 (i3-i2)+ where+ loop !h !j !k = readI swp j >>= loop2 h k j+ loop2 h1 !_ !j !_ | h1 < i1 = copyI swp 0 ma i1 (j+1)+ loop2 h1 k1 j y = do+ x <- readI ma h1+ if y `lt` x+ then do+ writeI ma k1 x+ loop2 (h1-1) (k1-1) j y+ else do+ writeI ma k1 y+ when (j > 0) $+ loop h1 (j-1) (k1-1)++ mergeStrategy merge getRun stk off end++ where+ lt !x !y = case cmp x y of LT -> True; _ -> False+ {-# INLINE lt #-}+ -- Note: Use lt instead of gt. Why? Because `compare` for types like Int and+ -- Word are defined in a way that needs one `<` op for LT but two (`<`,`==`)+ -- for GT.++ !end = off + len++ getRun = mkGetRun lt (readI ma) (writeI ma) (reverseI ma) end+{-# INLINE sortIntArrayBy' #-}++mkGetRun+ :: (a -> a -> Bool) -- comparison+ -> (Int -> ST s a) -- read+ -> (Int -> a -> ST s ()) -- write+ -> (Int -> Int -> ST s ()) -- reverse+ -> Int -- end+ -> (Int -> ST s Int)+mkGetRun lt rd wt rev !end = getRun+ where+ runAsc i | i >= end = pure i+ runAsc i = do+ x <- rd (i-1)+ y <- rd i+ if y `lt` x+ then pure i+ else runAsc (i+1)++ runDesc i | i >= end = pure i+ runDesc i = do+ x <- rd (i-1)+ y <- rd i+ if y `lt` x+ then runDesc (i+1)+ else pure i++ -- Insertion sort [i2,i3) into [i1,i3)+ -- Precondition: i1 < i2, i1 < i3+ insLoop !_ i2 i3 | i2 >= i3 = pure i2+ insLoop i1 i2 i3 = do+ x0 <- rd (i2-1)+ y <- rd i2+ when (y `lt` x0) $ do+ let ins j | j <= i1 = wt j y+ ins j = do+ x <- rd (j-1)+ if y `lt` x+ then wt j x *> ins (j-1)+ else wt j y+ wt i2 x0 *> ins (i2-1)+ insLoop i1 (i2+1) i3++ getRun i | i >= end || i+1 >= end = pure end+ getRun i = do+ x <- rd i+ y <- rd (i+1)+ !j <- if y `lt` x+ then do+ j <- runDesc (i+2)+ j <$ rev i (j-1)+ else runAsc (i+2)+ let k = i + minRunLen+ k' = if k <= 0 -- overflowed+ then end+ else min end k+ insLoop i j k'+{-# INLINE mkGetRun #-}++minRunLen :: Int+minRunLen = 8++mergeStrategy+ :: (Int -> Int -> Int -> ST s ()) -- merge+ -> (Int -> ST s Int) -- get next run+ -> MIA s -- stack+ -> Int -- offset+ -> Int -- end+ -> ST s ()+mergeStrategy merge getRun !stk !off !end = getRun off >>= mergeRuns (-1) off+ where+ -- [i,j) is the last run. Runs before it are on the stack.+ mergeRuns !top !i j+ | j >= end = finish top i+ | otherwise = getRun j >>= popPush top i j++ -- Maintain stack invariants+ popPush !top !i2 !i3 !i4+ | not (badYZ i2 i3 i4) = do+ writeI stk (top+1) i2+ mergeRuns (top+1) i3 i4+ | top < 0 = do+ merge i2 i3 i4+ mergeRuns top i2 i4+ | otherwise = do+ i1 <- readI stk top+ if mergeWithLeft i1 i2 i3 i4+ then do+ merge i1 i2 i3+ popPush (top-1) i1 i3 i4+ else do+ merge i2 i3 i4+ popPush (top-1) i1 i2 i4++ finish top !_ | top < 0 = pure ()+ finish top j = do+ i <- readI stk top+ merge i j end+ finish (top-1) i+{-# INLINE mergeStrategy #-}++badYZ :: Int -> Int -> Int -> Bool+badYZ i1 i2 i3 = (i2-i1) `shiftR` 1 < (i3-i2)+{-# INLINE badYZ #-}++mergeWithLeft :: Int -> Int -> Int -> Int -> Bool+mergeWithLeft i1 i2 i3 i4 = (i2-i1) < (i4-i3)+{-# INLINE mergeWithLeft #-}++reverseA+ :: MA s a+ -> Int -- ^ Start+ -> Int -- ^ End (inclusive)+ -> ST s ()+reverseA !ma = loop+ where+ loop i j | i >= j = pure ()+ loop i j = do+ x <- readA ma i+ readA ma j >>= writeA ma i+ writeA ma j x+ loop (i+1) (j-1)++reverseI+ :: MIA s+ -> Int -- ^ Start+ -> Int -- ^ End (inclusive)+ -> ST s ()+reverseI !ma = loop+ where+ loop i j | i >= j = pure ()+ loop i j = do+ x <- readI ma i+ readI ma j >>= writeI ma i+ writeI ma j x+ loop (i+1) (j-1)++lg :: Int -> Int+lg 0 = 0+lg i = finiteBitSize i - 1 - countLeadingZeros i+{-# INLINE lg #-}++errorElement :: a+errorElement = error "errorElement"++--------------------++-- The boxed wrappers MA, MIA, and functions operating on them are for the+-- convenience of working in ST. All of it should get optimized away.++data MA s a = MA (MutableArray# s a)++newA :: Int -> a -> ST s (MA s a)+newA (I# n#) x = ST $ \s ->+ case newArray# n# x s of (# s1, ma# #) -> (# s1, MA ma# #)+{-# INLINE newA #-}++readA :: MA s a -> Int -> ST s a+readA (MA ma#) (I# i#) = ST $ readArray# ma# i#+{-# INLINE readA #-}++writeA :: MA s a -> Int -> a -> ST s ()+writeA (MA ma#) (I# i#) x = ST $ \s ->+ case writeArray# ma# i# x s of s1 -> (# s1, () #)+{-# INLINE writeA #-}++copyA :: MA s a -> Int -> MA s a -> Int -> Int -> ST s ()+copyA (MA src#) (I# srcOff#) (MA dst#) (I# dstOff#) (I# len#) = ST $ \s ->+ case copyMutableArray# src# srcOff# dst# dstOff# len# s of s1 -> (# s1, () #)+{-# INLINE copyA #-}++data MIA s = MIA (MutableByteArray# s)++newI :: Int -> ST s (MIA s)+newI (I# n#) = ST $ \s ->+ case newByteArray# (n# *# intSize# (# #)) s of (# s1, ma# #) -> (# s1, MIA ma# #)+{-# INLINE newI #-}++readI :: MIA s -> Int -> ST s Int+readI (MIA ma#) (I# i#) = ST $ \s ->+ case readIntArray# ma# i# s of (# s1, x# #) -> (# s1, I# x# #)+{-# INLINE readI #-}++writeI :: MIA s -> Int -> Int -> ST s ()+writeI (MIA ma#) (I# i#) (I# x#) = ST $ \s ->+ case writeIntArray# ma# i# x# s of s1 -> (# s1, () #)+{-# INLINE writeI #-}++copyI :: MIA s -> Int -> MIA s -> Int -> Int -> ST s ()+copyI (MIA src#) (I# srcOff#) (MIA dst#) (I# dstOff#) (I# len#) = ST $ \s ->+ case copyMutableByteArray#+ src#+ (srcOff# *# intSize# (# #))+ dst#+ (dstOff# *# intSize# (# #))+ (len# *# intSize# (# #))+ s of+ s1 -> (# s1, () #)+{-# INLINE copyI #-}++intSize# :: (# #) -> Int#+intSize# _ = case finiteBitSize (0 :: Int) `shiftR` 3 of I# wsz# -> wsz#++--------------------++-- Note [Algorithm overview]+-- ~~~~~~~~~~~~~~~~~~~~~~~~~+-- Find non-decreasing and decreasing runs. Decreasing runs are reversed in+-- place. If a run is shorter than minRunLen, extend it to minRunLen using+-- insertion sort. Maintain a stack of runs. As each run is found, add it to+-- the stack and maintain stack invariants according to the 2-merge strategy.+-- This involves merging adjacent runs. Merging two runs is done by copying the+-- smaller run to a swap array, then merging into the main array. Elements of+-- the smaller array that can stay in place are skipped and not copied. After+-- all runs are found, runs on the stack are merged to get the final sorted+-- array.++-- Note [First iteration]+-- ~~~~~~~~~~~~~~~~~~~~~~+-- In certain places, the first iteration of a loop is pulled out of the loop+-- when many elements need to be compared with one element. This is to make GHC+-- aware that if the comparison is strict, the one element can be evaluated and+-- perhaps unboxed for subsequent comparisons. This could also be achieved by+-- being strict in the element, but we want to allow the comparison function+-- to be potentially lazy.++-- Note [Integer overflows]+-- ~~~~~~~~~~~~~~~~~~~~~~~~+-- We (reasonably) assume that end=off+len fits in an Int.+-- If that holds, this implementation /should/ work without encountering any+-- bugs due to overflow. But it is unclear how that can be tested without too+-- much trouble.
+ test/Main.hs view
@@ -0,0 +1,55 @@+{-# LANGUAGE MagicHash #-}++import qualified Data.Foldable as F+import qualified Data.List as L+import Data.Ord (comparing)+import Data.Primitive.Array+ ( MutableArray(..), arrayFromList, runArray, sizeofArray, thawArray )+import Data.Primitive.PrimArray+ ( MutablePrimArray(..)+ , primArrayFromList+ , primArrayToList+ , runPrimArray+ , sizeofPrimArray+ , thawPrimArray+ )++import Test.Tasty (defaultMain, localOption, testGroup)+import Test.Tasty.QuickCheck (QuickCheckTests(..), Fun, applyFun, testProperty, (===))+import Test.QuickCheck.Poly (A, OrdA)++import Data.SamSort (sortArrayBy, sortIntArrayBy)++main :: IO ()+main = defaultMain $ localOption (QuickCheckTests 5000) $ testGroup "Tests"+ [ testProperty "sortArrayBy" $ \xs ys zs ->+ sortViaMutableArray (comparing fst) (xs,ys,zs)+ ===+ ((xs :: [(OrdA, A)]) ++ L.sortBy (comparing fst) ys ++ zs)+ , testProperty "sortIntArrayBy" $ \f xs ys zs ->+ sortViaMutableIntArray+ (comparing (applyFun (f :: Fun Int OrdA)))+ (xs,ys,zs)+ ===+ (xs ++ L.sortBy (comparing (applyFun f)) ys ++ zs)+ ]++sortViaMutableArray+ :: (a -> a -> Ordering)+ -> ([a], [a], [a])+ -> [a]+sortViaMutableArray cmp (xs,ys,zs) = F.toList $ runArray $ do+ let a = arrayFromList (xs ++ ys ++ zs)+ ma@(MutableArray ma#) <- thawArray a 0 (sizeofArray a)+ sortArrayBy cmp ma# (length xs) (length ys)+ pure ma++sortViaMutableIntArray+ :: (Int -> Int -> Ordering)+ -> ([Int], [Int], [Int])+ -> [Int]+sortViaMutableIntArray cmp (xs,ys,zs) = primArrayToList $ runPrimArray $ do+ let a = primArrayFromList (xs ++ ys ++ zs)+ ma@(MutablePrimArray ma#) <- thawPrimArray a 0 (sizeofPrimArray a)+ sortIntArrayBy cmp ma# (length xs) (length ys)+ pure ma