From b6c5f5ba02e85e584ef41277e6c4398fd77ff7c3 Mon Sep 17 00:00:00 2001
From: Matthew Leon <ml@matthewleon.com>
Date: Sun, 28 May 2017 02:53:45 +0100
Subject: [PATCH] batch update and modify functions (#101)

* update and modify at indices

Also exports ST utility functions `modifySTArray` and `withArray`.

* change export ordering
---
 src/Data/Array.purs       | 19 ++++++++++++++++---
 src/Data/Array/ST.purs    | 22 ++++++++++++++++++++++
 test/Test/Data/Array.purs | 10 ++++++++++
 3 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/Data/Array.purs b/src/Data/Array.purs
index 366fd75d..3745d68f 100644
--- a/src/Data/Array.purs
+++ b/src/Data/Array.purs
@@ -59,7 +59,9 @@ module Data.Array
   , insertAt
   , deleteAt
   , updateAt
+  , updateAtIndices
   , modifyAt
+  , modifyAtIndices
   , alterAt
 
   , reverse
@@ -115,15 +117,15 @@ import Control.Alternative (class Alternative)
 import Control.Lazy (class Lazy, defer)
 import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM2)
 import Control.Monad.ST (pureST)
-import Data.Array.ST (unsafeFreeze, emptySTArray, pushSTArray)
+import Data.Array.ST (unsafeFreeze, emptySTArray, pokeSTArray, pushSTArray, modifySTArray, withArray)
 import Data.Array.ST.Iterator (iterator, iterate, pushWhile)
-import Data.Foldable (class Foldable, foldl, foldr)
+import Data.Foldable (class Foldable, foldl, foldr, traverse_)
 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.NonEmpty (NonEmpty, (:|))
 import Data.Traversable (scanl, scanr) as Exports
 import Data.Traversable (sequence, traverse)
-import Data.Tuple (Tuple(..))
+import Data.Tuple (Tuple(..), uncurry)
 import Data.Unfoldable (class Unfoldable, unfoldr)
 import Partial.Unsafe (unsafePartial)
 
@@ -445,6 +447,17 @@ mapWithIndex :: forall a b. (Int -> a -> b) -> Array a -> Array b
 mapWithIndex f xs =
   zipWith f (range 0 (length xs - 1)) xs
 
+-- | Change the elements at the specified indices in index/value pairs.
+-- | Out-of-bounds indices will have no effect.
+updateAtIndices :: forall t a. Foldable t => t (Tuple Int a) -> Array a -> Array a
+updateAtIndices us xs =
+  pureST (withArray (\res -> traverse_ (uncurry $ pokeSTArray res) us) xs)
+
+-- | Apply a function to the element at the specified indices,
+-- | creating a new array. Out-of-bounds indices will have no effect.
+modifyAtIndices :: forall t a. Foldable t => t Int -> (a -> a) -> Array a -> Array a
+modifyAtIndices is f xs =
+  pureST (withArray (\res -> traverse_ (\i -> modifySTArray res i f) is) xs)
 
 --------------------------------------------------------------------------------
 -- Sorting ---------------------------------------------------------------------
diff --git a/src/Data/Array/ST.purs b/src/Data/Array/ST.purs
index ce953116..08b0dd05 100644
--- a/src/Data/Array/ST.purs
+++ b/src/Data/Array/ST.purs
@@ -6,10 +6,12 @@ module Data.Array.ST
   ( STArray(..)
   , Assoc()
   , runSTArray
+  , withArray
   , emptySTArray
   , peekSTArray
   , pokeSTArray
   , pushSTArray
+  , modifySTArray
   , pushAllSTArray
   , spliceSTArray
   , freeze, thaw
@@ -47,6 +49,18 @@ foreign import runSTArray
    . (forall h. Eff (st :: ST h | r) (STArray h a))
   -> Eff r (Array a)
 
+-- Perform an effect requiring a mutable array on a copy of an immutable array,
+-- safely returning the result as an immutable array.
+withArray
+  :: forall a b r h
+   . (STArray h a -> Eff (st :: ST h | r) b)
+   -> Array a
+   -> Eff (st :: ST h | r) (Array a)
+withArray f xs = do
+  result <- thaw xs
+  _ <- f result
+  unsafeFreeze result
+
 -- | O(1). Convert a mutable array to an immutable array, without copying. The mutable
 -- | array must not be mutated afterwards.
 unsafeFreeze :: forall a r h. STArray h a -> Eff (st :: ST h | r) (Array a)
@@ -97,6 +111,14 @@ foreign import pushAllSTArray
   -> Array a
   -> Eff (st :: ST h | r) Int
 
+-- | Mutate the element at the specified index using the supplied function.
+modifySTArray :: forall a h r. STArray h a -> Int -> (a -> a) -> Eff (st :: ST h | r) Boolean
+modifySTArray xs i f = do
+  entry <- peekSTArray xs i
+  case entry of
+    Just x  -> pokeSTArray xs i (f x)
+    Nothing -> pure false
+
 -- | Remove and/or insert elements from/into a mutable array at the specified index.
 foreign import spliceSTArray
   :: forall a h r
diff --git a/test/Test/Data/Array.purs b/test/Test/Data/Array.purs
index 2e1c1e18..2436afb1 100644
--- a/test/Test/Data/Array.purs
+++ b/test/Test/Data/Array.purs
@@ -222,6 +222,16 @@ testArray = do
   log "mapWithIndex applies a function with an index for every element"
   assert $ A.mapWithIndex (\i x -> x - i) [9,8,7,6,5] == [9,7,5,3,1]
 
+  log "updateAtIndices changes the elements at specified indices"
+  assert $ A.updateAtIndices
+             [Tuple 0 false, Tuple 2 false, Tuple 8 false]
+             [true,  true, true,  true] ==
+             [false, true, false, true]
+
+  log "modifyAtIndices modifies the elements at specified indices"
+  assert $ A.modifyAtIndices [0, 2, 8] not [true,  true, true,  true] ==
+                                           [false, true, false, true]
+
   log "sort should reorder a list into ascending order based on the result of compare"
   assert $ A.sort [1, 3, 2, 5, 6, 4] == [1, 2, 3, 4, 5, 6]