internal/lsp: fix infinite recursion while fixing AST

We were recursing infinitely in cases like this:

switch true {
case true:
	go foo.F<>
}

There were three things that came together to cause this:
1. We recently starting recursively fixing broken go/defer statements.
2. In this case we were failing to swap in the correct ast.Node in for
   the *ast.BadStmt because we were only looking
   for *ast.BlockStmt (and *ast.CaseStmt has no block).
3. After 2), we weren't returning an error so the fix() code thought
   it should recurse.

Fix 2) by using reflection to swap AST nodes in a generic way. Perhaps
a bit overkill in this case, but I happened to have already written
this for an upcoming change, so I just pulled it in to fix this bug.

Fix 3) by returning an error if we fail to swap the AST nodes.

Fixes golang/go#34353.

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

internal/lsp/cache/parse.go

diff --git a/internal/lsp/cache/parse.go b/internal/lsp/cache/parse.go
index 537460a..724746a 100644
--- a/internal/lsp/cache/parse.go
+++ b/internal/lsp/cache/parse.go
@@ -402,15 +402,11 @@
 	case *ast.GoStmt:
 		stmt.Call = call
 	}
-	switch parent := parent.(type) {
-	case *ast.BlockStmt:
-		for i, s := range parent.List {
-			if s == bad {
-				parent.List[i] = stmt
-				break
-			}
-		}
+
+	if !replaceNode(parent, bad, stmt) {
+		return errors.Errorf("couldn't replace %T in %T", stmt, parent)
 	}
+
 	return nil
 }
 
@@ -444,3 +440,56 @@
 		return true
 	})
 }
+
+// replaceNode updates parent's child oldChild to be newChild. It
+// retuns whether it replaced successfully.
+func replaceNode(parent, oldChild, newChild ast.Node) bool {
+	if parent == nil || oldChild == nil || newChild == nil {
+		return false
+	}
+
+	parentVal := reflect.ValueOf(parent).Elem()
+	if parentVal.Kind() != reflect.Struct {
+		return false
+	}
+
+	newChildVal := reflect.ValueOf(newChild)
+
+	tryReplace := func(v reflect.Value) bool {
+		if !v.CanSet() || !v.CanInterface() {
+			return false
+		}
+
+		// If the existing value is oldChild, we found our child. Make
+		// sure our newChild is assignable and then make the swap.
+		if v.Interface() == oldChild && newChildVal.Type().AssignableTo(v.Type()) {
+			v.Set(newChildVal)
+			return true
+		}
+
+		return false
+	}
+
+	// Loop over parent's struct fields.
+	for i := 0; i < parentVal.NumField(); i++ {
+		f := parentVal.Field(i)
+
+		switch f.Kind() {
+		// Check interface and pointer fields.
+		case reflect.Interface, reflect.Ptr:
+			if tryReplace(f) {
+				return true
+			}
+
+		// Search through any slice fields.
+		case reflect.Slice:
+			for i := 0; i < f.Len(); i++ {
+				if tryReplace(f.Index(i)) {
+					return true
+				}
+			}
+		}
+	}
+
+	return false
+}

internal/lsp/testdata/badstmt/badstmt.go.in

diff --git a/internal/lsp/testdata/badstmt/badstmt.go.in b/internal/lsp/testdata/badstmt/badstmt.go.in
index 3aae7db..05b2c9a 100644
--- a/internal/lsp/testdata/badstmt/badstmt.go.in
+++ b/internal/lsp/testdata/badstmt/badstmt.go.in
@@ -11,6 +11,13 @@
 }
 
 func _() {
+	switch true {
+	case true:
+		go foo.F //@complete(" //", Foo)
+	}
+}
+
+func _() {
 	defer func() {
 		foo.F //@complete(" //", Foo),snippet(" //", Foo, "Foo()", "Foo()")
 		foo. //@complete(" //", Foo, IntFoo, StructFoo)

internal/lsp/tests/tests.go

diff --git a/internal/lsp/tests/tests.go b/internal/lsp/tests/tests.go
index 97e5709..85e54fe 100644
--- a/internal/lsp/tests/tests.go
+++ b/internal/lsp/tests/tests.go
@@ -30,7 +30,7 @@
 // We hardcode the expected number of test cases to ensure that all tests
 // are being executed. If a test is added, this number must be changed.
 const (
-	ExpectedCompletionsCount       = 164
+	ExpectedCompletionsCount       = 165
 	ExpectedCompletionSnippetCount = 35
 	ExpectedDiagnosticsCount       = 21
 	ExpectedFormatCount            = 6