From 88045f21b7c8af62c40d06f150ecf58e1d4ead01 Mon Sep 17 00:00:00 2001
From: Kz Ho <kz.ho@shopee.com>
Date: Thu, 9 Mar 2023 12:30:08 +0800
Subject: [PATCH] Combine Stream.WriteString's slow path into one with
 escapeHTML bool flag - refer to
 https://go.dev/src/encoding/json/encode.go#L1029

---
 stream_str.go | 91 ++++++++++++++-------------------------------------
 1 file changed, 24 insertions(+), 67 deletions(-)

diff --git a/stream_str.go b/stream_str.go
index 5614bb6e..4816de72 100644
--- a/stream_str.go
+++ b/stream_str.go
@@ -235,15 +235,36 @@ func (stream *Stream) WriteStringWithHTMLEscaped(s string) {
 		stream.buf = append(stream.buf, '"')
 		return
 	}
-	writeStringSlowPathWithHTMLEscaped(stream, i, s, valLen)
+	writeStringSlowPath(stream, i, s, valLen, true)
 }
 
-func writeStringSlowPathWithHTMLEscaped(stream *Stream, i int, s string, valLen int) {
+// WriteString write string to stream without html escape
+func (stream *Stream) WriteString(s string) {
+	valLen := len(s)
+	stream.buf = append(stream.buf, '"')
+	// write string, the fast path, without utf8 and escape support
+	i := 0
+	for ; i < valLen; i++ {
+		c := s[i]
+		if c < utf8.RuneSelf && safeSet[c] {
+			stream.buf = append(stream.buf, c)
+		} else {
+			break
+		}
+	}
+	if i == valLen {
+		stream.buf = append(stream.buf, '"')
+		return
+	}
+	writeStringSlowPath(stream, i, s, valLen, false)
+}
+
+func writeStringSlowPath(stream *Stream, i int, s string, valLen int, escapeHTML bool) {
 	start := i
 	// for the remaining parts, we process them char by char
 	for i < valLen {
 		if b := s[i]; b < utf8.RuneSelf {
-			if htmlSafeSet[b] {
+			if htmlSafeSet[b] || (!escapeHTML && safeSet[b]) {
 				i++
 				continue
 			}
@@ -306,67 +327,3 @@ func writeStringSlowPathWithHTMLEscaped(stream *Stream, i int, s string, valLen
 	}
 	stream.writeByte('"')
 }
-
-// WriteString write string to stream without html escape
-func (stream *Stream) WriteString(s string) {
-	valLen := len(s)
-	stream.buf = append(stream.buf, '"')
-	// write string, the fast path, without utf8 and escape support
-	i := 0
-	for ; i < valLen; i++ {
-		c := s[i]
-		if c < utf8.RuneSelf && safeSet[c] {
-			stream.buf = append(stream.buf, c)
-		} else {
-			break
-		}
-	}
-	if i == valLen {
-		stream.buf = append(stream.buf, '"')
-		return
-	}
-	writeStringSlowPath(stream, i, s, valLen)
-}
-
-func writeStringSlowPath(stream *Stream, i int, s string, valLen int) {
-	start := i
-	// for the remaining parts, we process them char by char
-	for i < valLen {
-		if b := s[i]; b < utf8.RuneSelf {
-			if safeSet[b] {
-				i++
-				continue
-			}
-			if start < i {
-				stream.WriteRaw(s[start:i])
-			}
-			switch b {
-			case '\\', '"':
-				stream.writeTwoBytes('\\', b)
-			case '\n':
-				stream.writeTwoBytes('\\', 'n')
-			case '\r':
-				stream.writeTwoBytes('\\', 'r')
-			case '\t':
-				stream.writeTwoBytes('\\', 't')
-			default:
-				// This encodes bytes < 0x20 except for \t, \n and \r.
-				// If escapeHTML is set, it also escapes <, >, and &
-				// because they can lead to security holes when
-				// user-controlled strings are rendered into JSON
-				// and served to some browsers.
-				stream.WriteRaw(`\u00`)
-				stream.writeTwoBytes(hex[b>>4], hex[b&0xF])
-			}
-			i++
-			start = i
-			continue
-		}
-		i++
-		continue
-	}
-	if start < len(s) {
-		stream.WriteRaw(s[start:])
-	}
-	stream.writeByte('"')
-}