internal/lsp/cache: improve completion in init statements

In cases like:

  if i := foo<>

we get an *ast.BadExpr because the parser is expecting the condition
expression, but "i := foo" is not a valid expression. Now we move the
statement into the "init" field and add a dummy "cond" expression.

We also needed a slight tweak to our missing curly brace fix. Now we
insert an extra semicolon in cases like:

for i := 0; i < foo

yielding

for i := 0; i < foo;{}

The parser doesn't like having only two clauses in the three clause
"for" statement.

Updates golang/go#31687.

Change-Id: I12d51e0d8af03436741227753f8e71452a463b05
Reviewed-on: https://go-review.googlesource.com/c/tools/+/216483
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

@@ -240,6 +240,15 @@ func fixAST(ctx context.Context, n ast.Node, tok *token.File, src []byte) error
 				err = fixAST(ctx, parent, tok, src)
 				return false
 			}
+
+			// Fix cases where parser interprets if/for/switch "init"
+			// statement as "cond" expression, e.g.:
+			//
+			//   // "i := foo" is init statement, not condition.
+			//   for i := foo
+			//
+			fixInitStmt(n, parent, tok, src)
+
 			return false
 		case *ast.SelectorExpr:
 			// Fix cases where a keyword prefix results in a phantom "_" selector, e.g.:
@@ -355,10 +364,26 @@ func fixMissingCurlies(f *ast.File, b *ast.BlockStmt, parent ast.Node, tok *toke
 		return nil
 	}
 
-	// Insert "{}" at insertPos.
 	var buf bytes.Buffer
-	buf.Grow(len(src) + 2)
+	buf.Grow(len(src) + 3)
 	buf.Write(src[:tok.Offset(insertPos)])
+
+	// Detect if we need to insert a semicolon to fix "for" loop situations like:
+	//
+	//   for i := foo(); foo<>
+	//
+	// Just adding curlies is not sufficient to make things parse well.
+	if fs, ok := parent.(*ast.ForStmt); ok {
+		if _, ok := fs.Cond.(*ast.BadExpr); !ok {
+			if xs, ok := fs.Post.(*ast.ExprStmt); ok {
+				if _, ok := xs.X.(*ast.BadExpr); ok {
+					buf.WriteByte(';')
+				}
+			}
+		}
+	}
+
+	// Insert "{}" at insertPos.
 	buf.WriteByte('{')
 	buf.WriteByte('}')
 	buf.Write(src[tok.Offset(insertPos):])
@@ -402,6 +427,48 @@ func isPhantomUnderscore(id *ast.Ident, tok *token.File, src []byte) bool {
 	return len(src) <= offset || src[offset] != '_'
 }
 
+// fixInitStmt fixes cases where the parser misinterprets an
+// if/for/switch "init" statement as the "cond" conditional. In cases
+// like "if i := 0" the user hasn't typed the semicolon yet so the
+// parser is looking for the conditional expression. However, "i := 0"
+// are not valid expressions, so we get a BadExpr.
+func fixInitStmt(bad *ast.BadExpr, parent ast.Node, tok *token.File, src []byte) {
+	if !bad.Pos().IsValid() || !bad.End().IsValid() {
+		return
+	}
+
+	// Try to extract a statement from the BadExpr.
+	stmtBytes := src[tok.Offset(bad.Pos()) : tok.Offset(bad.End()-1)+1]
+	stmt, err := parseStmt(bad.Pos(), stmtBytes)
+	if err != nil {
+		return
+	}
+
+	// If the parent statement doesn't already have an "init" statement,
+	// move the extracted statement into the "init" field and insert a
+	// dummy expression into the required "cond" field.
+	switch p := parent.(type) {
+	case *ast.IfStmt:
+		if p.Init != nil {
+			return
+		}
+		p.Init = stmt
+		p.Cond = &ast.Ident{Name: "_"}
+	case *ast.ForStmt:
+		if p.Init != nil {
+			return
+		}
+		p.Init = stmt
+		p.Cond = &ast.Ident{Name: "_"}
+	case *ast.SwitchStmt:
+		if p.Init != nil {
+			return
+		}
+		p.Init = stmt
+		p.Tag = nil
+	}
+}
+
 // readKeyword reads the keyword starting at pos, if any.
 func readKeyword(pos token.Pos, tok *token.File, src []byte) string {
 	var kwBytes []byte

internal/lsp/testdata/lsp/primarymod/danglingstmt/dangling_for_init.go

@@ -0,0 +1,9 @@
+package danglingstmt
+
+func _() {
+	for i := bar //@rank(" //", danglingBar2)
+}
+
+func bar2() int { //@item(danglingBar2, "bar2", "func() int", "func")
+	return 0
+}

internal/lsp/testdata/lsp/primarymod/danglingstmt/dangling_for_init_cond.go

@@ -0,0 +1,9 @@
+package danglingstmt
+
+func _() {
+	for i := bar3(); i > bar //@rank(" //", danglingBar3)
+}
+
+func bar3() int { //@item(danglingBar3, "bar3", "func() int", "func")
+	return 0
+}

internal/lsp/testdata/lsp/primarymod/danglingstmt/dangling_for_init_cond_post.go

@@ -0,0 +1,9 @@
+package danglingstmt
+
+func _() {
+	for i := bar4(); i > bar4(); i += bar //@rank(" //", danglingBar4)
+}
+
+func bar4() int { //@item(danglingBar4, "bar4", "func() int", "func")
+	return 0
+}

internal/lsp/testdata/lsp/primarymod/danglingstmt/dangling_if_init.go

@@ -0,0 +1,9 @@
+package danglingstmt
+
+func _() {
+	if i := foo //@rank(" //", danglingFoo2)
+}
+
+func foo2() bool { //@item(danglingFoo2, "foo2", "func() bool", "func")
+	return true
+}

internal/lsp/testdata/lsp/primarymod/danglingstmt/dangling_if_init_cond.go

@@ -0,0 +1,9 @@
+package danglingstmt
+
+func _() {
+	if i := 123; foo //@rank(" //", danglingFoo3)
+}
+
+func foo3() bool { //@item(danglingFoo3, "foo3", "func() bool", "func")
+	return true
+}

internal/lsp/testdata/lsp/primarymod/danglingstmt/dangling_switch_init.go

@@ -0,0 +1,9 @@
+package danglingstmt
+
+func _() {
+	switch i := baz //@rank(" //", danglingBaz)
+}
+
+func baz() int { //@item(danglingBaz, "baz", "func() int", "func")
+	return 0
+}

internal/lsp/testdata/lsp/primarymod/danglingstmt/dangling_switch_init_tag.go

@@ -0,0 +1,9 @@
+package danglingstmt
+
+func _() {
+	switch i := 0; baz //@rank(" //", danglingBaz2)
+}
+
+func baz2() int { //@item(danglingBaz2, "baz2", "func() int", "func")
+	return 0
+}

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

@@ -4,7 +4,7 @@ CompletionSnippetCount = 67
 UnimportedCompletionsCount = 11
 DeepCompletionsCount = 5
 FuzzyCompletionsCount = 8
-RankedCompletionsCount = 88
+RankedCompletionsCount = 95
 CaseSensitiveCompletionsCount = 4
 DiagnosticsCount = 38
 FoldingRangesCount = 2