cmd/compile: optimize []byte(string1 + string2)

This CL optimizes the compilation of string-to-bytes conversion in the
case of string additions.

Fixes #62407

Change-Id: Ic47df758478e5d061880620025c4ec7dbbff8a64
Reviewed-on: https://go-review.googlesource.com/c/go/+/527935
Reviewed-by: Cuong Manh Le <cuong.manhle.vn@gmail.com>
Reviewed-by: Keith Randall <khr@golang.org>
Auto-Submit: Keith Randall <khr@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Tim King <taking@google.com>

2 files changed, 46 insertions, 26 deletions

diff --git a/src/cmd/compile/internal/walk/convert.go b/src/cmd/compile/internal/walk/convert.go
index 280b3b65e8..3118233697 100644
--- a/src/cmd/compile/internal/walk/convert.go
+++ b/src/cmd/compile/internal/walk/convert.go
@@ -270,6 +270,11 @@ func walkRuneToString(n *ir.ConvExpr, init *ir.Nodes) ir.Node {
 // walkStringToBytes walks an OSTR2BYTES node.
 func walkStringToBytes(n *ir.ConvExpr, init *ir.Nodes) ir.Node {
 	s := n.X
+
+	if expr, ok := s.(*ir.AddStringExpr); ok {
+		return walkAddString(n.Type(), expr, init)
+	}
+
 	if ir.IsConst(s, constant.String) {
 		sc := ir.StringVal(s)
 
diff --git a/src/cmd/compile/internal/walk/expr.go b/src/cmd/compile/internal/walk/expr.go
index f73b5d9503..8c36e03aa0 100644
--- a/src/cmd/compile/internal/walk/expr.go
+++ b/src/cmd/compile/internal/walk/expr.go
@@ -273,7 +273,7 @@ func walkExpr1(n ir.Node, init *ir.Nodes) ir.Node {
 		return walkNew(n, init)
 
 	case ir.OADDSTR:
-		return walkAddString(n.(*ir.AddStringExpr), init)
+		return walkAddString(n.Type(), n.(*ir.AddStringExpr), init)
 
 	case ir.OAPPEND:
 		// order should make sure we only see OAS(node, OAPPEND), which we handle above.
@@ -464,49 +464,64 @@ func copyExpr(n ir.Node, t *types.Type, init *ir.Nodes) ir.Node {
 	return l
 }
 
-func walkAddString(n *ir.AddStringExpr, init *ir.Nodes) ir.Node {
+func walkAddString(typ *types.Type, n *ir.AddStringExpr, init *ir.Nodes) ir.Node {
 	c := len(n.List)
 
 	if c < 2 {
 		base.Fatalf("walkAddString count %d too small", c)
 	}
 
-	buf := typecheck.NodNil()
-	if n.Esc() == ir.EscNone {
-		sz := int64(0)
-		for _, n1 := range n.List {
-			if n1.Op() == ir.OLITERAL {
-				sz += int64(len(ir.StringVal(n1)))
+	// list of string arguments
+	var args []ir.Node
+
+	var fn, fnsmall, fnbig string
+
+	switch {
+	default:
+		base.FatalfAt(n.Pos(), "unexpected type: %v", typ)
+	case typ.IsString():
+		buf := typecheck.NodNil()
+		if n.Esc() == ir.EscNone {
+			sz := int64(0)
+			for _, n1 := range n.List {
+				if n1.Op() == ir.OLITERAL {
+					sz += int64(len(ir.StringVal(n1)))
+				}
 			}
-		}
 
-		// Don't allocate the buffer if the result won't fit.
-		if sz < tmpstringbufsize {
-			// Create temporary buffer for result string on stack.
-			buf = stackBufAddr(tmpstringbufsize, types.Types[types.TUINT8])
+			// Don't allocate the buffer if the result won't fit.
+			if sz < tmpstringbufsize {
+				// Create temporary buffer for result string on stack.
+				buf = stackBufAddr(tmpstringbufsize, types.Types[types.TUINT8])
+			}
 		}
-	}
 
-	// build list of string arguments
-	args := []ir.Node{buf}
-	for _, n2 := range n.List {
-		args = append(args, typecheck.Conv(n2, types.Types[types.TSTRING]))
+		args = []ir.Node{buf}
+		fnsmall, fnbig = "concatstring%d", "concatstrings"
+	case typ.IsSlice() && typ.Elem().IsKind(types.TUINT8): // Optimize []byte(str1+str2+...)
+		fnsmall, fnbig = "concatbyte%d", "concatbytes"
 	}
 
-	var fn string
 	if c <= 5 {
 		// small numbers of strings use direct runtime helpers.
 		// note: order.expr knows this cutoff too.
-		fn = fmt.Sprintf("concatstring%d", c)
+		fn = fmt.Sprintf(fnsmall, c)
+
+		for _, n2 := range n.List {
+			args = append(args, typecheck.Conv(n2, types.Types[types.TSTRING]))
+		}
 	} else {
 		// large numbers of strings are passed to the runtime as a slice.
-		fn = "concatstrings"
-
+		fn = fnbig
 		t := types.NewSlice(types.Types[types.TSTRING])
-		// args[1:] to skip buf arg
-		slice := ir.NewCompLitExpr(base.Pos, ir.OCOMPLIT, t, args[1:])
+
+		slargs := make([]ir.Node, len(n.List))
+		for i, n2 := range n.List {
+			slargs[i] = typecheck.Conv(n2, types.Types[types.TSTRING])
+		}
+		slice := ir.NewCompLitExpr(base.Pos, ir.OCOMPLIT, t, slargs)
 		slice.Prealloc = n.Prealloc
-		args = []ir.Node{buf, slice}
+		args = append(args, slice)
 		slice.SetEsc(ir.EscNone)
 	}
 
@@ -515,7 +530,7 @@ func walkAddString(n *ir.AddStringExpr, init *ir.Nodes) ir.Node {
 	r.Args = args
 	r1 := typecheck.Expr(r)
 	r1 = walkExpr(r1, init)
-	r1.SetType(n.Type())
+	r1.SetType(typ)
 
 	return r1
 }