internal/lsp/source: fix completion in empty switch statements

Now we properly offer "case" and "default" keyword completion in cases
like:

    switch {
      <>
    }

First I had to add an AST fix for empty switch statements to move the
closing brace down. For example, say the user is completing:

    switch {
    ca<>
    }

This gets parsed as:

    switch {
    }

Even if we manually recover the "ca" token, "<>" is not positioned
inside the switch statement anymore, so we don't know to offer "case"
and "default" candidates. To work around this, we move the closing
brace down one line yielding:

    switch {

    }

Second I had to add logic to manually extract the completion prefix
inside empty switch statements, and finally some logic to offer (only)
"case" and "default" candidates in empty switch statements.

Updates golang/go#34009.

Change-Id: I624f17da1c5e73faf91fe5f69e872d86f1cf5482
Reviewed-on: https://go-review.googlesource.com/c/tools/+/220579
Run-TryBot: Muir Manders <muir@mnd.rs>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/cache/parse.go

@@ -256,6 +256,16 @@ func fixAST(ctx context.Context, n ast.Node, tok *token.File, src []byte) error
 			//
 			fixPhantomSelector(n, tok, src)
 			return true
+
+		case *ast.BlockStmt:
+			switch parent.(type) {
+			case *ast.SwitchStmt, *ast.TypeSwitchStmt, *ast.SelectStmt:
+				// Adjust closing curly brace of empty switch/select
+				// statements so we can complete inside them.
+				fixEmptySwitch(n, tok, src)
+			}
+
+			return true
 		default:
 			return true
 		}
@@ -398,6 +408,51 @@ func fixMissingCurlies(f *ast.File, b *ast.BlockStmt, parent ast.Node, tok *toke
 	return buf.Bytes()
 }
 
+// fixEmptySwitch moves empty switch/select statements' closing curly
+// brace down one line. This allows us to properly detect incomplete
+// "case" and "default" keywords as inside the switch statement. For
+// example:
+//
+//   switch {
+//   def<>
+//   }
+//
+// gets parsed like:
+//
+//   switch {
+//   }
+//
+// Later we manually pull out the "def" token, but we need to detect
+// that our "<>" position is inside the switch block. To do that we
+// move the curly brace so it looks like:
+//
+//   switch {
+//
+//   }
+//
+func fixEmptySwitch(body *ast.BlockStmt, tok *token.File, src []byte) {
+	// We only care about empty switch statements.
+	if len(body.List) > 0 || !body.Rbrace.IsValid() {
+		return
+	}
+
+	// If the right brace is actually in the source code at the
+	// specified position, don't mess with it.
+	braceOffset := tok.Offset(body.Rbrace)
+	if braceOffset < len(src) && src[braceOffset] == '}' {
+		return
+	}
+
+	braceLine := tok.Line(body.Rbrace)
+	if braceLine >= tok.LineCount() {
+		// If we are the last line in the file, no need to fix anything.
+		return
+	}
+
+	// Move the right brace down one line.
+	body.Rbrace = tok.LineStart(braceLine + 1)
+}
+
 // fixDanglingSelector inserts real "_" selector expressions in place
 // of phantom "_" selectors. For example:
 //

internal/lsp/source/completion.go

@@ -501,26 +501,8 @@ func Completion(ctx context.Context, snapshot Snapshot, fh FileHandle, protoPos
 		c.deepState.maxDepth = -1
 	}
 
