Iterator groupby (#84)

* Add Data.Array.ST.Iterator

* Add runSTArray'

For instances where you want to return more than just one STArray from
an STArray computation. For example, a pair of STArrays (with
potentially different types), or an STArray of STArrays, or even an
STArray of nonempty STArrays. The possibilities are endless!

Also use unsafe-coerce instead of FFI.

* Faster `groupBy`

Implement `groupBy` using `Data.Array.ST.Iterator`; this should make it
faster.

* Fix groupBy tests

\x y -> odd x && odd y is not an equivalence relation! It doesn't
satisfy reflexivity.

* Restrict to a closed row in runSTArray

Author: hdgarrood

bower.json

@@ -22,7 +22,8 @@
     "purescript-st": "^2.0.0",
     "purescript-tailrec": "^2.0.0",
     "purescript-tuples": "^3.0.0",
-    "purescript-unfoldable": "^2.0.0"
+    "purescript-unfoldable": "^2.0.0",
+    "purescript-unsafe-coerce": "^2.0.0"
   },
   "devDependencies": {
     "purescript-assert": "^2.0.0",

src/Data/Array.purs

@@ -115,11 +115,16 @@ import Prelude
 import Control.Alt ((<|>))
 import Control.Alternative (class Alternative)
 import Control.Lazy (class Lazy, defer)
+import Control.Monad.Eff (Eff)
 import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM2)
+import Control.Monad.ST (ST)
 
