internal/lsp/source: support highlight of switch statements

Placing the cursor on a switch statement or corresponding break and call
document.Highlight will now highlight them.

Fixes golang/go#39275

Change-Id: Ib7e3ba0c6e78141ed3dd37cfd3b72b567b857247
Reviewed-on: https://go-review.googlesource.com/c/tools/+/238478
Run-TryBot: Daniel Martí <mvdan@mvdan.cc>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/source/highlight.go

@@ -90,6 +90,8 @@ func highlightPath(pkg Package, path []ast.Node) (map[posRange]struct{}, error)
 		highlightIdentifiers(pkg, path, result)
 	case *ast.ForStmt, *ast.RangeStmt:
 		highlightLoopControlFlow(path, result)
+	case *ast.SwitchStmt:
+		highlightSwitchFlow(path, result)
 	case *ast.BranchStmt:
 		// BREAK can exit a loop, switch or select, while CONTINUE exit a loop so
 		// these need to be handled separately. They can also be embedded in any
@@ -235,8 +237,11 @@ func highlightUnlabeledBreakFlow(path []ast.Node, result map[posRange]struct{})
 		case *ast.ForStmt, *ast.RangeStmt:
 			highlightLoopControlFlow(path, result)
 			return // only highlight the innermost statement
-		case *ast.SelectStmt, *ast.SwitchStmt:
-			// TODO: add highlight when breaking a select or switch.
+		case *ast.SwitchStmt:
+			highlightSwitchFlow(path, result)
+			return
+		case *ast.SelectStmt:
+			// TODO: add highlight when breaking a select.
 			return
 		}
 	}
@@ -254,11 +259,12 @@ func highlightLabeledFlow(node *ast.BranchStmt, result map[posRange]struct{}) {
 	switch label.Stmt.(type) {
 	case *ast.ForStmt, *ast.RangeStmt:
 		highlightLoopControlFlow([]ast.Node{label.Stmt, label}, result)
+	case *ast.SwitchStmt:
+		highlightSwitchFlow([]ast.Node{label.Stmt, label}, result)
 	}
 }
 
