From: John Benediktsson <mrjbq7@gmail.com>
Date: Fri, 6 Jul 2018 15:07:11 +0000 (-0700)
Subject: sequences.extras: fix bug in count-head, add some tests.
X-Git-Tag: 0.98~56
X-Git-Url: https://gitweb.factorcode.org/gitweb.cgi?p=factor.git;a=commitdiff_plain;h=6e89e4ecabe86b78a7ad479fcdb742a5cb676aca

sequences.extras: fix bug in count-head, add some tests.
---

diff --git a/extra/sequences/extras/extras-tests.factor b/extra/sequences/extras/extras-tests.factor
index cbf4b42457..aec14a8dc6 100644
--- a/extra/sequences/extras/extras-tests.factor
+++ b/extra/sequences/extras/extras-tests.factor
@@ -273,3 +273,11 @@ tools.test vectors vocabs ;
 { "a_b" } [ "ab" CHAR: _ interleaved ] unit-test
 { "a_b_c" } [ "abc" CHAR: _ interleaved ] unit-test
 { "a_b_c_d" } [ "abcd" CHAR: _ interleaved ] unit-test
+
+{ 0 } [ { 1 2 3 4 } [ 5 > ] count-head ] unit-test
+{ 2 } [ { 1 2 3 4 } [ 3 < ] count-head ] unit-test
+{ 4 } [ { 1 2 3 4 } [ 5 < ] count-head ] unit-test
+
+{ 0 } [ { 1 2 3 4 } [ 5 > ] count-tail ] unit-test
+{ 2 } [ { 1 2 3 4 } [ 2 > ] count-tail ] unit-test
+{ 4 } [ { 1 2 3 4 } [ 5 < ] count-tail ] unit-test
diff --git a/extra/sequences/extras/extras.factor b/extra/sequences/extras/extras.factor
index 85bf132719..72ee2f6aa9 100644
--- a/extra/sequences/extras/extras.factor
+++ b/extra/sequences/extras/extras.factor
@@ -629,7 +629,7 @@ PRIVATE>
     [ dup length ] unless* tail-slice ; inline
 
 : count-head ( seq quot -- n )
-    [ not ] compose find drop ; inline
+    [ not ] compose [ find drop ] 2keep drop length or ; inline
 
 : count-tail ( seq quot -- n )
     [ not ] compose [ find-last drop ] 2keep drop