Repository: sanctuary-js/sanctuary
Pull request #627 (2019-04-25 20:20:48)

array: generalize S.take and S.drop
The code is more complex than I would like it to be, but I believe the complexity is necessary. Please correct me if I'm wrong about this! :)


davidchambers: > I think it should be possible to do it in one pass

Could you elaborate, Brad? I believe we're traversing the structure just once, with `Z.reduce`.

Bradcomp: Ugh, I am tired today, but I see what is happening now. `reduce` is returning a `Just Pair (Int (f a))`, and we are rejecting if the int is > 0 because it means the foldable wasn't long enough. Otherwise we take the resultant array.

I had things inside out in my thinking.

Looks good to me! 

Patch:
diff --git a/index.js b/index.js
index 5b6b475..e900141 100644
--- a/index.js
+++ b/index.js
@@ -2916,57 +2916,90 @@
     impl: init
   };
 
-  //# take :: Integer -> Array a -> Maybe (Array a)
+  //# take :: (Applicative f, Foldable f, Monoid (f a)) => Integer -> f a -> Maybe (f a)
   //.
-  //. Returns Just the first N elements of the given array if N is greater
-  //. than or equal to zero and less than or equal to the length of the array;
+  //. Returns Just the first N elements of the given structure if N is
+  //. non-negative and less than or equal to the size of the structure;
   //. Nothing otherwise.
   //.
   //. ```javascript
-  //. > S.take (2) (['a', 'b', 'c', 'd', 'e'])
-  //. Just (['a', 'b'])
+  //. > S.take (0) (['foo', 'bar'])
+  //. Just ([])
   //.
-  //. > S.take (5) (['a', 'b', 'c', 'd', 'e'])
-  //. Just (['a', 'b', 'c', 'd', 'e'])
+  //. > S.take (1) (['foo', 'bar'])
+  //. Just (['foo'])
+  //.
+  //. > S.take (2) (['foo', 'bar'])
+  //. Just (['foo', 'bar'])
   //.
-  //. > S.take (6) (['a', 'b', 'c', 'd', 'e'])
+  //. > S.take (3) (['foo', 'bar'])
   //. Nothing
+  //.
+  //. > S.take (3) (Cons (1) (Cons (2) (Cons (3) (Cons (4) (Cons (5) (Nil))))))
+  //. Just (Cons (1) (Cons (2) (Cons (3) (Nil))))
   //. ```