-func highlightLoopControlFlow(path []ast.Node, result map[posRange]struct{}) {
-	labelFor := func(path []ast.Node) *ast.Ident {
+func labelFor(path []ast.Node) *ast.Ident {
 	if len(path) > 1 {
 		if n, ok := path[1].(*ast.LabeledStmt); ok {
 			return n.Label
@@ -267,6 +273,7 @@ func highlightLoopControlFlow(path []ast.Node, result map[posRange]struct{}) {
 	return nil
 }
 
+func highlightLoopControlFlow(path []ast.Node, result map[posRange]struct{}) {
 	var loop ast.Node
 	var loopLabel *ast.Ident
 	stmtLabel := labelFor(path)
@@ -344,6 +351,74 @@ Outer:
 	})
 }
 
+func highlightSwitchFlow(path []ast.Node, result map[posRange]struct{}) {
+	var switchNode ast.Node
+	var switchNodeLabel *ast.Ident
+	stmtLabel := labelFor(path)
+Outer:
+	// Reverse walk the path till we get to the switch statement.
+	for i := range path {
+		switch n := path[i].(type) {
+		case *ast.SwitchStmt:
+			switchNodeLabel = labelFor(path[i:])
+			if stmtLabel == nil || switchNodeLabel == stmtLabel {
+				switchNode = n
+				break Outer
+			}
+		}
+	}
+	// Cursor is not in a switch statement
+	if switchNode == nil {
+		return
+	}
+
+	// Add the switch statement.
+	rng := posRange{
+		start: switchNode.Pos(),
+		end:   switchNode.Pos() + token.Pos(len("switch")),
+	}
+	result[rng] = struct{}{}
+
+	// Traverse AST to find break statements within the same switch.
+	ast.Inspect(switchNode, func(n ast.Node) bool {
+		switch n.(type) {
+		case *ast.SwitchStmt:
+			return switchNode == n
+		case *ast.ForStmt, *ast.RangeStmt, *ast.SelectStmt:
+			return false
+		}
+
+		b, ok := n.(*ast.BranchStmt)
+		if !ok || b.Tok != token.BREAK {
+			return true
+		}
+
+		if b.Label == nil || labelDecl(b.Label) == switchNodeLabel {
+			result[posRange{start: b.Pos(), end: b.End()}] = struct{}{}
+		}
+		return true
+	})
+
+	// We don't need to check other switches if we aren't looking for labeled statements.
+	if switchNodeLabel == nil {
+		return
+	}
+
+	// Find labeled break statements in any switch
+	ast.Inspect(switchNode, func(n ast.Node) bool {
+		b, ok := n.(*ast.BranchStmt)
+		if !ok || b.Tok != token.BREAK {
+			return true
+		}
+
+		if b.Label != nil && labelDecl(b.Label) == switchNodeLabel {
+			result[posRange{start: b.Pos(), end: b.End()}] = struct{}{}
+		}
+
+		return true
+	})
+}
+
 func labelDecl(n *ast.Ident) *ast.Ident {
 	if n == nil {
 		return nil

internal/lsp/testdata/lsp/primarymod/highlights/highlights.go

@@ -80,9 +80,9 @@ Outer:
 			if i == 1 {
 				break Outer //@mark(brk6, "break Outer"),highlight(brk6, forDecl5, brk5, brk6, brk8)
 			}
-			switch i { //@mark(switch1, "switch"),highlight(switch1)
+			switch i { //@mark(switch1, "switch"),highlight(switch1, switch1, brk7)
 			case 5:
-				break //@mark(brk7, "break"),highlight(brk7)
+				break //@mark(brk7, "break"),highlight(brk7, switch1, brk7)
 			case 6:
 				continue //@mark(cont5, "continue"),highlight(cont5, forDecl6, cont5)
 			case 7:
@@ -92,6 +92,32 @@ Outer:
 	}
 }
 
+func testSwitch() {
+	var i, j int
+
+L1:
+	for { //@mark(forDecl7, "for"),highlight(forDecl7, forDecl7, brk10, cont6)
+	L2:
+		switch i { //@mark(switch2, "switch"),highlight(switch2, switch2, brk11, brk12, brk13)
+		case 1:
+			switch j { //@mark(switch3, "switch"),highlight(switch3, switch3, brk9)
+			case 1:
+				break //@mark(brk9, "break"),highlight(brk9, switch3, brk9)
+			case 2:
+				break L1 //@mark(brk10, "break L1"),highlight(brk10, forDecl7, brk10, cont6)
+			case 3:
+				break L2 //@mark(brk11, "break L2"),highlight(brk11, switch2, brk11, brk12, brk13)
+			default:
+				continue //@mark(cont6, "continue"),highlight(cont6, forDecl7, brk10, cont6)
+			}
+		case 2:
+			break //@mark(brk12, "break"),highlight(brk12, switch2, brk11, brk12, brk13)
+		default:
+			break L2 //@mark(brk13, "break L2"),highlight(brk13, switch2, brk11, brk12, brk13)
+		}
+	}
+}
+
 func testReturn() bool { //@mark(func1, "func"),mark(bool1, "bool"),highlight(func1, func1, fullRet11, fullRet12),highlight(bool1, bool1, false1, bool2, true1)
 	if 1 < 2 {
 		return false //@mark(ret11, "return"),mark(fullRet11, "return false"),mark(false1, "false"),highlight(ret11, func1, fullRet11, fullRet12)

internal/lsp/testdata/lsp/summary.txt.golden

@@ -14,7 +14,7 @@ ImportCount = 8
 SuggestedFixCount = 12
 DefinitionsCount = 53
 TypeDefinitionsCount = 2
-HighlightsCount = 60
+HighlightsCount = 69
 ReferencesCount = 11
 RenamesCount = 25
 PrepareRenamesCount = 7