+import Data.Array.ST (STArray, emptySTArray, pushSTArray, runSTArray')
+import Data.Array.ST.Iterator (iterate, iterator, pushWhile)
 import Data.Foldable (class Foldable, foldl, foldr)
 import Data.Foldable (foldl, foldr, foldMap, fold, intercalate, elem, notElem, find, findMap, any, all) as Exports
 import Data.Maybe (Maybe(..), maybe, isJust, fromJust)
+import Data.Newtype (class Newtype, unwrap)
 import Data.NonEmpty (NonEmpty, (:|))
 import Data.Traversable (scanl, scanr) as Exports
 import Data.Traversable (sequence, traverse)
@@ -548,14 +553,24 @@ group' = group <<< sort
 -- | Group equal, consecutive elements of an array into arrays, using the
 -- | specified equivalence relation to detemine equality.
 groupBy :: forall a. (a -> a -> Boolean) -> Array a -> Array (NonEmpty Array a)
-groupBy op = go []
+groupBy op xs =
+  runGroupedSTArray do
+    result <- emptySTArray
+    iter <- iterator (xs !! _)
+    iterate iter \x -> do
+      sub <- emptySTArray
+      pushSTArray result (x :| sub)
+      pushWhile (op x) iter sub
+    pure result
   where
-  go :: Array (NonEmpty Array a) -> Array a -> Array (NonEmpty Array a)
-  go acc xs = case uncons xs of
-    Just o ->
-      let sp = span (op o.head) o.tail
-      in go ((o.head :| sp.init) : acc) sp.rest
-    Nothing -> reverse acc
+  runGroupedSTArray
+    :: forall b
+     . (forall h. Eff (st :: ST h) (STArray h (NonEmpty (STArray h) b)))
+    -> Array (NonEmpty Array b)
+  runGroupedSTArray a = unwrap (runSTArray' (map Grouped a))
+
+newtype Grouped a arr = Grouped (arr (NonEmpty arr a))
+derive instance newtypeGrouped :: Newtype (Grouped a arr) _
 
 -- | Remove the duplicates from an array, creating a new array.
 nub :: forall a. Eq a => Array a -> Array a

src/Data/Array/ST.js

@@ -1,9 +1,5 @@
 "use strict";
 
-exports.runSTArray = function (f) {
-  return f;
-};
-
 exports.emptySTArray = function () {
   return [];
 };

src/Data/Array/ST.purs

@@ -6,6 +6,7 @@ module Data.Array.ST
   ( STArray(..)
   , Assoc()
   , runSTArray
+  , runSTArray'
   , emptySTArray
   , peekSTArray
   , pokeSTArray
@@ -16,10 +17,12 @@ module Data.Array.ST
   , toAssocArray
   ) where
 
-import Control.Monad.Eff (Eff)
+import Prelude
+import Control.Monad.Eff (Eff, Pure, runPure)
 import Control.Monad.ST (ST)
-
 import Data.Maybe (Maybe(..))
+import Data.Newtype (class Newtype, unwrap)
+import Unsafe.Coerce (unsafeCoerce)
 
 -- | A reference to a mutable array.
 -- |
@@ -33,14 +36,48 @@ foreign import data STArray :: * -> * -> *
 -- | An element and its index.
 type Assoc a = { value :: a, index :: Int }
 
+newtype Id a f = Id (f a)
+
+derive instance newtypeId :: Newtype (Id a f) _
+
 -- | Freeze a mutable array, creating an immutable array. Use this function as you would use
 -- | `runST` to freeze a mutable reference.
 -- |
--- | The rank-2 type prevents the reference from escaping the scope of `runSTArray`.
-foreign import runSTArray
-  :: forall a r
-   . (forall h. Eff (st :: ST h | r) (STArray h a))
-  -> Eff r (Array a)
+-- | The rank-2 type prevents the reference from escaping the scope of `runSTArray'`,
+-- | and the closed row on the `Eff` computation prevents the reference from
+-- | escaping into other parts of your program via native effects such as `setTimeout`.
+-- |
+-- | You can also return an immutable copy of an `STArray` from an `ST` computation
+-- | by using `freeze` combined with `runST`. However, when possible, you should
+-- | prefer this function, because it is `O(1)`. By contrast, `freeze` must copy the
+-- | underlying array and is therefore `O(n)`.
+runSTArray
+  :: forall a
+   . (forall h. Eff (st :: ST h) (STArray h a))
+  -> Array a
+runSTArray a = unwrap (runSTArray' (map Id a))
+
+-- | Freeze all mutable arrays in some structure, creating a version of the
+-- | same structure where all mutable arrays are replaced with immutable
+-- | arrays. Use this function as you would use `runST` to freeze a mutable
+-- | reference.
+-- |
+-- | The rank-2 type prevents the reference from escaping the scope of `runSTArray'`,
+-- | and the closed row on the `Eff` computation prevents the reference from
+-- | escaping into other parts of your program via native effects such as `setTimeout`.
+-- |
+-- | You can also return an immutable copy of an `STArray` from an `ST` computation
+-- | by using `freeze` combined with `runST`. However, when possible, you should
+-- | prefer this function, because it is `O(1)`. By contrast, `freeze` must copy the
+-- | underlying array and is therefore `O(n)`.
+runSTArray'
+  :: forall f
+   . (forall h. Eff (st :: ST h) (f (STArray h)))
+  -> f Array
+runSTArray' x = runPure (go x)
+  where
+  go :: (forall h. Eff (st :: ST h) (f (STArray h))) -> Pure (f Array)
+  go = unsafeCoerce
 
 -- | Create an empty mutable array.
 foreign import emptySTArray :: forall a h r. Eff (st :: ST h | r) (STArray h a)

src/Data/Array/ST/Iterator.purs

@@ -0,0 +1,77 @@
+module Data.Array.ST.Iterator
+  ( Iterator
+  , iterator
+  , iterate
+  , next
+  , peek
+  , exhausted
+  , pushWhile
+  , pushAll
+  ) where
+
+import Prelude
+import Control.Monad.Eff (Eff, whileE)
+import Control.Monad.ST (ST, STRef, newSTRef, readSTRef, writeSTRef, modifySTRef)
+import Data.Array.ST (STArray, pushSTArray)
+
+import Data.Maybe (Maybe(..), isNothing)
+
+-- | This type provides a slightly easier way of iterating over an array's
+-- | elements in an STArray computation, without having to keep track of
+-- | indices.
+data Iterator h a = Iterator (Int -> Maybe a) (STRef h Int)
+
+-- | Make an Iterator given an indexing function into an array (or anything
+-- | else). If `xs :: Array a`, the standard way to create an iterator over
+-- | `xs` is to use `iterator (xs !! _)`, where `(!!)` comes from `Data.Array`.
+iterator :: forall a h r. (Int -> Maybe a) -> Eff (st :: ST h | r) (Iterator h a)
+iterator f =
+  Iterator f <$> newSTRef 0
+
+-- | Perform an action once for each item left in an iterator. If the action
+-- | itself also advances the same iterator, `iterate` will miss those items
+-- | out.
+iterate :: forall a h r. Iterator h a -> (a -> Eff (st :: ST h | r) Unit) -> Eff (st :: ST h | r) Unit
+iterate iter f = do
+  break <- newSTRef false
+  whileE (not <$> readSTRef break) do
+    mx <- next iter
+    case mx of
+      Just x -> f x
+      Nothing -> void $ writeSTRef break true
+
+-- | Get the next item out of an iterator, advancing it. Returns Nothing if the
+-- | Iterator is exhausted.
+next :: forall a h r. Iterator h a -> Eff (st :: ST h | r) (Maybe a)
+next (Iterator f currentIndex) = do
+  i <- readSTRef currentIndex
+  modifySTRef currentIndex (_ + 1)
+  pure (f i)
+
+-- | Get the next item out of an iterator without advancing it.
+peek :: forall a h r. Iterator h a -> Eff (st :: ST h | r) (Maybe a)
+peek (Iterator f currentIndex) = do
+  i <- readSTRef currentIndex
+  pure (f i)
+
+-- | Check whether an iterator has been exhausted.
+exhausted :: forall a h r. Iterator h a -> Eff (st :: ST h | r) Boolean
+exhausted = map isNothing <<< peek
+
+-- | Extract elements from an iterator and push them on to an STArray for as
+-- | long as those elements satisfy a given predicate.
+pushWhile :: forall a h r. (a -> Boolean) -> Iterator h a -> STArray h a -> Eff (st :: ST h | r) Unit
+pushWhile p iter array = do
+  break <- newSTRef false
+  whileE (not <$> readSTRef break) do
+    mx <- peek iter
+    case mx of
+      Just x | p x -> do
+        pushSTArray array x
+        void $ next iter
+      _ ->
+        void $ writeSTRef break true
+
+-- | Push the entire remaining contents of an iterator onto an STArray.
+pushAll :: forall a h r. Iterator h a -> STArray h a -> Eff (st :: ST h | r) Unit
+pushAll = pushWhile (const true)

test/Test/Data/Array.purs

@@ -285,7 +285,7 @@ testArray = do
   assert $ A.group' [1, 2, 2, 3, 3, 3, 1] == [1 :| [1], 2 :| [2], 3 :| [3, 3]]
 
   log "groupBy should group consecutive equal elements into arrays based on an equivalence relation"
-  assert $ A.groupBy (\x y -> odd x && odd y) [1, 1, 2, 2, 3, 3] == [1 :| [1], NE.singleton 2, NE.singleton 2, 3 :| [3]]
+  assert $ A.groupBy eqParity [1, 1, 2, 2, 3, 5, 4] == [1 :| [1], 2 :| [2], 3 :| [5], NE.singleton 4]
 
   log "nub should remove duplicate elements from the list, keeping the first occurence"
   assert $ A.nub [1, 2, 2, 3, 4, 1] == [1, 2, 3, 4]
@@ -365,6 +365,11 @@ odd n = n `mod` 2 /= zero
 doubleAndOrig :: Int -> Array Int
 doubleAndOrig x = [x * 2, x]
 
+-- | An equivalence relation on integers. This relation splits the integers
+-- | into two equivalence classes: odds and evens.
+eqParity :: Int -> Int -> Boolean
+eqParity x y = x `mod` 2 == y `mod` 2
+
 data Replicated a = Replicated Int a
 
 instance foldableReplicated :: Foldable Replicated where

test/Test/Data/Array/ST.purs

@@ -17,11 +17,11 @@ testArrayST = do
 
   log "emptySTArray should produce an empty array"
 
-  assert $ runPure (runSTArray emptySTArray) == nil
+  assert $ runSTArray emptySTArray == nil
 
   log "thaw should produce an STArray from a standard array"
 
-  assert $ runPure (runSTArray (thaw [1, 2, 3])) == [1, 2, 3]
+  assert $ runSTArray (thaw [1, 2, 3]) == [1, 2, 3]
 
   log "freeze should produce a standard array from an STArray"
 
@@ -31,28 +31,28 @@ testArrayST = do
 
   log "pushSTArray should append a value to the end of the array"
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- emptySTArray
     pushSTArray arr 1
     pushSTArray arr 2
-    pure arr)) == [1, 2]
+    pure arr) == [1, 2]
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- thaw [1, 2, 3]
     pushSTArray arr 4
