From 1b2566f8228214176d6beb64293decfa5269607e Mon Sep 17 00:00:00 2001
From: ivahaev <ivahaev@gmail.com>
Date: Thu, 11 Aug 2016 01:38:01 +0500
Subject: [PATCH] Implementing ArrayEachLimited method. If called callback
 returns false, ArrayEachLimited will stop processing.

---
 parser.go      | 15 +++++++++++-
 parser_test.go | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/parser.go b/parser.go
index a86cdd5..57fbe45 100644
--- a/parser.go
+++ b/parser.go
@@ -456,6 +456,15 @@ func Get(data []byte, keys ...string) (value []byte, dataType ValueType, offset
 
 // ArrayEach is used when iterating arrays, accepts a callback function with the same return arguments as `Get`.
 func ArrayEach(data []byte, cb func(value []byte, dataType ValueType, offset int, err error), keys ...string) (err error) {
+	return ArrayEachLimited(data, func(_ int, value []byte, dataType ValueType, offset int, err error) bool {
+		cb(value, dataType, offset, err)
+		return true
+	}, keys...)
+}
+
+// ArrayEachLimited is used when iterating arrays, accepts a callback function with the same return arguments as `Get` with value index.
+// It allows to break iterating when necessary value found by returning false from callback func.
+func ArrayEachLimited(data []byte, cb func(idx int, value []byte, dataType ValueType, offset int, err error) bool, keys ...string) (err error) {
 	if len(data) == 0 {
 		return MalformedObjectError
 	}
@@ -482,6 +491,7 @@ func ArrayEach(data []byte, cb func(value []byte, dataType ValueType, offset int
 		offset++
 	}
 
+	var idx int
 	for true {
 		v, t, o, e := Get(data[offset:])
 
@@ -490,7 +500,10 @@ func ArrayEach(data []byte, cb func(value []byte, dataType ValueType, offset int
 		}
 
 		if t != NotExist {
-			cb(v, t, o, e)
+			if !cb(idx, v, t, o, e) {
+				break
+			}
+			idx++
 		}
 
 		if e != nil {
diff --git a/parser_test.go b/parser_test.go
index 74d4f33..3b69dfb 100644
--- a/parser_test.go
+++ b/parser_test.go
@@ -650,6 +650,71 @@ func TestArrayEach(t *testing.T) {
 	}, "a", "b")
 }
 
+func TestArrayEachLimited(t *testing.T) {
+	mock := []byte(`{"a": { "b":[{"x": 1} ,{"x":2},{ "x":3}, {"x":4} ]}}`)
+	count := 0
+
+	ArrayEachLimited(mock, func(idx int, value []byte, dataType ValueType, offset int, err error) bool {
+		count++
+
+		if count != idx+1 {
+			t.Errorf("Wrong index: %v", idx)
+		}
+
+		switch count {
+		case 1:
+			if string(value) != `{"x": 1}` {
+				t.Errorf("Wrong first item: %s", string(value))
+			}
+		case 2:
+			if string(value) != `{"x":2}` {
+				t.Errorf("Wrong second item: %s", string(value))
+			}
+		case 3:
+			if string(value) != `{ "x":3}` {
+				t.Errorf("Wrong third item: %s", string(value))
+			}
+		case 4:
+			if string(value) != `{"x":4}` {
+				t.Errorf("Wrong forth item: %s", string(value))
+			}
+		default:
+			t.Errorf("Should process only 4 items")
+		}
+
+		return true
+	}, "a", "b")
+
+	count = 0
+	ArrayEachLimited(mock, func(idx int, value []byte, dataType ValueType, offset int, err error) bool {
+		count++
+
+		if count != idx+1 {
+			t.Errorf("Wrong index: %v", idx)
+		}
+
+		switch count {
+		case 1:
+			if string(value) != `{"x": 1}` {
+				t.Errorf("Wrong first item: %s", string(value))
+			}
+		case 2:
+			if string(value) != `{"x":2}` {
+				t.Errorf("Wrong second item: %s", string(value))
+			}
+			return false
+		case 3:
+			t.Errorf("Wrong behavior. Should process only two values from array")
+		case 4:
+			t.Errorf("Wrong behavior. Should process only two values from array")
+		default:
+			t.Errorf("Wrong behavior. Should process only two values from array")
+		}
+
+		return true
+	}, "a", "b")
+}
+
 var testJson = []byte(`{"name": "Name", "order": "Order", "sum": 100, "len": 12, "isPaid": true, "nested": {"a":"test", "b":2, "nested3":{"a":"test3","b":4}, "c": "unknown"}, "nested2": {"a":"test2", "b":3}, "arr": [{"a":"zxc", "b": 1}, {"a":"123", "b":2}], "arrInt": [1,2,3,4], "intPtr": 10}`)
 
 func TestEachKey(t *testing.T) {