-  function take(n) {
-    return function(xs) {
-      return n >= 0 && n <= xs.length ? Just (xs.slice (0, n)) : Nothing;
+  function _takeDrop(arrayCase, generalCase) {
+    return function(n) {
+      return function(xs) {
+        if (n < 0) return Nothing;
+
+        //  Fast path for arrays.
+        if (Array.isArray (xs)) {
+          return n <= xs.length ? Just (arrayCase (n, xs)) : Nothing;
+        }
+
+        //  m :: Maybe (Pair Integer (f a))
+        var m = Z.reduce (function(m, x) {
+          return Z.map (function(pair) {
+            var n = pair.fst;
+            var xs = pair.snd;
+            return Pair (n - 1) (generalCase (n, xs, x));
+          }, m);
+        }, Just (Pair (n) (Z.empty (xs.constructor))), xs);
+
+        return Z.map (Pair.snd, Z.reject (B (gt (0)) (Pair.fst), m));
+      };
     };
   }
+  var take = _takeDrop (
+    function(n, xs) { return xs.slice (0, n); },
+    function(n, xs, x) { return n > 0 ? Z.append (x, xs) : xs; }
+  );
   _.take = {
-    consts: {},
-    types: [$.Integer, $.Array (a), $.Maybe ($.Array (a))],
+    consts: {f: [Z.Applicative, Z.Foldable, Z.Monoid]},
+    types: [$.Integer, f (a), $.Maybe (f (a))],
     impl: take
   };
 
-  //# drop :: Integer -> Array a -> Maybe (Array a)
+  //# drop :: (Applicative f, Foldable f, Monoid (f a)) => Integer -> f a -> Maybe (f a)
   //.
-  //. Returns Just all but the first N elements of the given array if N is
-  //. greater than or equal to zero and less than or equal to the length of
-  //. the array; Nothing otherwise.
+  //. Returns Just all but the first N elements of the given structure if
+  //. N is non-negative and less than or equal to the size of the structure;
+  //. Nothing otherwise.
   //.
   //. ```javascript
-  //. > S.drop (2) (['a', 'b', 'c', 'd', 'e'])
-  //. Just (['c', 'd', 'e'])
+  //. > S.drop (0) (['foo', 'bar'])
+  //. Just (['foo', 'bar'])
+  //.
+  //. > S.drop (1) (['foo', 'bar'])
+  //. Just (['bar'])
   //.
-  //. > S.drop (5) (['a', 'b', 'c', 'd', 'e'])
+  //. > S.drop (2) (['foo', 'bar'])
   //. Just ([])
   //.
-  //. > S.drop (6) (['a', 'b', 'c', 'd', 'e'])
+  //. > S.drop (3) (['foo', 'bar'])
   //. Nothing
+  //.
+  //. > S.drop (3) (Cons (1) (Cons (2) (Cons (3) (Cons (4) (Cons (5) (Nil))))))
+  //. Just (Cons (4) (Cons (5) (Nil)))
   //. ```
-  function drop(n) {
-    return function(xs) {
-      return n >= 0 && n <= xs.length ? Just (xs.slice (n)) : Nothing;
-    };
-  }
+  var drop = _takeDrop (
+    function(n, xs) { return xs.slice (n); },
+    function(n, xs, x) { return n > 0 ? xs : Z.append (x, xs); }
+  );
   _.drop = {
-    consts: {},
-    types: [$.Integer, $.Array (a), $.Maybe ($.Array (a))],
+    consts: {f: [Z.Applicative, Z.Foldable, Z.Monoid]},
+    types: [$.Integer, f (a), $.Maybe (f (a))],
     impl: drop
   };
 


Test patch:
diff --git a/test/drop.js b/test/drop.js
index 2648344..8e9d6b8 100644
--- a/test/drop.js
+++ b/test/drop.js
@@ -1,7 +1,8 @@
 'use strict';
 
-const S = require ('..');
+const S = require ('./internal/sanctuary');
 
+const {Nil, Cons} = require ('./internal/List');
 const eq = require ('./internal/eq');
 
 
@@ -9,7 +10,7 @@ test ('drop', () => {
 
   eq (typeof S.drop) ('function');
   eq (S.drop.length) (1);
-  eq (S.show (S.drop)) ('drop :: Integer -> Array a -> Maybe (Array a)');
+  eq (S.show (S.drop)) ('drop :: (Applicative f, Foldable f, Monoid f) => Integer -> f a -> Maybe (f a)');
 
   eq (S.drop (0) ([1, 2, 3, 4, 5])) (S.Just ([1, 2, 3, 4, 5]));
   eq (S.drop (1) ([1, 2, 3, 4, 5])) (S.Just ([2, 3, 4, 5]));
@@ -21,4 +22,12 @@ test ('drop', () => {
 
   eq (S.drop (-1) ([1, 2, 3, 4, 5])) (S.Nothing);
 
+  eq (S.drop (0) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Cons (1) (Cons (2) (Cons (3) (Nil)))));
+  eq (S.drop (1) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Cons (2) (Cons (3) (Nil))));
+  eq (S.drop (2) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Cons (3) (Nil)));
+  eq (S.drop (3) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Nil));
+  eq (S.drop (4) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Nothing);
+
+  eq (S.drop (-1) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Nothing);
+
 });
diff --git a/test/take.js b/test/take.js
index ebdcb27..d52a667 100644
--- a/test/take.js
+++ b/test/take.js
@@ -1,7 +1,8 @@
 'use strict';
 
-const S = require ('..');
+const S = require ('./internal/sanctuary');
 
+const {Nil, Cons} = require ('./internal/List');
 const eq = require ('./internal/eq');
 
 
@@ -9,7 +10,7 @@ test ('take', () => {
 
   eq (typeof S.take) ('function');
   eq (S.take.length) (1);
-  eq (S.show (S.take)) ('take :: Integer -> Array a -> Maybe (Array a)');
+  eq (S.show (S.take)) ('take :: (Applicative f, Foldable f, Monoid f) => Integer -> f a -> Maybe (f a)');
 
   eq (S.take (0) ([1, 2, 3, 4, 5])) (S.Just ([]));
   eq (S.take (1) ([1, 2, 3, 4, 5])) (S.Just ([1]));
@@ -21,4 +22,12 @@ test ('take', () => {
 
   eq (S.take (-1) ([1, 2, 3, 4, 5])) (S.Nothing);
 
+  eq (S.take (0) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Nil));
+  eq (S.take (1) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Cons (1) (Nil)));
+  eq (S.take (2) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Cons (1) (Cons (2) (Nil))));
+  eq (S.take (3) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Just (Cons (1) (Cons (2) (Cons (3) (Nil)))));
+  eq (S.take (4) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Nothing);
+
+  eq (S.take (-1) (Cons (1) (Cons (2) (Cons (3) (Nil))))) (S.Nothing);
+
 });