Speed up 'difference', add findIndexSTArray

Author: hdgarrood

src/Data/Array.purs

@@ -114,6 +114,7 @@ import Control.Alt ((<|>))
 import Control.Alternative (class Alternative)
 import Control.Lazy (class Lazy, defer)
 import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM2)
+import Control.Monad.Eff (runPure, foreachE)
 
 import Data.Foldable (class Foldable, foldl, foldr)
 import Data.Foldable (foldl, foldr, foldMap, fold, intercalate, elem, notElem, find, findMap, any, all) as Exports
@@ -123,6 +124,7 @@ import Data.Traversable (scanl, scanr) as Exports
 import Data.Traversable (sequence)
 import Data.Tuple (Tuple(..))
 import Data.Unfoldable (class Unfoldable, unfoldr)
+import Data.Array.ST (runSTArray, pokeSTArray, findIndexSTArray, thaw)
 
 import Partial.Unsafe (unsafePartial)
 
@@ -583,9 +585,16 @@ deleteBy eq x ys = maybe ys (\i -> unsafePartial $ fromJust (deleteAt i ys)) (fi
 -- | Delete the first occurrence of each element in the second array from the
 -- | first array, creating a new array.
 difference :: forall a. Eq a => Array a -> Array a -> Array a
-difference xs ys
-  | null xs = []
-  | otherwise = uncons' (const xs) (\z zs -> delete z xs \\ zs) ys
+difference xs ys = catMaybes $ runPure $ runSTArray do
+  zs <- thaw (map Just xs)
+  foreachE ys \y -> do
+    idx <- findIndexSTArray (_ == Just y) zs
+    case idx of
+      Just i ->
+        void $ pokeSTArray zs i Nothing
+      _ ->
+        pure unit
+  pure zs
 
 infix 5 difference as \\
 

src/Data/Array/ST.js

@@ -58,6 +58,21 @@ exports.copyImpl = function (xs) {
   };
 };
 
+exports.findIndexSTArrayImpl = function (just) {
+  return function (nothing) {
+    return function (f) {
+      return function (xs) {
+        return function () {
+          for (var i = 0, l = xs.length; i < l; i++) {
+            if (f(xs[i])) return just(i);
+          }
+          return nothing;
+        };
+      };
+    };
+  };
+};
+
 exports.toAssocArray = function (xs) {
   return function () {
     var n = xs.length;

src/Data/Array/ST.purs

@@ -12,6 +12,7 @@ module Data.Array.ST
   , pushSTArray
   , pushAllSTArray
   , spliceSTArray
+  , findIndexSTArray
   , freeze, thaw
   , toAssocArray
   ) where
@@ -96,6 +97,22 @@ foreign import spliceSTArray
   -> Array a
   -> Eff (st :: ST h | r) (Array a)
 
+-- | Find the first index for which a predicate holds.
+findIndexSTArray
+  :: forall a h r
+   . (a -> Boolean)
+  -> STArray h a
+  -> Eff (st :: ST h | r) (Maybe Int)
+findIndexSTArray = findIndexSTArrayImpl Just Nothing
+
+foreign import findIndexSTArrayImpl
+  :: forall a h r
+   . (forall t. t -> Maybe t)
+  -> (forall t. Maybe t)
+  -> (a -> Boolean)
+  -> STArray h a
+  -> Eff (st :: ST h | r) (Maybe Int)
+
 -- | Create an immutable copy of a mutable array, where each element
 -- | is labelled with its index in the original array.
 foreign import toAssocArray