internal/lsp: add control flow highlighting inside for loops

When the cursor is on a "for" statement or on any branch statement inside
the for loop. It will highlight the control flow inside the for loop.

Updates #34496

Change-Id: Idef14e3c89bc161d305d4a49fd784095a93bbc03
Reviewed-on: https://go-review.googlesource.com/c/tools/+/208337
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/source/highlight.go

@@ -10,6 +10,7 @@ import (
 
 	"golang.org/x/tools/go/ast/astutil"
 	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/telemetry/log"
 	"golang.org/x/tools/internal/telemetry/trace"
 	errors "golang.org/x/xerrors"
 )
@@ -55,33 +56,94 @@ func Highlight(ctx context.Context, snapshot Snapshot, f File, pos protocol.Posi
 	if len(path) == 0 {
 		return nil, errors.Errorf("no enclosing position found for %v:%v", int(pos.Line), int(pos.Character))
 	}
-	result := []protocol.Range{}
-	id, ok := path[0].(*ast.Ident)
-	if !ok {
-		// If the cursor is not within an identifier, return empty results.
-		return result, nil
-	}
-	idObj := pkg.GetTypesInfo().ObjectOf(id)
 
-	ast.Inspect(path[len(path)-1], func(node ast.Node) bool {
-		n, ok := node.(*ast.Ident)
-		if !ok {
-			return true
+	switch path[0].(type) {
+	case *ast.Ident:
+		return highlightIdentifiers(ctx, snapshot, m, path, pkg)
+	case *ast.BranchStmt, *ast.ForStmt, *ast.RangeStmt:
+		return highlightControlFlow(ctx, snapshot, m, path)
+	}
+
+	// If the cursor is in an unidentified area, return empty results.
+	return nil, nil
+}
+
+func highlightControlFlow(ctx context.Context, snapshot Snapshot, m *protocol.ColumnMapper, path []ast.Node) ([]protocol.Range, error) {
+	// Reverse walk the path till we get to the for loop.
+	var loop ast.Node
+Outer:
+	for _, n := range path {
+		switch n.(type) {
+		case *ast.ForStmt, *ast.RangeStmt:
+			loop = n
+			break Outer
 		}
-		if n.Name != id.Name {
-			return true
+	}
+	if loop == nil {
+		// Cursor is not in a for loop.
+		return nil, nil
+	}
+
+	var result []protocol.Range
+
+	// Add the for statement.
+	forStmt, err := posToRange(ctx, snapshot.View(), m, loop.Pos(), loop.Pos()+3)
+	if err != nil {
+		return nil, err
+	}
+	rng, err := forStmt.Range()
+	if err != nil {
+		return nil, err
+	}
+	result = append(result, rng)
+
+	ast.Inspect(loop, func(n ast.Node) bool {
+		// Don't traverse any other for loops.
+		switch n.(type) {
+		case *ast.ForStmt, *ast.RangeStmt:
+			return loop == n
 		}
-		if n.Obj != id.Obj {
-			return true
-		}
-		nodeObj := pkg.GetTypesInfo().ObjectOf(n)
-		if nodeObj != idObj {
-			return false
-		}
-		if rng, err := nodeToProtocolRange(ctx, snapshot.View(), m, n); err == nil {
+
+		if n, ok := n.(*ast.BranchStmt); ok {
+			// Add all branch statements in same scope as the identified one.
+			rng, err := nodeToProtocolRange(ctx, snapshot.View(), m, n)
+			if err != nil {
+				log.Error(ctx, "Error getting range for node", err)
+				return false
+			}
 			result = append(result, rng)
 		}
 		return true
 	})
 	return result, nil
 }
+
+func highlightIdentifiers(ctx context.Context, snapshot Snapshot, m *protocol.ColumnMapper, path []ast.Node, pkg Package) ([]protocol.Range, error) {
+	var result []protocol.Range
+	id, ok := path[0].(*ast.Ident)
+	if !ok {
+		return nil, errors.Errorf("highlightIdentifiers called with an ast.Node of type %T", id)
+	}
+
+	idObj := pkg.GetTypesInfo().ObjectOf(id)
+	ast.Inspect(path[len(path)-1], func(node ast.Node) bool {
+		n, ok := node.(*ast.Ident)
+		if !ok {
+			return true
+		}
+		if n.Name != id.Name || n.Obj != id.Obj {
+			return false
+		}
+		if nObj := pkg.GetTypesInfo().ObjectOf(n); nObj != idObj {
+			return false
+		}
+
+		if rng, err := nodeToProtocolRange(ctx, snapshot.View(), m, n); err == nil {
+			result = append(result, rng)
+		} else {
+			log.Error(ctx, "Error getting range for node", err)
+		}
+		return false
+	})
+	return result, nil
+}

internal/lsp/testdata/highlights/highlights.go

@@ -36,3 +36,35 @@ func toProtocolHighlight(rngs []protocol.Range) []protocol.DocumentHighlight { /
 	}
 	return result
 }
+
+func testForLoops() {
+	for i := 0; i < 10; i++ { //@mark(forDecl1, "for"),highlight(forDecl1, forDecl1, brk1, cont1)
+		if i > 8 {
+			break //@mark(brk1, "break"),highlight(brk1, forDecl1, brk1, cont1)
+		}
+		if i < 2 {
+			for j := 1; j < 10; j++ { //@mark(forDecl2, "for"),highlight(forDecl2, forDecl2, cont2)
+				if j < 3 {
+					for k := 1; k < 10; k++ { //@mark(forDecl3, "for"),highlight(forDecl3, forDecl3, cont3)
+						if k < 3 {
+							continue //@mark(cont3, "continue"),highlight(cont3, forDecl3, cont3)
+						}
+					}
+					continue //@mark(cont2, "continue"),highlight(cont2, forDecl2, cont2)
+				}
+			}
+			continue //@mark(cont1, "continue"),highlight(cont1, forDecl1, brk1, cont1)
+		}
+	}
+
+	arr := []int{}
+
+	for i := range arr { //@mark(forDecl4, "for"),highlight(forDecl4, forDecl4, brk4, cont4)
+		if i > 8 {
+			break //@mark(brk4, "break"),highlight(brk4, forDecl4, brk4, cont4)
+		}
+		if i < 4 {
+			continue //@mark(cont4, "continue"),highlight(cont4, forDecl4, brk4, cont4)
+		}
+	}
+}

internal/lsp/testdata/summary.txt.golden

@@ -13,7 +13,7 @@ ImportCount = 7
 SuggestedFixCount = 1
 DefinitionsCount = 37
 TypeDefinitionsCount = 2
-HighlightsCount = 12
+HighlightsCount = 22
 ReferencesCount = 7
 RenamesCount = 22
 PrepareRenamesCount = 8