diff --git a/CHANGELOG.md b/CHANGELOG.md index 387f4a3..66189be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ Breaking changes: New features: Bugfixes: +- Make `groupBy` use pairwise function (#230 by @JordanMartinez) Other improvements: diff --git a/src/Data/Array.purs b/src/Data/Array.purs index 6bc035a..4ffd04b 100644 --- a/src/Data/Array.purs +++ b/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 diff --git a/test/Test/Data/Array.purs b/test/Test/Data/Array.purs index 9325c0f..b9804e6 100644 --- a/test/Test/Data/Array.purs +++ b/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]]