tools/fix: merge calls to list.Concat

Previously, `cue fix` would correctly convert list addition to calls to
`list.Concat`. However, chains of list-addition would not be merged,
which causes `cue fix` to create somewhat unreadable code, such as:

    l: list.Concat([list.Concat([list.Concat([["a"], ["b"]]), ["c"]), ["d"]])

But it's not tricky to detect that an argument to + is itself a call to
`list.Concat` and thus flatten such trees to a single call to
list.Concat.

Fixes #3523

Change-Id: I9bb57fe3d576725ad5d2ae850b809d6b858e9f51
Signed-off-by: Matthew Sackman <[email protected]>
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1202868
TryBot-Result: CUEcueckoo <[email protected]>
Reviewed-by: Daniel Martí <[email protected]>

Author: cuematthew

tools/fix/fix.go

@@ -69,19 +69,27 @@ func File(f *ast.File, o ...Option) *ast.File {
 				x, y := n.X, n.Y
 				_, xIsList := x.(*ast.ListLit)
 				_, yIsList := y.(*ast.ListLit)
-				if !(xIsList || yIsList) {
-					break
-				}
+				_, xIsConcat := concatCallArgs(x)
+				_, yIsConcat := concatCallArgs(y)
+
 				if n.Op == token.ADD {
+					if !(xIsList || xIsConcat || yIsList || yIsConcat) {
+						break
+					}
 					// Rewrite list addition to use list.Concat
+					exprs := expandConcats(x, y)
 					ast.SetRelPos(x, token.NoSpace)
 					c.Replace(ast.NewCall(
 						ast.NewSel(&ast.Ident{
 							Name: "list",
 							Node: ast.NewImport(nil, "list"),
-						}, "Concat"), ast.NewList(x, y)),
+						}, "Concat"), ast.NewList(exprs...)),
 					)
+
 				} else {
+					if !(xIsList || yIsList) {
+						break
+					}
 					// Rewrite list multiplication to use list.Repeat
 					if !xIsList {
 						x, y = y, x
@@ -115,3 +123,42 @@ func File(f *ast.File, o ...Option) *ast.File {
 	}
 	return f
 }
+
+func expandConcats(exprs ...ast.Expr) (result []ast.Expr) {
+	for _, expr := range exprs {
+		list, ok := concatCallArgs(expr)
+		if ok {
+			result = append(result, expandConcats(list.Elts...)...)
+		} else {
+			result = append(result, expr)
+		}
+	}
+	return result
+}
+
+func concatCallArgs(expr ast.Expr) (*ast.ListLit, bool) {
+	call, ok := expr.(*ast.CallExpr)
+	if !ok {
+		return nil, false
+	}
+	sel, ok := call.Fun.(*ast.SelectorExpr)
+	if !ok {
+		return nil, false
+	}
+	name, ok := sel.X.(*ast.Ident)
+	if !ok || name.Name != "list" {
+		return nil, false
+	}
+	name, ok = sel.Sel.(*ast.Ident)
+	if !ok || name.Name != "Concat" {
+		return nil, false
+	}
+	if len(call.Args) != 1 {
+		return nil, false
+	}
+	list, ok := call.Args[0].(*ast.ListLit)
+	if !ok {
+		return nil, false
+	}
+	return list, true
+}

tools/fix/fix_test.go

@@ -73,6 +73,10 @@ b: a + a
 c: a + [8]
 d: [9] + a
 e: [0] + [1]
+f: [0] + [1] + [2]
+g: list.Concat([[0], [1, 2]]) + list.Concat([[3, 4], [5]])
+h: list.Concat([[0], [1, 2]]) + [3] + [4] + list.Concat([[5, 6], [7]])
+i: list.Concat(list.Concat([[0], [1, 2]]), list.Concat([[3, 4], [5]]))
 `,
 			out: `import "list"
 
@@ -81,6 +85,10 @@ b: a + a
 c: list.Concat([a, [8]])
 d: list.Concat([[9], a])
 e: list.Concat([[0], [1]])
+f: list.Concat([[0], [1], [2]])
+g: list.Concat([[0], [1, 2], [3, 4], [5]])
+h: list.Concat([[0], [1, 2], [3], [4], [5, 6], [7]])
+i: list.Concat(list.Concat([[0], [1, 2]]), list.Concat([[3, 4], [5]]))
 `,
 		},
 
@@ -92,6 +100,7 @@ c: 4
 d: [7] * c
 e: c * [8]
 f: [9] * 5
+g: ([9] * 5) + (6 * [10])
 `,
 			out: `import "list"
 
@@ -101,6 +110,7 @@ c: 4
 d: list.Repeat([7], c)
 e: list.Repeat([8], c)
 f: list.Repeat([9], 5)
+g: (list.Repeat([9], 5)) + (list.Repeat([10], 6))
 `,
 		},
 	}