From 6567bb862d0cf192d21358f8baea72168b79d069 Mon Sep 17 00:00:00 2001
From: 8084 <klntsky@gmail.com>
Date: Tue, 2 Apr 2019 19:22:05 +0300
Subject: [PATCH] added: pop, shift, unshift, unshiftAll to Data.Array.ST

---
 src/Data/Array/ST.js         | 28 ++++++++++++++
 src/Data/Array/ST.purs       | 31 ++++++++++++++-
 test/Test/Data/Array/ST.purs | 75 ++++++++++++++++++++++++++++++++++++
 3 files changed, 133 insertions(+), 1 deletion(-)

diff --git a/src/Data/Array/ST.js b/src/Data/Array/ST.js
index 54ddf4d1..4dc6f484 100644
--- a/src/Data/Array/ST.js
+++ b/src/Data/Array/ST.js
@@ -28,6 +28,16 @@ exports.poke = function (i) {
   };
 };
 
+exports.popImpl = function (just) {
+  return function (nothing) {
+    return function (xs) {
+      return function () {
+        return xs.length ? just(xs.pop()) : nothing;
+      };
+    };
+  };
+};
+
 exports.pushAll = function (as) {
   return function (xs) {
     return function () {
@@ -36,6 +46,24 @@ exports.pushAll = function (as) {
   };
 };
 
+exports.shiftImpl = function (just) {
+  return function (nothing) {
+    return function (xs) {
+      return function () {
+        return xs.length ? just(xs.shift()) : nothing;
+      };
+    };
+  };
+};
+
+exports.unshiftAll = function (as) {
+  return function (xs) {
+    return function () {
+      return xs.unshift.apply(xs, as);
+    };
+  };
+};
+
 exports.splice = function (i) {
   return function (howMany) {
     return function (bs) {
diff --git a/src/Data/Array/ST.purs b/src/Data/Array/ST.purs
index 186f3e78..cda10a06 100644
--- a/src/Data/Array/ST.purs
+++ b/src/Data/Array/ST.purs
@@ -10,9 +10,13 @@ module Data.Array.ST
   , empty
   , peek
   , poke
-  , push
   , modify
+  , pop
+  , push
   , pushAll
+  , shift
+  , unshift
+  , unshiftAll
   , splice
   , sort
   , sortBy
@@ -83,6 +87,12 @@ thaw = copyImpl
 sort :: forall a h. Ord a => STArray h a -> ST h (STArray h a)
 sort = sortBy compare
 
+-- | Remove the first element from an array and return that element.
+shift :: forall h a. STArray h a -> ST h (Maybe a)
+shift = shiftImpl Just Nothing
+
+foreign import shiftImpl :: forall h a. (a -> Maybe a) -> Maybe a -> STArray h a -> ST h (Maybe a)
+
 -- | Sort a mutable array in place using a comparison function.
 sortBy
   :: forall a h
@@ -136,6 +146,12 @@ foreign import peekImpl
 -- | Change the value at the specified index in a mutable array.
 foreign import poke :: forall h a. Int -> a -> STArray h a -> ST h Boolean
 
+-- | Remove the last element from an array and return that element.
+pop :: forall h a. STArray h a -> ST h (Maybe a)
+pop = popImpl Just Nothing
+
+foreign import popImpl :: forall h a. (a -> Maybe a) -> Maybe a -> STArray h a -> ST h (Maybe a)
+
 -- | Append an element to the end of a mutable array. Returns the new length of
 -- | the array.
 push :: forall h a. a -> STArray h a -> ST h Int
@@ -149,6 +165,19 @@ foreign import pushAll
   -> STArray h a
   -> ST h Int
 
+-- | Append an element to the front of a mutable array. Returns the new length of
+-- | the array.
+unshift :: forall h a. a -> STArray h a -> ST h Int
+unshift a = unshiftAll [a]
+
+-- | Append the values in an immutable array to the front of a mutable array.
+-- | Returns the new length of the mutable array.
+foreign import unshiftAll
+  :: forall h a
+   . Array a
+  -> STArray h a
+  -> ST h Int
+
 -- | Mutate the element at the specified index using the supplied function.
 modify :: forall h a. Int -> (a -> a) -> STArray h a -> ST h Boolean
 modify i f xs = do
diff --git a/test/Test/Data/Array/ST.purs b/test/Test/Data/Array/ST.purs
index c3a79b4d..35cf3471 100644
--- a/test/Test/Data/Array/ST.purs
+++ b/test/Test/Data/Array/ST.purs
@@ -46,6 +46,25 @@ testArrayST = do
 
   assert $ STA.run (STA.unsafeThaw [1, 2, 3]) == [1, 2, 3]
 
+  log "pop should remove elements from an STArray"
+
+  assert $ STA.run (do
+    arr <- STA.thaw [1, 2, 3]
+    void $ STA.pop arr
+    pure arr) == [1, 2]
+
+  log "pop should return the last element of an STArray"
+
+  assert $ ST.run (do
+    arr <- STA.thaw [1, 2, 3]
+    STA.pop arr) == Just 3
+
+  log "pop should return Nothing when given an empty array"
+
+  assert $ isNothing $ ST.run (do
+    arr <- STA.empty
+    STA.pop arr)
+
   log "push should append a value to the end of the array"
 
   assert $ STA.run (do
@@ -145,6 +164,62 @@ testArrayST = do
     void $ STA.poke 1 2 arr
     pure arr) == [1]
 
+  log "shift should remove elements from an STArray"
+
+  assert $ STA.run (do
+    arr <- STA.thaw [1, 2, 3]
+    void $ STA.shift arr
+    pure arr) == [2, 3]
+
+  log "shift should return the first element of an STArray"
+
+  assert $ ST.run (do
+    arr <- STA.thaw [1, 2, 3]
+    STA.shift arr) == Just 1
+
+  log "shift should return Nothing when given an empty array"
+
+  assert $ isNothing $ ST.run (do
+    arr <- STA.empty
+    STA.shift arr)
+
+  log "unshift should append a value to the front of the array"
+
+  assert $ STA.run (do
+    arr <- STA.empty
+    void $ STA.unshift 1 arr
+    void $ STA.unshift 2 arr
+    pure arr) == [2, 1]
+
+  assert $ STA.run (do
+    arr <- STA.thaw [1, 2, 3]
+    void $ STA.unshift 4 arr
+    pure arr) == [4, 1, 2, 3]
+
+  log "unshift should return the new length of the array"
+
+  assert $ ST.run (do
+    arr <- STA.thaw [unit, unit, unit]
+    STA.unshift unit arr) == 4
+
+  log "unshiftAll should append multiple values to the front of the array"
+
+  assert $ STA.run (do
+    arr <- STA.empty
+    void $ STA.unshiftAll [1, 2] arr
+    pure arr) == [1, 2]
+
+  assert $ STA.run (do
+    arr <- STA.thaw [1, 2, 3]
+    void $ STA.unshiftAll [4, 5, 6] arr
+    pure arr) == [4, 5, 6, 1, 2, 3]
+
+  log "unshiftAll should return the new length of the array"
+
+  assert $ ST.run (do
+    arr <- STA.thaw [unit, unit, unit]
+    STA.unshiftAll [unit, unit] arr) == 5
+
   log "sort should reorder a list into ascending order based on the result of compare"
   assert $ STA.run (
     STA.sort =<< STA.unsafeThaw [1, 3, 2, 5, 6, 4]