Avoid doing type switching in inner functions

The CPU profiler tells me that the encoder is spending a lot of time on
type switches.

1 files changed, 43 insertions, 6 deletions

diff --git a/reencode.go b/reencode.go
index 876af62..393e8c6 100644
--- a/reencode.go
+++ b/reencode.go
@@ -71,7 +71,7 @@ type ReEncoderConfig struct {
 func NewReEncoder(out io.Writer, cfg ReEncoderConfig) *ReEncoder {
 	return &ReEncoder{
 		ReEncoderConfig: cfg,
-		out:             out,
+		out:             internal.NewAllWriter(out),
 	}
 }
 
@@ -85,9 +85,9 @@ func NewReEncoder(out io.Writer, cfg ReEncoderConfig) *ReEncoder {
 // The memory use of a ReEncoder is O( (CompactIfUnder+1)^2 + depth).
 type ReEncoder struct {
 	ReEncoderConfig
-	out io.Writer
+	out internal.AllWriter
 
-	// state: .Write's utf8-decoding buffer
+	// state: .Write's and .WriteString's utf8-decoding buffer
 	buf    [utf8.UTFMax]byte
 	bufLen int
 
@@ -119,6 +119,11 @@ type speculation struct {
 
 // public API //////////////////////////////////////////////////////////////////
 
+var (
+	_ internal.AllWriter = (*ReEncoder)(nil)
+	_ io.Closer          = (*ReEncoder)(nil)
+)
+
 // Write implements io.Writer; it does what you'd expect.
 //
 // It is worth noting that Write returns the number of bytes consumed
@@ -152,6 +157,38 @@ func (enc *ReEncoder) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+// WriteString implements io.StringWriter; it does what you'd expect,
+// but see the notes on the Write method.
+func (enc *ReEncoder) WriteString(p string) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	var n int
+	if enc.bufLen > 0 {
+		copy(enc.buf[enc.bufLen:], p)
+		c, size := utf8.DecodeRune(enc.buf[:])
+		n += size - enc.bufLen
+		enc.bufLen = 0
+		if _, err := enc.WriteRune(c); err != nil {
+			return 0, err
+		}
+	}
+	for utf8.FullRuneInString(p[n:]) {
+		c, size := utf8.DecodeRuneInString(p[n:])
+		if _, err := enc.WriteRune(c); err != nil {
+			return n, err
+		}
+		n += size
+	}
+	enc.bufLen = copy(enc.buf[:], p[n:])
+	return len(p), nil
+}
+
+// WriteByte implements io.ByteWriter; it does what you'd expect.
+func (enc *ReEncoder) WriteByte(b byte) error {
+	return internal.WriteByte(enc, b)
+}
+
 // Close implements io.Closer; it does what you'd expect, mostly.
 //
 // The *ReEncoder may continue to be written to with new JSON values
@@ -471,7 +508,7 @@ func (enc *ReEncoder) handleRuneMain(c rune, t internal.RuneType) error {
 }
 
 func (enc *ReEncoder) emitByte(c byte) error {
-	err := writeByte(enc.out, c)
+	err := enc.out.WriteByte(c)
 	if err == nil {
 		enc.written++
 	}
@@ -488,12 +525,12 @@ func (enc *ReEncoder) emitNlIndent() error {
 		return err
 	}
 	if enc.Prefix != "" {
-		if err := enc.emit(io.WriteString(enc.out, enc.Prefix)); err != nil {
+		if err := enc.emit(enc.out.WriteString(enc.Prefix)); err != nil {
 			return err
 		}
 	}
 	for i := 0; i < enc.handleRuneState.curIndent; i++ {
-		if err := enc.emit(io.WriteString(enc.out, enc.Indent)); err != nil {
+		if err := enc.emit(enc.out.WriteString(enc.Indent)); err != nil {
 			return err
 		}
 	}