Improve padding performance

padding/padding.go

@@ -3,7 +3,8 @@ package padding
 import (
 	"bytes"
 	"io"
-	"strings"
+	"sync"
+	"unicode/utf8"
 
 	"github.com/mattn/go-runewidth"
 	"github.com/muesli/reflow/ansi"
@@ -46,11 +47,16 @@ func NewWriterPipe(forward io.Writer, width uint, paddingFunc PaddingFunc) *Writ
 // Bytes is shorthand for declaring a new default padding-writer instance,
 // used to immediately pad a byte slice.
 func Bytes(b []byte, width uint) []byte {
-	f := NewWriter(width, nil)
+	f := acquireWriter(width)
+	defer wp.Put(f)
+
 	_, _ = f.Write(b)
-	_ = f.Flush()
 
-	return f.Bytes()
+	if f.lineLen != 0 {
+		_ = f.pad()
+	}
+
+	return f.buf.Bytes()
 }
 
 // String is shorthand for declaring a new default padding-writer instance,
@@ -71,7 +77,7 @@ func (w *Writer) Write(b []byte) (int, error) {
 				w.ansi = false
 			}
 		} else {
-			w.lineLen += runewidth.StringWidth(string(c))
+			w.lineLen += runewidth.RuneWidth(c)
 
 			if c == '\n' {
 				// end of current line
@@ -84,7 +90,7 @@ func (w *Writer) Write(b []byte) (int, error) {
 			}
 		}
 
-		_, err := w.ansiWriter.Write([]byte(string(c)))
+		_, err := w.writeRune(c)
 		if err != nil {
 			return 0, err
 		}
@@ -100,7 +106,7 @@ func (w *Writer) pad() error {
 				w.PadFunc(w.ansiWriter)
 			}
 		} else {
-			_, err := w.ansiWriter.Write([]byte(strings.Repeat(" ", int(w.Padding)-w.lineLen)))
+			_, err := w.ansiWriter.Write(bytes.Repeat([]byte(" "), int(w.Padding)-w.lineLen))
 			if err != nil {
 				return err
 			}
@@ -110,6 +116,12 @@ func (w *Writer) pad() error {
 	return nil
 }
 
+func (w *Writer) writeRune(r rune) (int, error) {
+	bb := make([]byte, utf8.UTFMax)
+	n := utf8.EncodeRune(bb, r)
+	return w.ansiWriter.Write(bb[:n])
+}
+
 // Close will finish the padding operation.
 func (w *Writer) Close() (err error) {
 	return w.Flush()
@@ -141,3 +153,23 @@ func (w *Writer) Flush() (err error) {
 
 	return
 }
+
+var wp = sync.Pool{
+	New: func() interface{} {
+		w := &Writer{}
+		w.ansiWriter = &ansi.Writer{
+			Forward: &w.buf,
+		}
+		return w
+	},
+}
+
+func acquireWriter(width uint) *Writer {
+	w := wp.Get().(*Writer)
+	w.Padding = width
+	w.lineLen = 0
+	w.ansi = false
+	w.buf.Reset()
+
+	return w
+}

padding/padding_test.go

@@ -100,6 +100,7 @@ func TestPaddingString(t *testing.T) {
 	}
 }
 
+// go test -bench=BenchmarkPaddingString -benchmem -count=4
 func BenchmarkPaddingString(b *testing.B) {
 	b.RunParallel(func(pb *testing.PB) {
 		b.ReportAllocs()