From d7ea7371f297a2cc0ea8d00d909b94843f833b4c Mon Sep 17 00:00:00 2001 From: itchyny Date: Wed, 13 Jan 2021 02:12:59 +0900 Subject: [PATCH] Improve performance by reducing allocations --- prettyjson.go | 187 ++++++++++++++++++++++++++++++++------------------ 1 file changed, 120 insertions(+), 67 deletions(-) diff --git a/prettyjson.go b/prettyjson.go index 10c79e4..b682964 100644 --- a/prettyjson.go +++ b/prettyjson.go @@ -5,9 +5,10 @@ import ( "bytes" "encoding/json" "fmt" + "io" "sort" "strconv" - "strings" + _ "unsafe" // go:linkname "github.com/fatih/color" ) @@ -77,14 +78,23 @@ func (f *Formatter) Format(data []byte) ([]byte, error) { return nil, err } - return []byte(f.pretty(v, 1)), nil + buf := newBytesBuffer() + f.pretty(buf, v, 1) + return buf.Bytes(), nil } -func (f *Formatter) sprintfColor(c *color.Color, format string, args ...interface{}) string { +func (f *Formatter) setColor(buf *bytesBuffer, c *color.Color) { if f.DisabledColor || c == nil { - return fmt.Sprintf(format, args...) + return } - return c.SprintfFunc()(format, args...) + setWriter(c, buf) +} + +func (f *Formatter) unsetColor(buf *bytesBuffer, c *color.Color) { + if f.DisabledColor || c == nil { + return + } + unsetWriter(c, buf) } func (f *Formatter) sprintColor(c *color.Color, s string) string { @@ -94,99 +104,119 @@ func (f *Formatter) sprintColor(c *color.Color, s string) string { return c.SprintFunc()(s) } -func (f *Formatter) pretty(v interface{}, depth int) string { +//go:linkname setWriter github.com/fatih/color.(*Color).setWriter +func setWriter(*color.Color, io.Writer) *color.Color + +//go:linkname unsetWriter github.com/fatih/color.(*Color).unsetWriter +func unsetWriter(*color.Color, io.Writer) + +func (f *Formatter) pretty(buf *bytesBuffer, v interface{}, depth int) { switch val := v.(type) { case string: - return f.processString(val) + f.setColor(buf, f.StringColor) + f.writeString(buf, val) + f.unsetColor(buf, f.StringColor) case float64: - return f.sprintColor(f.NumberColor, strconv.FormatFloat(val, 'f', -1, 64)) + f.setColor(buf, f.NumberColor) + buf.writeFloat64(val) + f.unsetColor(buf, f.NumberColor) case json.Number: - return f.sprintColor(f.NumberColor, string(val)) + f.setColor(buf, f.NumberColor) + buf.WriteString(val.String()) + f.unsetColor(buf, f.NumberColor) case bool: - return f.sprintColor(f.BoolColor, strconv.FormatBool(val)) + f.setColor(buf, f.BoolColor) + if val { + buf.WriteString("true") + } else { + buf.WriteString("false") + } + f.unsetColor(buf, f.BoolColor) case nil: - return f.sprintColor(f.NullColor, "null") + f.setColor(buf, f.NullColor) + buf.WriteString("null") + f.unsetColor(buf, f.NullColor) case map[string]interface{}: - return f.processMap(val, depth) + f.writeMap(buf, val, depth) case []interface{}: - return f.processArray(val, depth) + f.writeArray(buf, val, depth) } - - return "" } -func (f *Formatter) processString(s string) string { - r := []rune(s) - if f.StringMaxLength != 0 && len(r) >= f.StringMaxLength { - s = string(r[0:f.StringMaxLength]) + "..." +func (f *Formatter) writeString(buf *bytesBuffer, s string) { + if f.StringMaxLength != 0 { + if r := []rune(s); len(r) >= f.StringMaxLength { + s = string(r[0:f.StringMaxLength]) + "..." + } } - - buf := &bytes.Buffer{} - encoder := json.NewEncoder(buf) - encoder.SetEscapeHTML(false) - encoder.Encode(s) - s = string(buf.Bytes()) - s = strings.TrimSuffix(s, "\n") - - return f.sprintColor(f.StringColor, s) + buf.writeString(s) } -func (f *Formatter) processMap(m map[string]interface{}, depth int) string { +func (f *Formatter) writeMap(buf *bytesBuffer, m map[string]interface{}, depth int) { if len(m) == 0 { - return "{}" + buf.WriteString("{}") + return } - currentIndent := f.generateIndent(depth - 1) - nextIndent := f.generateIndent(depth) - rows := []string{} - keys := []string{} - + keys := make([]string, len(m)) + var i int for key := range m { - keys = append(keys, key) + keys[i] = key + i++ } - sort.Strings(keys) - for _, key := range keys { - val := m[key] - buf := &bytes.Buffer{} - encoder := json.NewEncoder(buf) - encoder.SetEscapeHTML(false) - encoder.Encode(key) - s := strings.TrimSuffix(string(buf.Bytes()), "\n") - k := f.sprintColor(f.KeyColor, s) - v := f.pretty(val, depth+1) - - valueIndent := " " - if f.Newline == "" { - valueIndent = "" + buf.WriteByte('{') + for i, key := range keys { + if i > 0 { + buf.WriteByte(',') + } + buf.WriteString(f.Newline) + f.appendIndent(buf, depth) + f.setColor(buf, f.KeyColor) + buf.writeString(key) + f.unsetColor(buf, f.KeyColor) + buf.WriteByte(':') + if f.Newline != "" { + buf.WriteByte(' ') } - row := fmt.Sprintf("%s%s:%s%s", nextIndent, k, valueIndent, v) - rows = append(rows, row) + f.pretty(buf, m[key], depth+1) } - - return fmt.Sprintf("{%s%s%s%s}", f.Newline, strings.Join(rows, ","+f.Newline), f.Newline, currentIndent) + buf.WriteString(f.Newline) + f.appendIndent(buf, depth-1) + buf.WriteByte('}') } -func (f *Formatter) processArray(a []interface{}, depth int) string { +func (f *Formatter) writeArray(buf *bytesBuffer, a []interface{}, depth int) { if len(a) == 0 { - return "[]" + buf.WriteString("[]") + return } - currentIndent := f.generateIndent(depth - 1) - nextIndent := f.generateIndent(depth) - rows := []string{} - - for _, val := range a { - c := f.pretty(val, depth+1) - row := nextIndent + c - rows = append(rows, row) + buf.WriteByte('[') + for i, val := range a { + if i > 0 { + buf.WriteByte(',') + } + buf.WriteString(f.Newline) + f.appendIndent(buf, depth) + f.pretty(buf, val, depth+1) } - return fmt.Sprintf("[%s%s%s%s]", f.Newline, strings.Join(rows, ","+f.Newline), f.Newline, currentIndent) + buf.WriteString(f.Newline) + f.appendIndent(buf, depth-1) + buf.WriteByte(']') } -func (f *Formatter) generateIndent(depth int) string { - return strings.Repeat(" ", f.Indent*depth) +const spaces = " " + +func (f *Formatter) appendIndent(buf *bytesBuffer, depth int) { + if n := f.Indent * depth; n > 0 { + for n > len(spaces) { + buf.Write([]byte(spaces)) + n -= len(spaces) + } + buf.Write([]byte(spaces)[:n]) + } } // Marshal JSON data with default options. @@ -198,3 +228,26 @@ func Marshal(v interface{}) ([]byte, error) { func Format(data []byte) ([]byte, error) { return NewFormatter().Format(data) } + +type bytesBuffer struct { + *bytes.Buffer + enc *json.Encoder + scratch [64]byte +} + +func newBytesBuffer() *bytesBuffer { + buf := new(bytes.Buffer) + enc := json.NewEncoder(buf) + enc.SetEscapeHTML(false) + return &bytesBuffer{Buffer: buf, enc: enc} +} + +func (buf *bytesBuffer) writeString(str string) { + if buf.enc.Encode(str) == nil { + buf.Truncate(len(buf.Bytes()) - 1) + } +} + +func (buf *bytesBuffer) writeFloat64(f float64) { + buf.Write(strconv.AppendFloat(buf.scratch[:0], f, 'f', -1, 64)) +}