gopls/internal/lsp/source: highlight returns correctly

Rewrite the control flow highlighting logic to be more understandable,
and to fix a bug where multi-name result parameters were not properly
counted.

Fixes golang/go#60589

Change-Id: Id12ada78852be0a8376fbd482326322fc1b87fcf
Reviewed-on: https://go-review.googlesource.com/c/tools/+/503439
Auto-Submit: Robert Findley <rfindley@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

gopls/internal/golang/highlight.go

@@ -65,6 +65,8 @@ func Highlight(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle, po
 	return ranges, nil
 }
 
+// highlightPath returns ranges to highlight for the given enclosing path,
+// which should be the result of astutil.PathEnclosingInterval.
 func highlightPath(path []ast.Node, file *ast.File, info *types.Info) (map[posRange]struct{}, error) {
 	result := make(map[posRange]struct{})
 	switch node := path[0].(type) {
@@ -133,115 +135,167 @@ type posRange struct {
 	start, end token.Pos
 }
 
-func highlightFuncControlFlow(path []ast.Node, result map[posRange]struct{}) {
-	var enclosingFunc ast.Node
-	var returnStmt *ast.ReturnStmt
-	var resultsList *ast.FieldList
-	inReturnList := false
-
-Outer:
-	// Reverse walk the path till we get to the func block.
+// highlightFuncControlFlow adds highlight ranges to the result map to
+// associate results and result parameters.
+//
+// Specifically, if the cursor is in a result or result parameter, all
+// results and result parameters with the same index are highlighted. If the
+// cursor is in a 'func' or 'return' keyword, the func keyword as well as all
+// returns from that func are highlighted.
+//
+// As a special case, if the cursor is within a complicated expression, control
+// flow highlighting is disabled, as it would highlight too much.
+func highlightFuncControlFlow(path []ast.Node, result map[posRange]unit) {
+
+	var (
+		funcType   *ast.FuncType   // type of enclosing func, or nil
+		funcBody   *ast.BlockStmt  // body of enclosing func, or nil
+		returnStmt *ast.ReturnStmt // enclosing ReturnStmt within the func, or nil
+	)
+
+findEnclosingFunc:
 	for i, n := range path {
-		switch node := n.(type) {
+		switch n := n.(type) {
+		// TODO(rfindley, low priority): these pre-existing cases for KeyValueExpr
+		// and CallExpr appear to avoid highlighting when the cursor is in a
+		// complicated expression. However, the basis for this heuristic is
+		// unclear. Can we formalize a rationale?
 		case *ast.KeyValueExpr:
-			// If cursor is in a key: value expr, we don't want control flow highlighting
+			// If cursor is in a key: value expr, we don't want control flow highlighting.
 			return
+
 		case *ast.CallExpr:
 			// If cursor is an arg in a callExpr, we don't want control flow highlighting.
 			if i > 0 {
-				for _, arg := range node.Args {
+				for _, arg := range n.Args {
 					if arg == path[i-1] {
 						return
 					}
 				}
 			}
-		case *ast.Field:
-			inReturnList = true
+
 		case *ast.FuncLit:
-			enclosingFunc = n
-			resultsList = node.Type.Results
-			break Outer
+			funcType = n.Type
+			funcBody = n.Body
+			break findEnclosingFunc
+
 		case *ast.FuncDecl:
-			enclosingFunc = n
-			resultsList = node.Type.Results
-			break Outer
+			funcType = n.Type
+			funcBody = n.Body
+			break findEnclosingFunc
+
 		case *ast.ReturnStmt:
-			returnStmt = node
-			// If the cursor is not directly in a *ast.ReturnStmt, then
-			// we need to know if it is within one of the values that is being returned.
-			inReturnList = inReturnList || path[0] != returnStmt
-		}
-	}
-	// Cursor is not in a function.
-	if enclosingFunc == nil {
-		return
-	}
-	// If the cursor is on a "return" or "func" keyword, we should highlight all of the exit
-	// points of the function, including the "return" and "func" keywords.
-	highlightAllReturnsAndFunc := path[0] == returnStmt || path[0] == enclosingFunc
-	switch path[0].(type) {
-	case *ast.Ident, *ast.BasicLit:
-		// Cursor is in an identifier and not in a return statement or in the results list.
-		if returnStmt == nil && !inReturnList {
-			return
+			returnStmt = n
 		}
-	case *ast.FuncType:
-		highlightAllReturnsAndFunc = true
 	}
-	// The user's cursor may be within the return statement of a function,
-	// or within the result section of a function's signature.
-	// index := -1
-	var nodes []ast.Node
-	if returnStmt != nil {
-		for _, n := range returnStmt.Results {
-			nodes = append(nodes, n)
-		}
-	} else if resultsList != nil {
-		for _, n := range resultsList.List {
-			nodes = append(nodes, n)
-		}
-	}
-	_, index := nodeAtPos(nodes, path[0].Pos())
 
-	// Highlight the correct argument in the function declaration return types.
-	if resultsList != nil && -1 < index && index < len(resultsList.List) {
-		rng := posRange{
-			start: resultsList.List[index].Pos(),
-			end:   resultsList.List[index].End(),
-		}
-		result[rng] = struct{}{}
+	if funcType == nil {
+		return // cursor is not in a function
 	}
-	// Add the "func" part of the func declaration.
-	if highlightAllReturnsAndFunc {
-		r := posRange{
-			start: enclosingFunc.Pos(),
-			end:   enclosingFunc.Pos() + token.Pos(len("func")),
+
+	// Helper functions for inspecting the current location.
+	var (
+		pos    = path[0].Pos()
+		inSpan = func(start, end token.Pos) bool { return start <= pos && pos < end }
+		inNode = func(n ast.Node) bool { return inSpan(n.Pos(), n.End()) }
+	)
+
+	inResults := funcType.Results != nil && inNode(funcType.Results)
+
+	// If the cursor is on a "return" or "func" keyword, but not highlighting any
+	// specific field or expression, we should highlight all of the exit points
+	// of the function, including the "return" and "func" keywords.
+	funcEnd := funcType.Func + token.Pos(len("func"))
+	highlightAll := path[0] == returnStmt || inSpan(funcType.Func, funcEnd)
+	var highlightIndexes map[int]bool
+
+	if highlightAll {
+		// Add the "func" part of the func declaration.
+		result[posRange{
+			start: funcType.Func,
+			end:   funcEnd,
+		}] = unit{}
+	} else if returnStmt == nil && !inResults {
+		return // nothing to highlight
+	} else {
+		// If we're not highighting the entire return statement, we need to collect
+		// specific result indexes to highlight. This may be more than one index if
+		// the cursor is on a multi-name result field, but not in any specific name.
+		if !highlightAll {
+			highlightIndexes = make(map[int]bool)
+			if returnStmt != nil {
+				for i, n := range returnStmt.Results {
+					if inNode(n) {
+						highlightIndexes[i] = true
+						break
+					}
+				}
+			}
+
+			// Scan fields, either adding highlights according to the highlightIndexes
+			// computed above, or accounting for the cursor position within the result
+			// list.
+			// (We do both at once to avoid repeating the cumbersome field traversal.)
+			i := 0
+		findField:
+			for _, field := range funcType.Results.List {
+				for j, name := range field.Names {
+					if inNode(name) || highlightIndexes[i+j] {
+						result[posRange{name.Pos(), name.End()}] = unit{}
+						highlightIndexes[i+j] = true
+						break findField // found/highlighted the specific name
+					}
+				}
+				// If the cursor is in a field but not in a name (e.g. in the space, or
+				// the type), highlight the whole field.
+				//
+				// Note that this may not be ideal if we're at e.g.
+				//
+				//  (x,‸y int, z int8)
+				//
+				// ...where it would make more sense to highlight only y. But we don't
+				// reach this function if not in a func, return, ident, or basiclit.
+				if inNode(field) || highlightIndexes[i] {
+					result[posRange{field.Pos(), field.End()}] = unit{}
+					highlightIndexes[i] = true
+					if inNode(field) {
+						for j := range field.Names {
+							highlightIndexes[i+j] = true
+						}
+					}
+					break findField // found/highlighted the field
+				}
+
+				n := len(field.Names)
+				if n == 0 {
+					n = 1
+				}
+				i += n
+			}
 		}
-		result[r] = struct{}{}
 	}
-	ast.Inspect(enclosingFunc, func(n ast.Node) bool {
-		// Don't traverse any other functions.
-		switch n.(type) {
+
+	ast.Inspect(funcBody, func(n ast.Node) bool {
+		switch n := n.(type) {
 		case *ast.FuncDecl, *ast.FuncLit:
-			return enclosingFunc == n
-		}
-		ret, ok := n.(*ast.ReturnStmt)
-		if !ok {
-			return true
-		}
-		var toAdd ast.Node
-		// Add the entire return statement, applies when highlight the word "return" or "func".
-		if highlightAllReturnsAndFunc {
-			toAdd = n
-		}
-		// Add the relevant field within the entire return statement.
-		if -1 < index && index < len(ret.Results) {
-			toAdd = ret.Results[index]
-		}
-		if toAdd != nil {
-			result[posRange{start: toAdd.Pos(), end: toAdd.End()}] = struct{}{}
+			// Don't traverse into any functions other than enclosingFunc.
+			return false
+		case *ast.ReturnStmt:
+			if highlightAll {
+				// Add the entire return statement.
+				result[posRange{n.Pos(), n.End()}] = unit{}
+			} else {
+				// Add the highlighted indexes.
+				for i, expr := range n.Results {
+					if highlightIndexes[i] {
+						result[posRange{expr.Pos(), expr.End()}] = unit{}
+					}
+				}
+			}
+			return false
+
 		}
-		return false
+		return true
 	})
 }
 

gopls/internal/test/marker/testdata/highlight/issue60589.txt

@@ -0,0 +1,30 @@
+This test verifies that control flow lighlighting correctly accounts for
+multi-name result parameters. In golang/go#60589, it did not.
+
+-- go.mod --
+module mod.com
+
+go 1.18
+
+-- p.go --
+package p
+
+func _() (foo int, bar, baz string) { //@ loc(func, "func"), loc(foo, "foo"), loc(fooint, "foo int"), loc(int, "int"), loc(bar, "bar"), loc(beforebaz, " baz"), loc(baz, "baz"), loc(barbazstring, "bar, baz string"), loc(beforestring, re`() string`), loc(string, "string")
+	return 0, "1", "2" //@ loc(return, `return 0, "1", "2"`), loc(l0, "0"), loc(l1, `"1"`), loc(l2, `"2"`)
+}
+
+// Assertions, expressed here to avoid clutter above.
+// Note that when the cursor is over the field type, there is some
+// (likely harmless) redundancy.
+
+//@ highlight(func, func, return)
+//@ highlight(foo, foo, l0)
+//@ highlight(int, fooint, int, l0)
+//@ highlight(bar, bar, l1)
+//@ highlight(beforebaz)
+//@ highlight(baz, baz, l2)
+//@ highlight(beforestring, baz, l2)
+//@ highlight(string, barbazstring, string, l1, l2)
+//@ highlight(l0, foo, l0)
+//@ highlight(l1, bar, l1)
+//@ highlight(l2, baz, l2)