Implement io.StringWriter on some more types (#2023)

In theory this can optimize some code paths where a string first needs
to be converted to a []byte to use the normal Write method. By
implementing WriteString this extra copy isn't needed. Internall we
don't do the copy and just use s2b instead.

Author: erikdubbelboer

compress.go

@@ -351,6 +351,11 @@ func (w *byteSliceWriter) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+func (w *byteSliceWriter) WriteString(s string) (int, error) {
+	w.b = append(w.b, s...)
+	return len(s), nil
+}
+
 type byteSliceReader struct {
 	b []byte
 }

fasthttputil/pipeconns.go

@@ -142,6 +142,10 @@ func (c *pipeConn) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+func (c *pipeConn) WriteString(s string) (int, error) {
+	return c.Write(s2b(s))
+}
+
 func (c *pipeConn) Read(p []byte) (int, error) {
 	mayBlock := true
 	nn := 0

fasthttputil/s2b.go

@@ -0,0 +1,8 @@
+package fasthttputil
+
+import "unsafe"
+
+// s2b converts string to a byte slice without memory allocation.
+func s2b(s string) []byte {
+	return unsafe.Slice(unsafe.StringData(s), len(s))
+}

http.go

@@ -378,6 +378,11 @@ func (w *responseBodyWriter) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+func (w *responseBodyWriter) WriteString(s string) (int, error) {
+	w.r.AppendBodyString(s)
+	return len(s), nil
+}
+
 type requestBodyWriter struct {
 	r *Request
 }
@@ -387,6 +392,11 @@ func (w *requestBodyWriter) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+func (w *requestBodyWriter) WriteString(s string) (int, error) {
+	w.r.AppendBodyString(s)
+	return len(s), nil
+}
+
 func (resp *Response) ParseNetConn(conn net.Conn) {
 	resp.raddr = conn.RemoteAddr()
 	resp.laddr = conn.LocalAddr()
@@ -1544,6 +1554,12 @@ func (w *statsWriter) Write(p []byte) (int, error) {
 	return n, err
 }
 
+func (w *statsWriter) WriteString(s string) (int, error) {
+	n, err := w.w.Write(s2b(s))
+	w.bytesWritten += int64(n)
+	return n, err
+}
+
 func acquireStatsWriter(w io.Writer) *statsWriter {
 	v := statsWriterPool.Get()
 	if v == nil {
@@ -1969,6 +1985,10 @@ func (w *flushWriter) Write(p []byte) (int, error) {
 	return n, nil
 }
 
+func (w *flushWriter) WriteString(s string) (int, error) {
+	return w.Write(s2b(s))
+}
+
 // Write writes response to w.
 //
 // Write doesn't flush response to w for performance reasons.

stackless/s2b.go

@@ -0,0 +1,8 @@
+package stackless
+
+import "unsafe"
+
+// s2b converts string to a byte slice without memory allocation.
+func s2b(s string) []byte {
+	return unsafe.Slice(unsafe.StringData(s), len(s))
+}

stackless/writer.go

@@ -67,6 +67,13 @@ func (w *writer) Write(p []byte) (int, error) {
 	return w.n, err
 }
 
+func (w *writer) WriteString(s string) (int, error) {
+	w.p = s2b(s)
+	err := w.do(opWrite)
+	w.p = nil
+	return w.n, err
+}
+
 func (w *writer) Flush() error {
 	return w.do(opFlush)
 }