-    pure arr)) == [1, 2, 3, 4]
+    pure arr) == [1, 2, 3, 4]
 
   log "pushAllSTArray should append multiple values to the end of the array"
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- emptySTArray
     pushAllSTArray arr [1, 2]
-    pure arr)) == [1, 2]
+    pure arr) == [1, 2]
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- thaw [1, 2, 3]
     pushAllSTArray arr [4, 5, 6]
-    pure arr)) == [1, 2, 3, 4, 5, 6]
+    pure arr) == [1, 2, 3, 4, 5, 6]
 
   log "peekSTArray should return Nothing when peeking a value outside the array bounds"
 
@@ -104,24 +104,24 @@ testArrayST = do
 
   log "pokeSTArray should replace the value at the specified index"
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- thaw [1]
     pokeSTArray arr 0 10
-    pure arr)) == [10]
+    pure arr) == [10]
 
   log "pokeSTArray should do nothing when attempting to modify a value outside the array bounds"
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- thaw [1]
     pokeSTArray arr 1 2
-    pure arr)) == [1]
+    pure arr) == [1]
 
   log "spliceSTArray should be able to delete multiple items at a specified index"
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- thaw [1, 2, 3, 4, 5]
     spliceSTArray arr 1 3 []
-    pure arr)) == [1, 5]
+    pure arr) == [1, 5]
 
   log "spliceSTArray should return the items removed"
 
@@ -131,17 +131,17 @@ testArrayST = do
 
   log "spliceSTArray should be able to insert multiple items at a specified index"
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- thaw [1, 2, 3, 4, 5]
     spliceSTArray arr 1 0 [0, 100]
-    pure arr)) == [1, 0, 100, 2, 3, 4, 5]
+    pure arr) == [1, 0, 100, 2, 3, 4, 5]
 
   log "spliceSTArray should be able to delete and insert at the same time"
 
-  assert $ runPure (runSTArray (do
+  assert $ runSTArray (do
     arr <- thaw [1, 2, 3, 4, 5]
     spliceSTArray arr 1 2 [0, 100]
-    pure arr)) == [1, 0, 100, 4, 5]
+    pure arr) == [1, 0, 100, 4, 5]
 
   log "toAssocArray should return all items in the array with the correct indices and values"