Make groupBy use pairwise function

CHANGELOG.md

@@ -9,6 +9,7 @@ Breaking changes:
 New features:
 
 Bugfixes:
+- Make `groupBy` use pairwise function (#230 by @JordanMartinez)
 
 Other improvements:
 

src/Data/Array.purs

@@ -137,6 +137,7 @@ import Control.Alternative (class Alternative)
 import Control.Lazy (class Lazy, defer)
 import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM2)
 import Control.Monad.ST as ST
+import Control.Monad.ST.Ref as STRef
 import Data.Array.NonEmpty.Internal (NonEmptyArray(..))
 import Data.Array.ST as STA
 import Data.Array.ST.Iterator as STAI
@@ -990,16 +991,42 @@ groupAll = groupAllBy compare
 -- | ```
 -- |
 groupBy :: forall a. (a -> a -> Boolean) -> Array a -> Array (NonEmptyArray a)
-groupBy op xs =
-  ST.run do
+groupBy op xs = case xs !! 0 of
+  Nothing -> 
+    []
+  Just h -> ST.run do
+    firstGroup <- STA.new
+    void $ STA.push h firstGroup
+    currentGroupRef <- STRef.new firstGroup
+
     result <- STA.new
-    iter <- STAI.iterator (xs !! _)
-    STAI.iterate iter \x -> void do
-      sub <- STA.new
-      _ <- STA.push x sub
-      STAI.pushWhile (op x) iter sub
-      grp <- STA.unsafeFreeze sub
-      STA.push (NonEmptyArray grp) result
+    prevRef <- STRef.new h
+    currentIdxRef <- STRef.new 1
+    ST.while (notEq 0 <$> STRef.read currentIdxRef) do
+      currentIdx <- STRef.read currentIdxRef
+      case xs !! currentIdx of
+        Just nextEl -> do
+          _ <- STRef.write (currentIdx + 1) currentIdxRef
+          prevEl <- STRef.read prevRef
+          currentGroup <- STRef.read currentGroupRef
+          if op prevEl nextEl then do
+            void $ STA.push nextEl currentGroup
+          else do
+            grp <- STA.unsafeFreeze currentGroup
+            _ <- STA.push (NonEmptyArray grp) result
+            nextGroup <- STA.new
+            _ <- STA.push nextEl nextGroup
+            void $ STRef.write nextGroup currentGroupRef
+          -- I would put this line immediately below 
+          --    `prevEl <- STRef.read prefRef`
+          -- but doing so causes the first `groupBy` test to fail,
+          -- likely due to a bug in how the compiler handles STRef.
+          void $ STRef.write nextEl prevRef
+        Nothing -> do
+          _ <- STRef.write 0 currentIdxRef
+          currentGroup <- STRef.read currentGroupRef
+          grp <- STA.unsafeFreeze currentGroup
+          void $ STA.push (NonEmptyArray grp) result
     STA.unsafeFreeze result
 
 -- | Group equal elements of an array into arrays, using the specified

test/Test/Data/Array.purs

@@ -366,7 +366,11 @@ testArray = do
   assert $ A.groupAll [1, 2, 2, 3, 3, 3, 1] == [nea [1, 1], nea [2, 2], nea [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] == [nea [1, 1], nea [2], nea [2], nea [3, 3]]
+  assert $ A.groupBy (\x y -> odd x && odd y) [1, 1, 2, 2, 3, 3] == [nea [1, 1], NEA.singleton 2, NEA.singleton 2, nea [3, 3]]
+
+  log "groupBy uses a pairwise function"
+  assert $ A.groupBy (\l r -> l + 1 == r) [1, 2, 3, 4] == [nea [1, 2, 3, 4]]
+  assert $ A.groupBy (\l r -> l + 1 == r) [1, 3, 4, 6] == [nea [1], nea [3, 4], nea [6]]
 
   log "groupBy should be stable"
   assert $ A.groupBy (\_ _ -> true) [1, 2, 3] == [nea [1, 2, 3]]