From 1961e23a289f312ca8aa94790f41351a18633e05 Mon Sep 17 00:00:00 2001
From: "M. J. Fromberger" <michael.j.fromberger@gmail.com>
Date: Thu, 12 Sep 2024 14:45:26 -0700
Subject: [PATCH] shell: pool buffers and avoid reallocations in Quote and Join
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This is a mild improvement to speed, and a substantial amortized improvement
for memory usage. Pooling helps across many uses, but for Join in particular,
appending to a buffer avoids the double-allocation of a slice for quoted
strings followed by the string join to add the spaces.

                   │ before.txt │ after.txt
                   │ sec/op     │ sec/op    vs base
    Quote/Quote-10   751.6µ       737.3µ    -1.89%
    Quote/Join-10    394.3µ       340.5µ    -13.64%

                   │ before.txt │ after.txt
                   │ B/op       │ B/op      vs base
    Quote/Quote-10    512.0Ki     200.3Ki   -60.88%
    Quote/Join-10    165.78Ki     56.02Ki   -66.21%

                   │ before.txt │ after.txt
                   │ allocs/op  │ allocs/op vs base
    Quote/Quote-10      3.000     1.000     -66.67%
    Quote/Join-10    1880.000     1.000     -99.95%
---
 shell/shell.go | 52 +++++++++++++++++++++++++++++++++++---------------
 1 file changed, 37 insertions(+), 15 deletions(-)

diff --git a/shell/shell.go b/shell/shell.go
index 8bd91f8..d6d0850 100644
--- a/shell/shell.go
+++ b/shell/shell.go
@@ -290,26 +290,42 @@ func quotable(s string) (hasQ, hasOther bool) {
 	return v&quote != 0, v&other != 0
 }
 
+var bufPool = &sync.Pool{
+	New: func() any { return new(bytes.Buffer) },
+}
+
 // Quote returns a copy of s in which shell metacharacters are quoted to
 // protect them from evaluation.
 func Quote(s string) string {
-	var buf bytes.Buffer
-	return quote(s, &buf)
+	if s == "" {
+		return "''"
+	}
+	if hasQ, hasOther := quotable(s); !hasQ && !hasOther {
+		return s // no quotation needed
+	}
+	buf := bufPool.Get().(*bytes.Buffer)
+	defer bufPool.Put(buf)
+	buf.Reset()
+
+	quote(s, buf)
+	return buf.String()
 }
 
-// quote implements quotation, using the provided buffer as scratch space.  The
-// existing contents of the buffer are clobbered.
-func quote(s string, buf *bytes.Buffer) string {
+// quote appends the quoted representation of s to buf.
+// Any existing contents in buf are not modified.
+func quote(s string, buf *bytes.Buffer) {
 	if s == "" {
-		return "''"
+		buf.WriteString("''")
+		return
 	}
+
 	hasQ, hasOther := quotable(s)
 	if !hasQ && !hasOther {
-		return s // fast path: nothing needs quotation
+		buf.WriteString(s)
+		return // fast path: nothing needs quotation
 	}
 
-	buf.Reset()
-	buf.Grow(len(s))
+	buf.Grow(len(s) + 2) // +2 for expected quotes; it may be more
 	inq := false
 	for i := range len(s) {
 		ch := s[i]
@@ -328,16 +344,22 @@ func quote(s string, buf *bytes.Buffer) string {
 	if inq {
 		buf.WriteByte('\'')
 	}
-	return buf.String()
 }
 
 // Join quotes each element of ss with Quote and concatenates the resulting
 // strings separated by spaces.
 func Join(ss []string) string {
-	quoted := make([]string, len(ss))
-	var buf bytes.Buffer
-	for i, s := range ss {
-		quoted[i] = quote(s, &buf)
+	if len(ss) == 0 {
+		return ""
 	}
-	return strings.Join(quoted, " ")
+	buf := bufPool.Get().(*bytes.Buffer)
+	defer bufPool.Put(buf)
+	buf.Reset()
+
+	quote(ss[0], buf)
+	for _, s := range ss[1:] {
+		buf.WriteByte(' ')
+		quote(s, buf)
+	}
+	return buf.String()
 }