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-05-05 + +* First version. diff --git a/HOWTO.md b/HOWTO.md new file mode 100644 --- /dev/null +++ b/HOWTO.md @@ -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). 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,53 @@ +# samsort + +[![Hackage](https://img.shields.io/hackage/v/samsort?logo=haskell&color=blue)](https://hackage.haskell.org/package/samsort) +[![Haskell-CI](https://github.com/meooow25/samsort/actions/workflows/haskell-ci.yml/badge.svg)](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! diff --git a/samsort.cabal b/samsort.cabal new file mode 100644 --- /dev/null +++ b/samsort.cabal @@ -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 diff --git a/src/Data/SamSort.hs b/src/Data/SamSort.hs new file mode 100644 --- /dev/null +++ b/src/Data/SamSort.hs @@ -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. diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -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