From 81d73150002b0fa1c26cdd8b3e8729735c204612 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 8 Aug 2014 23:47:51 +0200
Subject: [PATCH] bugfix: take would fail on 0

---
 sequence.ml            | 2 +-
 tests/test_sequence.ml | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/sequence.ml b/sequence.ml
index d9bc87e..54ed7e8 100644
--- a/sequence.ml
+++ b/sequence.ml
@@ -336,9 +336,9 @@ let take n seq k =
   let count = ref 0 in
   try
     seq (fun x ->
+      if !count = n then raise ExitSequence;
       incr count;
       k x;
-      if !count = n then raise ExitSequence
     )
   with ExitSequence -> ()
 
diff --git a/tests/test_sequence.ml b/tests/test_sequence.ml
index bfd7dac..1b6c862 100644
--- a/tests/test_sequence.ml
+++ b/tests/test_sequence.ml
@@ -190,6 +190,13 @@ let test_int_range () =
   OUnit.assert_equal ~printer:pp_ilist [] S.(to_list (10 --^ 60));
   ()
 
+let test_take () =
+  let l = S.(to_list (take 0 (of_list [1]))) in
+  OUnit.assert_equal ~printer:pp_ilist [] l;
+  let l = S.(to_list (take 5 (of_list [1;2;3;4;5;6;7;8;9;10]))) in
+  OUnit.assert_equal ~printer:pp_ilist [1;2;3;4;5] l;
+  ()
+
 let suite =
   "test_sequence" >:::
     [ "test_empty" >:: test_empty;
@@ -217,4 +224,5 @@ let suite =
       "test_unfoldr" >:: test_unfoldr;
       "test_hashtbl" >:: test_hashtbl;
       "test_int_range" >:: test_int_range;
+      "test_take" >:: test_take;
     ]