-	// Detect our surrounding identifier.
+	if surrounding := c.containingIdent(src); surrounding != nil {
-	switch leaf := path[0].(type) {
+		c.setSurrounding(surrounding)
-	case *ast.Ident:
-		// In the normal case, our leaf AST node is the identifier being completed.
-		c.setSurrounding(leaf)
-	case *ast.BadDecl:
-		// You don't get *ast.Idents at the file level, so look for bad
-		// decls and manually extract the surrounding token.
-		pos, _, lit := c.scanToken(ctx, src)
-		if pos.IsValid() {
-			c.setSurrounding(&ast.Ident{Name: lit, NamePos: pos})
-		}
-	default:
-		// Otherwise, manually extract the prefix if our containing token
-		// is a keyword. This improves completion after an "accidental
-		// keyword", e.g. completing to "variance" in "someFunc(var<>)".
-		pos, tkn, lit := c.scanToken(ctx, src)
-		if pos.IsValid() && tkn.IsKeyword() {
-			c.setSurrounding(&ast.Ident{Name: lit, NamePos: pos})
-		}
 	}
 
 	c.inference = expectedCandidate(c)
@@ -552,6 +534,12 @@ func Completion(ctx context.Context, snapshot Snapshot, fh FileHandle, protoPos
 	// contexts, as opposed to a single object.
 	c.addStatementCandidates()
 
+	if c.emptySwitchStmt() {
+		// Empty switch statements only admit "default" and "case" keywords.
+		c.addKeywordItems(map[string]bool{}, highScore, CASE, DEFAULT)
+		return c.items, c.getSurrounding(), nil
+	}
+
 	switch n := path[0].(type) {
 	case *ast.Ident:
 		// Is this the Sel part of a selector?
@@ -604,8 +592,43 @@ func Completion(ctx context.Context, snapshot Snapshot, fh FileHandle, protoPos
 	return c.items, c.getSurrounding(), nil
 }
 
+// containingIdent returns the *ast.Ident containing pos, if any. It
+// synthesizes an *ast.Ident to allow completion in the face of
+// certain syntax errors.
+func (c *completer) containingIdent(src []byte) *ast.Ident {
+	// In the normal case, our leaf AST node is the identifer being completed.
+	if ident, ok := c.path[0].(*ast.Ident); ok {
+		return ident
+	}
+
+	pos, tkn, lit := c.scanToken(src)
+	if !pos.IsValid() {
+		return nil
+	}
+
+	fakeIdent := &ast.Ident{Name: lit, NamePos: pos}
+
+	if _, isBadDecl := c.path[0].(*ast.BadDecl); isBadDecl {
+		// You don't get *ast.Idents at the file level, so look for bad
+		// decls and use the manually extracted token.
+		return fakeIdent
+	} else if c.emptySwitchStmt() {
+		// Only keywords are allowed in empty switch statements.
+		// *ast.Idents are not parsed, so we must use the manually
+		// extracted token.
+		return fakeIdent
+	} else if tkn.IsKeyword() {
+		// Otherwise, manually extract the prefix if our containing token
+		// is a keyword. This improves completion after an "accidental
+		// keyword", e.g. completing to "variance" in "someFunc(var<>)".
+		return fakeIdent
+	}
+
+	return nil
+}
+
 // scanToken scans pgh's contents for the token containing pos.
-func (c *completer) scanToken(ctx context.Context, contents []byte) (token.Pos, token.Token, string) {
+func (c *completer) scanToken(contents []byte) (token.Pos, token.Token, string) {
 	tok := c.snapshot.View().Session().Cache().FileSet().File(c.pos)
 
 	var s scanner.Scanner
@@ -635,6 +658,22 @@ func (c *completer) sortItems() {
 	})
 }
 
+// emptySwitchStmt reports whether pos is in an empty switch or select
+// statement.
+func (c *completer) emptySwitchStmt() bool {
+	block, ok := c.path[0].(*ast.BlockStmt)
+	if !ok || len(block.List) > 0 || len(c.path) == 1 {
+		return false
+	}
+
+	switch c.path[1].(type) {
+	case *ast.SwitchStmt, *ast.TypeSwitchStmt, *ast.SelectStmt:
+		return true
+	default:
+		return false
+	}
+}
+
 // populateCommentCompletions yields completions for an exported
 // variable immediately preceding comment.
 func (c *completer) populateCommentCompletions(comment *ast.CommentGroup) {

internal/lsp/source/completion_keywords.go

@@ -38,26 +38,6 @@ const (
 func (c *completer) addKeywordCompletions() {
 	seen := make(map[string]bool)
 
-	// addKeywords dedupes and adds completion items for the specified
-	// keywords with the specified score.
-	addKeywords := func(score float64, kws ...string) {
-		for _, kw := range kws {
-			if seen[kw] {
-				continue
-			}
-			seen[kw] = true
-
-			if matchScore := c.matcher.Score(kw); matchScore > 0 {
-				c.items = append(c.items, CompletionItem{
-					Label:      kw,
-					Kind:       protocol.KeywordCompletion,
-					InsertText: kw,
-					Score:      score * float64(matchScore),
-				})
-			}
-		}
-	}
-
 	if c.wantTypeName() {
 		// If we expect a type name, include "interface", "struct",
 		// "func", "chan", and "map".
@@ -71,15 +51,15 @@ func (c *completer) addKeywordCompletions() {
 			}
 		}
 
-		addKeywords(structIntf, STRUCT, INTERFACE)
+		c.addKeywordItems(seen, structIntf, STRUCT, INTERFACE)
-		addKeywords(funcChanMap, FUNC, CHAN, MAP)
+		c.addKeywordItems(seen, funcChanMap, FUNC, CHAN, MAP)
 	}
 
 	// If we are at the file scope, only offer decl keywords. We don't
 	// get *ast.Idents at the file scope because non-keyword identifiers
 	// turn into *ast.BadDecl, not *ast.Ident.
 	if len(c.path) == 1 || isASTFile(c.path[1]) {
-		addKeywords(stdScore, TYPE, CONST, VAR, FUNC, IMPORT)
+		c.addKeywordItems(seen, stdScore, TYPE, CONST, VAR, FUNC, IMPORT)
 		return
 	} else if _, ok := c.path[0].(*ast.Ident); !ok {
 		// Otherwise only offer keywords if the client is completing an identifier.
@@ -90,7 +70,7 @@ func (c *completer) addKeywordCompletions() {
 		// Offer "range" if we are in ast.ForStmt.Init. This is what the
 		// AST looks like before "range" is typed, e.g. "for i := r<>".
 		if loop, ok := c.path[2].(*ast.ForStmt); ok && nodeContains(loop.Init, c.pos) {
-			addKeywords(stdScore, RANGE)
+			c.addKeywordItems(seen, stdScore, RANGE)
 		}
 	}
 
@@ -109,7 +89,7 @@ func (c *completer) addKeywordCompletions() {
 		case *ast.CaseClause:
 			// only recommend "fallthrough" and "break" within the bodies of a case clause
 			if c.pos > node.Colon {
-				addKeywords(stdScore, BREAK)
+				c.addKeywordItems(seen, stdScore, BREAK)
 				// "fallthrough" is only valid in switch statements.
 				// A case clause is always nested within a block statement in a switch statement,
 				// that block statement is nested within either a TypeSwitchStmt or a SwitchStmt.
@@ -117,23 +97,42 @@ func (c *completer) addKeywordCompletions() {
 					continue
 				}
 				if _, ok := path[i+2].(*ast.SwitchStmt); ok {
-					addKeywords(stdScore, FALLTHROUGH)
+					c.addKeywordItems(seen, stdScore, FALLTHROUGH)
 				}
 			}
 		case *ast.CommClause:
 			if c.pos > node.Colon {
-				addKeywords(stdScore, BREAK)
+				c.addKeywordItems(seen, stdScore, BREAK)
 			}
 		case *ast.TypeSwitchStmt, *ast.SelectStmt, *ast.SwitchStmt:
-			addKeywords(stdScore, CASE, DEFAULT)
+			c.addKeywordItems(seen, stdScore, CASE, DEFAULT)
 		case *ast.ForStmt:
-			addKeywords(stdScore, BREAK, CONTINUE)
+			c.addKeywordItems(seen, stdScore, BREAK, CONTINUE)
 		// This is a bit weak, functions allow for many keywords
 		case *ast.FuncDecl:
 			if node.Body != nil && c.pos > node.Body.Lbrace {
-				addKeywords(stdScore, DEFER, RETURN, FOR, GO, SWITCH, SELECT, IF, ELSE, VAR, CONST, GOTO, TYPE)
+				c.addKeywordItems(seen, stdScore, DEFER, RETURN, FOR, GO, SWITCH, SELECT, IF, ELSE, VAR, CONST, GOTO, TYPE)
 			}
 		}
 	}
-
+}
+
+// addKeywordItems dedupes and adds completion items for the specified
+// keywords with the specified score.
+func (c *completer) addKeywordItems(seen map[string]bool, score float64, kws ...string) {
+	for _, kw := range kws {
+		if seen[kw] {
+			continue
+		}
+		seen[kw] = true
+
+		if matchScore := c.matcher.Score(kw); matchScore > 0 {
+			c.items = append(c.items, CompletionItem{
+				Label:      kw,
+				Kind:       protocol.KeywordCompletion,
+				InsertText: kw,
+				Score:      score * float64(matchScore),
+			})
+		}
+	}
 }

internal/lsp/testdata/lsp/primarymod/keywords/empty_select.go

@@ -0,0 +1,7 @@
+package keywords
+
+func _() {
+	select {
+		c //@complete(" //", case)
+	}
+}

internal/lsp/testdata/lsp/primarymod/keywords/empty_switch.go

@@ -0,0 +1,11 @@
+package keywords
+
+func _() {
+	switch {
+		//@complete("", case, default)
+	}
+
+	switch test.(type) {
+		d //@complete(" //", default)
+	}
+}

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

@@ -1,6 +1,6 @@
 -- summary --
 CodeLensCount = 0
-CompletionsCount = 231
+CompletionsCount = 234
 CompletionSnippetCount = 74
 UnimportedCompletionsCount = 11
 DeepCompletionsCount = 5