feat: add Capture for readers

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

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)))
 }