avoid using internal package in ext encoders

shamaton · Sep 2, 2024 · 208705f · 208705f
1 parent 9d4237b
commit 208705f
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 18 deletions.
diff --git a/ext/encode_stream.go b/ext/encode_stream.go
@@ -1,15 +1,17 @@
 package ext
 
 import (
-	"github.com/shamaton/msgpack/v2/internal/common"
 	"io"
 	"reflect"
+
+	"github.com/shamaton/msgpack/v2/internal/common"
 )
 
+// StreamEncoder is interface that extended encoder should implement
 type StreamEncoder interface {
 	Code() int8
 	Type() reflect.Type
-	Write(w io.Writer, value reflect.Value, buf *common.Buffer) error
+	Write(w StreamWriter, value reflect.Value) error
 }
 
 type StreamEncoderCommon struct{}
@@ -118,3 +120,118 @@ func (c *StreamEncoderCommon) WriteByte4Uint32(w io.Writer, value uint32, buf *c
 func (c *StreamEncoderCommon) WriteBytes(w io.Writer, bs []byte, buf *common.Buffer) error {
 	return buf.Write(w, bs...)
 }
+
+// StreamWriter is provided some writing functions for extended format by user
+type StreamWriter struct {
+	w   io.Writer
+	buf *common.Buffer
+}
+
+func CreateStreamWriter(w io.Writer, buf *common.Buffer) StreamWriter {
+	return StreamWriter{w, buf}
+}
+
+func (w *StreamWriter) WriteByte1Int64(value int64) error {
+	return w.buf.Write(w.w,
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte2Int64(value int64) error {
+	return w.buf.Write(w.w,
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte4Int64(value int64) error {
+	return w.buf.Write(w.w,
+		byte(value>>24),
+		byte(value>>16),
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte8Int64(value int64) error {
+	return w.buf.Write(w.w,
+		byte(value>>56),
+		byte(value>>48),
+		byte(value>>40),
+		byte(value>>32),
+		byte(value>>24),
+		byte(value>>16),
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte1Uint64(value uint64) error {
+	return w.buf.Write(w.w,
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte2Uint64(value uint64) error {
+	return w.buf.Write(w.w,
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte4Uint64(value uint64) error {
+	return w.buf.Write(w.w,
+		byte(value>>24),
+		byte(value>>16),
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte8Uint64(value uint64) error {
+	return w.buf.Write(w.w,
+		byte(value>>56),
+		byte(value>>48),
+		byte(value>>40),
+		byte(value>>32),
+		byte(value>>24),
+		byte(value>>16),
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte1Int(value int) error {
+	return w.buf.Write(w.w,
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte2Int(value int) error {
+	return w.buf.Write(w.w,
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte4Int(value int) error {
+	return w.buf.Write(w.w,
+		byte(value>>24),
+		byte(value>>16),
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteByte4Uint32(value uint32) error {
+	return w.buf.Write(w.w,
+		byte(value>>24),
+		byte(value>>16),
+		byte(value>>8),
+		byte(value),
+	)
+}
+
+func (w *StreamWriter) WriteBytes(bs []byte) error {
+	return w.buf.Write(w.w, bs...)
+}
diff --git a/internal/stream/encoding/struct.go b/internal/stream/encoding/struct.go
@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/shamaton/msgpack/v2/def"
+	"github.com/shamaton/msgpack/v2/ext"
 	"github.com/shamaton/msgpack/v2/internal/common"
 )
 
@@ -24,7 +25,8 @@ func (e *encoder) getStructWriter(typ reflect.Type) structWriteFunc {
 	for i := range extCoders {
 		if extCoders[i].Type() == typ {
 			return func(rv reflect.Value) error {
-				return extCoders[i].Write(e.w, rv, e.buf)
+				w := ext.CreateStreamWriter(e.w, e.buf)
+				return extCoders[i].Write(w, rv)
 			}
 		}
 	}
@@ -39,7 +41,8 @@ func (e *encoder) writeStruct(rv reflect.Value) error {
 
 	for i := range extCoders {
 		if extCoders[i].Type() == rv.Type() {
-			return extCoders[i].Write(e.w, rv, e.buf)
+			w := ext.CreateStreamWriter(e.w, e.buf)
+			return extCoders[i].Write(w, rv)
 		}
 	}
 

diff --git a/time/encode_stream.go b/time/encode_stream.go
@@ -1,8 +1,6 @@
 package time
 
 import (
-	"github.com/shamaton/msgpack/v2/internal/common"
-	"io"
 	"reflect"
 	"time"
 
@@ -26,50 +24,50 @@ func (timeStreamEncoder) Type() reflect.Type {
 	return typeOf
 }
 
-func (e timeStreamEncoder) Write(w io.Writer, value reflect.Value, buf *common.Buffer) error {
+func (e timeStreamEncoder) Write(w ext.StreamWriter, value reflect.Value) error {
 	t := value.Interface().(time.Time)
 
 	secs := uint64(t.Unix())
 	if secs>>34 == 0 {
 		data := uint64(t.Nanosecond())<<34 | secs
 		if data&0xffffffff00000000 == 0 {
-			if err := e.WriteByte1Int(w, def.Fixext4, buf); err != nil {
+			if err := w.WriteByte1Int(def.Fixext4); err != nil {
 				return err
 			}
-			if err := e.WriteByte1Int(w, def.TimeStamp, buf); err != nil {
+			if err := w.WriteByte1Int(def.TimeStamp); err != nil {
 				return err
 			}
-			if err := e.WriteByte4Uint64(w, data, buf); err != nil {
+			if err := w.WriteByte4Uint64(data); err != nil {
 				return err
 			}
 			return nil
 		}
 
-		if err := e.WriteByte1Int(w, def.Fixext8, buf); err != nil {
+		if err := w.WriteByte1Int(def.Fixext8); err != nil {
 			return err
 		}
-		if err := e.WriteByte1Int(w, def.TimeStamp, buf); err != nil {
+		if err := w.WriteByte1Int(def.TimeStamp); err != nil {
 			return err
 		}
-		if err := e.WriteByte8Uint64(w, data, buf); err != nil {
+		if err := w.WriteByte8Uint64(data); err != nil {
 			return err
 		}
 		return nil
 	}
 
-	if err := e.WriteByte1Int(w, def.Ext8, buf); err != nil {
+	if err := w.WriteByte1Int(def.Ext8); err != nil {
 		return err
 	}
-	if err := e.WriteByte1Int(w, 12, buf); err != nil {
+	if err := w.WriteByte1Int(12); err != nil {
 		return err
 	}
-	if err := e.WriteByte1Int(w, def.TimeStamp, buf); err != nil {
+	if err := w.WriteByte1Int(def.TimeStamp); err != nil {
 		return err
 	}
-	if err := e.WriteByte4Int(w, t.Nanosecond(), buf); err != nil {
+	if err := w.WriteByte4Int(t.Nanosecond()); err != nil {
 		return err
 	}
-	if err := e.WriteByte8Uint64(w, secs, buf); err != nil {
+	if err := w.WriteByte8Uint64(secs); err != nil {
 		return err
 	}
 	return nil