Improve wordwrap performance

ansi/buffer.go

@@ -2,6 +2,7 @@ package ansi
 
 import (
 	"bytes"
+	"unsafe"
 
 	"github.com/mattn/go-runewidth"
 )
@@ -13,7 +14,7 @@ type Buffer struct {
 
 // PrintableRuneWidth returns the width of all printable runes in the buffer.
 func (w Buffer) PrintableRuneWidth() int {
-	return PrintableRuneWidth(w.String())
+	return PrintableRuneWidth(b2s(w.Bytes()))
 }
 
 func PrintableRuneWidth(s string) int {
@@ -36,3 +37,13 @@ func PrintableRuneWidth(s string) int {
 
 	return n
 }
+
+// b2s converts byte slice to a string without memory allocation.
+// See https://groups.google.com/forum/#!msg/Golang-Nuts/ENgbUzYvCuU/90yGx7GUAgAJ .
+//
+// Note it may break if string and/or slice header will change
+// in the future go versions.
+func b2s(b []byte) string {
+	/* #nosec G103 */
+	return *(*string)(unsafe.Pointer(&b))
+}

ansi/buffer_test.go

@@ -17,6 +17,19 @@ func TestBuffer_PrintableRuneWidth(t *testing.T) {
 	}
 }
 
+// go test -bench=BenchmarkBuffer_PrintableRuneWidth -benchmem -count=4
+func BenchmarkBuffer_PrintableRuneWidth(b *testing.B) {
+	var buf Buffer
+	buf.Buffer.WriteString("\x1B[38;2;249;38;114m(\x1B[0m\x1B[38;2;248;248;242mjust another test\x1B[38;2;249;38;114m)\x1B[0m")
+	b.RunParallel(func(pb *testing.PB) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		for pb.Next() {
+			buf.PrintableRuneWidth()
+		}
+	})
+}
+
 // go test -bench=Benchmark_PrintableRuneWidth -benchmem -count=4
 func Benchmark_PrintableRuneWidth(b *testing.B) {
 	s := "\x1B[38;2;249;38;114mfoo"

wordwrap/wordwrap.go

@@ -2,8 +2,11 @@ package wordwrap
 
 import (
 	"bytes"
+	"reflect"
 	"strings"
+	"sync"
 	"unicode"
+	"unsafe"
 
 	"github.com/muesli/reflow/ansi"
 )
@@ -44,36 +47,38 @@ func NewWriter(limit int) *WordWrap {
 // Bytes is shorthand for declaring a new default WordWrap instance,
 // used to immediately word-wrap a byte slice.
 func Bytes(b []byte, limit int) []byte {
-	f := NewWriter(limit)
+	f := acquireWordWrap(limit)
+	defer wp.Put(f)
 	_, _ = f.Write(b)
-	f.Close()
 
-	return f.Bytes()
+	f.addWord()
+
+	return f.buf.Bytes()
 }
 
 // String is shorthand for declaring a new default WordWrap instance,
 // used to immediately word-wrap a string.
 func String(s string, limit int) string {
-	return string(Bytes([]byte(s), limit))
+	return b2s(Bytes(s2b(s), limit))
 }
 
 func (w *WordWrap) addSpace() {
 	w.lineLen += w.space.Len()
-	w.buf.Write(w.space.Bytes())
+	_, _ = w.buf.Write(w.space.Bytes())
 	w.space.Reset()
 }
 
 func (w *WordWrap) addWord() {
 	if w.word.Len() > 0 {
 		w.addSpace()
 		w.lineLen += w.word.PrintableRuneWidth()
-		w.buf.Write(w.word.Bytes())
+		_, _ = w.buf.Write(w.word.Bytes())
 		w.word.Reset()
 	}
 }
 
 func (w *WordWrap) addNewLine() {
-	w.buf.WriteRune('\n')
+	_ = w.buf.WriteByte('\n')
 	w.lineLen = 0
 	w.space.Reset()
 }
@@ -93,18 +98,18 @@ func (w *WordWrap) Write(b []byte) (int, error) {
 		return w.buf.Write(b)
 	}
 
-	s := string(b)
+	s := b2s(b)
 	if !w.KeepNewlines {
 		s = strings.Replace(strings.TrimSpace(s), "\n", " ", -1)
 	}
 
 	for _, c := range s {
 		if c == '\x1B' {
 			// ANSI escape sequence
-			w.word.WriteRune(c)
+			_, _ = w.word.WriteRune(c)
 			w.ansi = true
 		} else if w.ansi {
-			w.word.WriteRune(c)
+			_, _ = w.word.WriteRune(c)
 			if (c >= 0x40 && c <= 0x5a) || (c >= 0x61 && c <= 0x7a) {
 				// ANSI sequence terminated
 				w.ansi = false
@@ -117,7 +122,7 @@ func (w *WordWrap) Write(b []byte) (int, error) {
 					w.lineLen = 0
 				} else {
 					// preserve whitespace
-					w.buf.Write(w.space.Bytes())
+					_, _ = w.buf.Write(w.space.Bytes())
 				}
 				w.space.Reset()
 			}
@@ -127,20 +132,20 @@ func (w *WordWrap) Write(b []byte) (int, error) {
 		} else if unicode.IsSpace(c) {
 			// end of current word
 			w.addWord()
-			w.space.WriteRune(c)
+			_, _ = w.space.WriteRune(c)
 		} else if inGroup(w.Breakpoints, c) {
 			// valid breakpoint
 			w.addSpace()
 			w.addWord()
-			w.buf.WriteRune(c)
+			_, _ = w.buf.WriteRune(c)
 		} else {
 			// any other character
-			w.word.WriteRune(c)
+			_, _ = w.word.WriteRune(c)
 
 			// add a line break if the current word would exceed the line's
 			// character limit
-			if w.lineLen+w.space.Len()+w.word.PrintableRuneWidth() > w.Limit &&
-				w.word.PrintableRuneWidth() < w.Limit {
+			if wordWidth := w.word.PrintableRuneWidth(); w.lineLen+w.space.Len()+wordWidth > w.Limit &&
+				wordWidth < w.Limit {
 				w.addNewLine()
 			}
 		}
@@ -165,3 +170,48 @@ func (w *WordWrap) Bytes() []byte {
 func (w *WordWrap) String() string {
 	return w.buf.String()
 }
+
+var wp = sync.Pool{
+	New: func() interface{} {
+		return &WordWrap{
+			Breakpoints:  defaultBreakpoints,
+			Newline:      defaultNewline,
+			KeepNewlines: true,
+		}
+	},
+}
+
+func acquireWordWrap(limit int) *WordWrap {
+	w := wp.Get().(*WordWrap)
+	w.Limit = limit
+	w.buf.Reset()
+	w.lineLen = 0
+	w.ansi = false
+
+	return w
+}
+
+// b2s converts byte slice to a string without memory allocation.
+// See https://groups.google.com/forum/#!msg/Golang-Nuts/ENgbUzYvCuU/90yGx7GUAgAJ .
+//
+// Note it may break if string and/or slice header will change
+// in the future go versions.
+func b2s(b []byte) string {
+	/* #nosec G103 */
+	return *(*string)(unsafe.Pointer(&b))
+}
+
+// s2b converts string to a byte slice without memory allocation.
+//
+// Note it may break if string and/or slice header will change
+// in the future go versions.
+func s2b(s string) (b []byte) {
+	/* #nosec G103 */
+	bh := (*reflect.SliceHeader)(unsafe.Pointer(&b))
+	/* #nosec G103 */
+	sh := *(*reflect.StringHeader)(unsafe.Pointer(&s))
+	bh.Data = sh.Data
+	bh.Len = sh.Len
+	bh.Cap = sh.Len
+	return b
+}

wordwrap/wordwrap_test.go

@@ -153,3 +153,15 @@ func TestWordWrapString(t *testing.T) {
 		t.Errorf("expected:\n\n`%s`\n\nActual Output:\n\n`%s`", expected, actual)
 	}
 }
+
+// go test -bench=BenchmarkWordWrapString -benchmem -count=4
+func BenchmarkWordWrapString(b *testing.B) {
+	s := "\x1B[38;2;249;38;114m(\x1B[0m\x1B[38;2;248;248;242mjust 另一个 test\x1B[38;2;249;38;114m)\x1B[0m"
+	b.RunParallel(func(pb *testing.PB) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		for pb.Next() {
+			String(s, 7)
+		}
+	})
+}