From ef23ce9e92d82a8868210ab27d4777c0980fcd7a Mon Sep 17 00:00:00 2001
From: Jordan Krage <jmank88@gmail.com>
Date: Tue, 27 Jun 2017 20:24:37 -0500
Subject: [PATCH] WriteTo string concat allocation reduction (#177)

* reduce string concat allocs in Tree.writeTo
* fix failingWriter and usages
---
 tomltree_write.go      | 21 +++++++++++++++------
 tomltree_write_test.go | 22 +++++++++++-----------
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/tomltree_write.go b/tomltree_write.go
index cd03f9d..ca763ed 100644
--- a/tomltree_write.go
+++ b/tomltree_write.go
@@ -118,8 +118,7 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64) (
 			return bytesCount, err
 		}
 
-		kvRepr := indent + k + " = " + repr + "\n"
-		writtenBytesCount, err := w.Write([]byte(kvRepr))
+		writtenBytesCount, err := writeStrings(w, indent, k, " = ", repr, "\n")
 		bytesCount += int64(writtenBytesCount)
 		if err != nil {
 			return bytesCount, err
@@ -137,8 +136,7 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64) (
 		switch node := v.(type) {
 		// node has to be of those two types given how keys are sorted above
 		case *Tree:
-			tableName := "\n" + indent + "[" + combinedKey + "]\n"
-			writtenBytesCount, err := w.Write([]byte(tableName))
+			writtenBytesCount, err := writeStrings(w, "\n", indent, "[", combinedKey, "]\n")
 			bytesCount += int64(writtenBytesCount)
 			if err != nil {
 				return bytesCount, err
@@ -149,8 +147,7 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64) (
 			}
 		case []*Tree:
 			for _, subTree := range node {
-				tableArrayName := "\n" + indent + "[[" + combinedKey + "]]\n"
-				writtenBytesCount, err := w.Write([]byte(tableArrayName))
+				writtenBytesCount, err := writeStrings(w, "\n", indent, "[[", combinedKey, "]]\n")
 				bytesCount += int64(writtenBytesCount)
 				if err != nil {
 					return bytesCount, err
@@ -167,6 +164,18 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64) (
 	return bytesCount, nil
 }
 
+func writeStrings(w io.Writer, s ...string) (int, error) {
+	var n int
+	for i := range s {
+		b, err := io.WriteString(w, s[i])
+		n += b
+		if err != nil {
+			return n, err
+		}
+	}
+	return n, nil
+}
+
 // WriteTo encode the Tree as Toml and writes it to the writer w.
 // Returns the number of bytes written in case of success, or an error if anything happened.
 func (t *Tree) WriteTo(w io.Writer) (int64, error) {
diff --git a/tomltree_write_test.go b/tomltree_write_test.go
index 0edf1be..c2a1ce3 100644
--- a/tomltree_write_test.go
+++ b/tomltree_write_test.go
@@ -16,26 +16,26 @@ type failingWriter struct {
 	buffer  bytes.Buffer
 }
 
-func (f failingWriter) Write(p []byte) (n int, err error) {
+func (f *failingWriter) Write(p []byte) (n int, err error) {
 	count := len(p)
-	toWrite := f.failAt - count + f.written
+	toWrite := f.failAt - (count + f.written)
 	if toWrite < 0 {
 		toWrite = 0
 	}
 	if toWrite > count {
 		f.written += count
-		f.buffer.WriteString(string(p))
+		f.buffer.Write(p)
 		return count, nil
 	}
 
-	f.buffer.WriteString(string(p[:toWrite]))
+	f.buffer.Write(p[:toWrite])
 	f.written = f.failAt
-	return f.written, fmt.Errorf("failingWriter failed after writting %d bytes", f.written)
+	return toWrite, fmt.Errorf("failingWriter failed after writting %d bytes", f.written)
 }
 
 func assertErrorString(t *testing.T, expected string, err error) {
 	expectedErr := errors.New(expected)
-	if err.Error() != expectedErr.Error() {
+	if err == nil || err.Error() != expectedErr.Error() {
 		t.Errorf("expecting error %s, but got %s instead", expected, err)
 	}
 }
@@ -175,7 +175,7 @@ func TestTreeWriteToInvalidTreeTomlValueArray(t *testing.T) {
 func TestTreeWriteToFailingWriterInSimpleValue(t *testing.T) {
 	toml, _ := Load(`a = 2`)
 	writer := failingWriter{failAt: 0, written: 0}
-	_, err := toml.WriteTo(writer)
+	_, err := toml.WriteTo(&writer)
 	assertErrorString(t, "failingWriter failed after writting 0 bytes", err)
 }
 
@@ -184,11 +184,11 @@ func TestTreeWriteToFailingWriterInTable(t *testing.T) {
 [b]
 a = 2`)
 	writer := failingWriter{failAt: 2, written: 0}
-	_, err := toml.WriteTo(writer)
+	_, err := toml.WriteTo(&writer)
 	assertErrorString(t, "failingWriter failed after writting 2 bytes", err)
 
 	writer = failingWriter{failAt: 13, written: 0}
-	_, err = toml.WriteTo(writer)
+	_, err = toml.WriteTo(&writer)
 	assertErrorString(t, "failingWriter failed after writting 13 bytes", err)
 }
 
@@ -197,11 +197,11 @@ func TestTreeWriteToFailingWriterInArray(t *testing.T) {
 [[b]]
 a = 2`)
 	writer := failingWriter{failAt: 2, written: 0}
-	_, err := toml.WriteTo(writer)
+	_, err := toml.WriteTo(&writer)
 	assertErrorString(t, "failingWriter failed after writting 2 bytes", err)
 
 	writer = failingWriter{failAt: 15, written: 0}
-	_, err = toml.WriteTo(writer)
+	_, err = toml.WriteTo(&writer)
 	assertErrorString(t, "failingWriter failed after writting 15 bytes", err)
 }