From 17fc788c854ad86c7f223fd4769251266e32b295 Mon Sep 17 00:00:00 2001
From: tdakkota <tanc13@yandex.ru>
Date: Tue, 23 Nov 2021 12:44:00 +0300
Subject: [PATCH] feat: add Capture for readers

---
 dec_capture.go      | 18 ++++++++++++------
 dec_capture_test.go | 42 ++++++++++++++++++++++--------------------
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/dec_capture.go b/dec_capture.go
index 6a5beab..5a9e643 100644
--- a/dec_capture.go
+++ b/dec_capture.go
@@ -1,19 +1,25 @@
 package jx
 
 import (
-	"github.com/go-faster/errors"
+	"bytes"
+	"io"
 )
 
 // Capture calls f and then rolls back to state before call.
-//
-// Does not work with reader.
 func (d *Decoder) Capture(f func(d *Decoder) error) error {
-	if d.reader != nil {
-		return errors.New("capture is not supported with reader")
-	}
 	if f == nil {
 		return nil
 	}
+
+	if d.reader != nil {
+		// TODO(tdakkota): May it be more efficient?
+		var buf bytes.Buffer
+		reader := io.TeeReader(d.reader, &buf)
+		defer func() {
+			d.reader = io.MultiReader(&buf, d.reader)
+		}()
+		d.reader = reader
+	}
 	head, tail, depth := d.head, d.tail, d.depth
 	err := f(d)
 	d.head, d.tail, d.depth = head, tail, depth
diff --git a/dec_capture_test.go b/dec_capture_test.go
index bf6a566..69eab1e 100644
--- a/dec_capture_test.go
+++ b/dec_capture_test.go
@@ -1,7 +1,7 @@
 package jx
 
 import (
-	"bytes"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -84,25 +84,27 @@ func BenchmarkIterator_Skip(b *testing.B) {
 }
 
 func TestDecoder_Capture(t *testing.T) {
-	i := DecodeStr(`["foo", "bar", "baz"]`)
-	var elems int
-	if err := i.Capture(func(i *Decoder) error {
-		return i.Arr(func(i *Decoder) error {
-			elems++
-			return i.Skip()
-		})
-	}); err != nil {
-		t.Fatal(err)
+	test := func(i *Decoder) func(t *testing.T) {
+		return func(t *testing.T) {
+			var elems int
+			if err := i.Capture(func(i *Decoder) error {
+				return i.Arr(func(i *Decoder) error {
+					elems++
+					return i.Skip()
+				})
+			}); err != nil {
+				t.Fatal(err)
+			}
+			require.Equal(t, Array, i.Next())
+			require.Equal(t, 3, elems)
+			t.Run("Nil", func(t *testing.T) {
+				require.NoError(t, i.Capture(nil))
+				require.Equal(t, Array, i.Next())
+			})
+		}
 	}
-	require.Equal(t, Array, i.Next())
-	require.Equal(t, 3, elems)
-	t.Run("Nil", func(t *testing.T) {
-		require.NoError(t, i.Capture(nil))
-		require.Equal(t, Array, i.Next())
-	})
-}
 
-func TestDecoder_Capture_reader(t *testing.T) {
-	i := Decode(new(bytes.Buffer), 0)
-	require.Error(t, i.Capture(nil))
+	testData := `["foo", "bar", "baz"]`
+	t.Run("Str", test(DecodeStr(testData)))
+	t.Run("Reader", test(Decode(strings.NewReader(testData), 0)))
 }