STArray: in-place sort

Author: matthewleon

src/Data/Array/ST.js

@@ -58,6 +58,16 @@ exports.copyImpl = function (xs) {
   };
 };
 
+exports.sortByImpl = function (comp) {
+  return function (xs) {
+    return function () {
+      return xs.sort(function (x, y) {
+        return comp(x)(y);
+      });
+    };
+  };
+};
+
 exports.toAssocArray = function (xs) {
   return function () {
     var n = xs.length;

src/Data/Array/ST.purs

@@ -14,6 +14,9 @@ module Data.Array.ST
   , modifySTArray
   , pushAllSTArray
   , spliceSTArray
+  , sort
+  , sortBy
+  , sortWith
   , freeze
   , thaw
   , unsafeFreeze
@@ -80,6 +83,38 @@ foreign import emptySTArray :: forall a h r. Eff (st :: ST h | r) (STArray h a)
 thaw :: forall a h r. Array a -> Eff (st :: ST h | r) (STArray h a)
 thaw = copyImpl
 
+-- | Sort a mutable array in place.
+sort :: forall a h r. Ord a => STArray h a -> Eff (st :: ST h | r) (STArray h a)
+sort = sortBy compare
+
+-- | Sort a mutable array in place using a comparison function.
+sortBy
+  :: forall a h r
+   . (a -> a -> Ordering)
+  -> STArray h a
+  -> Eff (st :: ST h | r) (STArray h a)
+sortBy comp = sortByImpl comp'
+  where
+  comp' x y = case comp x y of
+    GT -> 1
+    EQ -> 0
+    LT -> -1
+
+foreign import sortByImpl
+  :: forall a h r
+   . (a -> a -> Int)
+  -> STArray h a
+  -> Eff (st :: ST h | r) (STArray h a)
+
+-- | Sort a mutable array in place based on a projection.
+sortWith
+  :: forall a b h r
+   . Ord b
+  => (a -> b)
+  -> STArray h a
+  -> Eff (st :: ST h | r) (STArray h a)
+sortWith f = sortBy (comparing f)
+
 -- | Create an immutable copy of a mutable array.
 freeze :: forall a h r. STArray h a -> Eff (st :: ST h | r) (Array a)
 freeze = copyImpl

test/Test/Data/Array/ST.purs

@@ -4,7 +4,7 @@ import Prelude
 import Control.Monad.Eff (Eff)
 import Control.Monad.Eff.Console (log, CONSOLE)
 import Control.Monad.ST (ST, pureST)
-import Data.Array.ST (STArray, emptySTArray, freeze, peekSTArray, pokeSTArray, pushAllSTArray, pushSTArray, spliceSTArray, thaw, toAssocArray, unsafeThaw, unsafeFreeze)
+import Data.Array.ST (STArray, emptySTArray, freeze, peekSTArray, pokeSTArray, pushAllSTArray, pushSTArray, sort, sortBy, sortWith, spliceSTArray, thaw, toAssocArray, unsafeThaw, unsafeFreeze)
 import Data.Foldable (all)
 import Data.Maybe (Maybe(..), isNothing)
 import Test.Assert (assert, ASSERT)
@@ -132,6 +132,21 @@ testArrayST = do
     void $ pokeSTArray arr 1 2
     pure arr) == [1]
 
+  log "sort should reorder a list into ascending order based on the result of compare"
+  assert $ run (
+    sort =<< unsafeThaw [1, 3, 2, 5, 6, 4]
+  ) == [1, 2, 3, 4, 5, 6]
+
+  log "sortBy should reorder a list into ascending order based on the result of a comparison function"
+  assert $ run (
+    sortBy (flip compare) =<< unsafeThaw [1, 3, 2, 5, 6, 4]
+  ) == [6, 5, 4, 3, 2, 1]
+
+  log "sortWith should reorder a list into ascending order based on the result of compare over a projection"
+  assert $ run (
+    sortWith id =<< unsafeThaw [1, 3, 2, 5, 6, 4]
+  ) == [1, 2, 3, 4, 5, 6]
+
   log "spliceSTArray should be able to delete multiple items at a specified index"
 
   assert $ run (do