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 cbf4b42457617c605102e0afb78580a17a1e064c..aec14a8dc6ede9a23543171d5fb445a393e7c749 100644 (file)
--- 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 85bf1327191b2d026e55a298587944d21d0eef00..72ee2f6aa92d2e5fd21467e379344f89e27ed08c 100644 (file)
--- 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