Repository: fxamacker/cbor
Pull request #266 (2021-04-15 03:51:21)

Add stream encoding for CBOR array and tag
StreamEncoder provides low-level API for sequential encoding.
    
It can be used to avoid creating large collection objects.
    
Users of StreamEncoder should be familiar with CBOR data models.  When in doubt, please use Encoder instead.
    
Users should use Valid() to check whether generated CBOR data is complete and well-formed.


Patch:
diff --git a/encode.go b/encode.go
index 95d2c23..ebd903a 100644
--- a/encode.go
+++ b/encode.go
@@ -488,6 +488,7 @@ func (opts EncOptions) encMode() (*encMode, error) {
 type EncMode interface {
 	Marshal(v interface{}) ([]byte, error)
 	NewEncoder(w io.Writer) *Encoder
+	NewStreamEncoder(w io.Writer) *StreamEncoder
 	EncOptions() EncOptions
 }
 
@@ -553,6 +554,13 @@ func (em *encMode) NewEncoder(w io.Writer) *Encoder {
 	return &Encoder{w: w, em: em, e: getEncoderBuffer()}
 }
 
+// NewStreamEncoder returns a new stream encoder that writes sequentially to w using em EncMode.
+func (em *encMode) NewStreamEncoder(w io.Writer) *StreamEncoder {
+	return &StreamEncoder{
+		Encoder: &Encoder{w: w, em: em, e: getEncoderBuffer(), stream: true},
+	}
+}
+
 type encoderBuffer struct {
 	bytes.Buffer
 	scratch [16]byte
@@ -562,7 +570,7 @@ type encoderBuffer struct {
 var encoderBufferPool = sync.Pool{
 	New: func() interface{} {
 		e := new(encoderBuffer)
-		e.Grow(32) // TODO: make this configurable
+		e.Grow(64) // TODO: make this configurable
 		return e
 	},
 }
diff --git a/stream.go b/stream.go
index 29172b9..8c7bdcc 100644
--- a/stream.go
+++ b/stream.go
@@ -86,6 +86,7 @@ type Encoder struct {
 	em         *encMode
 	e          *encoderBuffer
 	indefTypes []cborType
+	stream     bool
 }
 
 // NewEncoder returns a new encoder that writes to w using the default encoding options.
@@ -112,10 +113,10 @@ func (enc *Encoder) Encode(v interface{}) error {
 	}
 
 	err := encode(enc.e, enc.em, reflect.ValueOf(v))
-	if err == nil {
+	if err == nil && !enc.stream {
 		_, err = enc.e.WriteTo(enc.w)
+		enc.e.Reset()
 	}
-	enc.e.Reset()
 	return err
 }
 
diff --git a/stream_encode.go b/stream_encode.go
new file mode 100644
index 0000000..f8d9c73
--- /dev/null
+++ b/stream_encode.go
@@ -0,0 +1,40 @@
+// Copyright (c) Faye Amacker. All rights reserved.
+// Licensed under the MIT License. See LICENSE in the project root for license information.
+
+package cbor
+
+import "io"
+
+// StreamEncoder provides low-level API for sequential encoding.
+//
+// Users of StreamEncoder should be familiar with CBOR data models.
+// When in doubt, please use Encoder instead.
+//
+// IMPORTANT: Use Valid() to check whether generated CBOR data is
+// complete and well-formed.
+type StreamEncoder struct {
+	*Encoder
+}
+
+// NewStreamEncoder returns a new StreamEncoder for sequential encoding.
+func NewStreamEncoder(w io.Writer) *StreamEncoder {
+	enc := defaultEncMode.NewEncoder(w)
+	enc.stream = true
+	return &StreamEncoder{enc}
+}
+
+// Flush writes streamed data to underlying io.Writer.
+func (se *StreamEncoder) Flush() error {
+	_, err := se.Encoder.e.WriteTo(se.Encoder.w)
+	return err
+}
+
+// EncodeArrayHead encodes CBOR array head of specified size.
+func (se *StreamEncoder) EncodeArrayHead(size uint64) {
+	encodeHead(se.Encoder.e, byte(cborTypeArray), size)
+}
+
+// EncodeTagHead encodes CBOR tag head with num as tag number.
+func (se *StreamEncoder) EncodeTagHead(num uint64) {
+	encodeHead(se.Encoder.e, byte(cborTypeTag), num)
+}


Test patch:
diff --git a/stream_encode_test.go b/stream_encode_test.go
new file mode 100644
index 0000000..60c1ab2
--- /dev/null
+++ b/stream_encode_test.go
@@ -0,0 +1,196 @@
+// Copyright (c) Faye Amacker. All rights reserved.
+// Licensed under the MIT License. See LICENSE in the project root for license information.
+
+package cbor
+
+import (
+	"bytes"
+	"math/big"
+	"testing"
+)
+
+func TestStreamEncodeArray(t *testing.T) {
+	t.Run("default mode", func(t *testing.T) {
+		expected := []byte{
+			// array, 2 items follow
+			0x82,
+			// UTF-8 string, length 5
+			0x65,
+			// h, e, l, l, o
+			0x68, 0x65, 0x6c, 0x6c, 0x6f,
+			// array, 1 items follow
+			0x81,
+			// 1
+			0x01,
+		}
+
+		var buf bytes.Buffer
+		se := NewStreamEncoder(&buf)
+
+		se.EncodeArrayHead(2)
+
+		err := se.Encode("hello")
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		se.EncodeArrayHead(1)
+
+		err = se.Encode(big.NewInt(1))
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		err = se.Flush()
+		if err != nil {
+			t.Errorf("StreamEncoder.Flush() returned error %v", err)
+		}
+
+		if !bytes.Equal(buf.Bytes(), expected) {
+			t.Errorf("StreamEncoder encoded 0x%x, want 0x%x", buf.Bytes(), expected)
+		}
+	})
+
+	t.Run("BigIntConvertNone mode", func(t *testing.T) {
+		expected := []byte{
+			// array, 2 items follow
+			0x82,
+			// UTF-8 string, length 5
+			0x65,
+			// h, e, l, l, o
+			0x68, 0x65, 0x6c, 0x6c, 0x6f,
+			// array, 1 items follow
+			0x81,
+			// positive bignum
+			0xc2,
+			// byte string, length 1
+			0x41,
+			0x01,
+		}
+
+		opts := EncOptions{BigIntConvert: BigIntConvertNone}
+		em, err := opts.encMode()
+		if err != nil {
+			panic(err)
+		}
+
+		var buf bytes.Buffer
+		se := em.NewStreamEncoder(&buf)
+
+		se.EncodeArrayHead(2)
+
+		err = se.Encode("hello")
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		se.EncodeArrayHead(1)
+
+		err = se.Encode(big.NewInt(1))
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		err = se.Flush()
+		if err != nil {
+			t.Errorf("StreamEncoder.Flush() returned error %v", err)
+		}
+
+		if !bytes.Equal(buf.Bytes(), expected) {
+			t.Errorf("StreamEncoder encoded 0x%x, want 0x%x", buf.Bytes(), expected)
+		}
+	})
+}
+
+func TestStreamEncodeTag(t *testing.T) {
+	t.Run("default mode", func(t *testing.T) {
+		expected := []byte{
+			// tag 128
+			0xd8, 0x80,
+			// array, 2 items follow
+			0x82,
+			// UTF-8 string, length 5
+			0x65,
+			// h, e, l, l, o
+			0x68, 0x65, 0x6c, 0x6c, 0x6f,
+			// 1
+			0x01,
+		}
+
+		var buf bytes.Buffer
+		se := NewStreamEncoder(&buf)
+
+		se.EncodeTagHead(128)
+
+		se.EncodeArrayHead(2)
+
+		err := se.Encode("hello")
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		err = se.Encode(big.NewInt(1))
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		err = se.Flush()
+		if err != nil {
+			t.Errorf("StreamEncoder.Flush() returned error %v", err)
+		}
+
+		if !bytes.Equal(buf.Bytes(), expected) {
+			t.Errorf("StreamEncoder encoded 0x%x, want 0x%x", buf.Bytes(), expected)
+		}
+	})
+
+	t.Run("BigIntConvertNone", func(t *testing.T) {
+		expected := []byte{
+			// tag 128
+			0xd8, 0x80,
+			// array, 2 items follow
+			0x82,
+			// UTF-8 string, length 5
+			0x65,
+			// h, e, l, l, o
+			0x68, 0x65, 0x6c, 0x6c, 0x6f,
+			// positive bignum
+			0xc2,
+			// byte string, length 1
+			0x41,
+			0x01,
+		}
+
+		opts := EncOptions{BigIntConvert: BigIntConvertNone}
+		em, err := opts.encMode()
+		if err != nil {
+			panic(err)
+		}
+
+		var buf bytes.Buffer
+		se := em.NewStreamEncoder(&buf)
+
+		se.EncodeTagHead(128)
+
+		se.EncodeArrayHead(2)
+
+		err = se.Encode("hello")
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		err = se.Encode(big.NewInt(1))
+		if err != nil {
+			t.Errorf("StreamEncoder.Encode() returned error %v", err)
+		}
+
+		err = se.Flush()
+		if err != nil {
+			t.Errorf("StreamEncoder.Flush() returned error %v", err)
+		}
+
+		if !bytes.Equal(buf.Bytes(), expected) {
+			t.Errorf("StreamEncoder encoded 0x%x, want 0x%x", buf.Bytes(), expected)
+		}
+	})
+}