gopls/internal/analysis/fillswitch/fillswitch.go

// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package fillswitch

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/token"
	"go/types"

	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/internal/refactor"
	"golang.org/x/tools/internal/typesinternal"
)

// Diagnose computes diagnostics for switch statements with missing cases
// overlapping with the provided start and end position of file f.
//
// If either start or end is invalid, the entire file is inspected.
func Diagnose(f *ast.File, start, end token.Pos, pkg *types.Package, info *types.Info) []analysis.Diagnostic {
	var diags []analysis.Diagnostic
	qual := typesinternal.FileQualifier(f, pkg)
	ast.Inspect(f, func(n ast.Node) bool {
		if n == nil {
			return true // pop
		}
		if start.IsValid() && n.End() < start ||
			end.IsValid() && n.Pos() > end {
			return false // skip non-overlapping subtree
		}
		var fix *analysis.SuggestedFix
		switch n := n.(type) {
		case *ast.SwitchStmt:
			fix = suggestedFixSwitch(f, n, pkg, info, qual)
		case *ast.TypeSwitchStmt:
			fix = suggestedFixTypeSwitch(f, n, pkg, info, qual)
		}
		if fix != nil {
			diags = append(diags, analysis.Diagnostic{
				Message:        fix.Message,
				Pos:            n.Pos(),
				End:            n.Pos() + token.Pos(len("switch")),
				SuggestedFixes: []analysis.SuggestedFix{*fix},
			})
		}
		return true
	})

	return diags
}

func suggestedFixTypeSwitch(f *ast.File, stmt *ast.TypeSwitchStmt, pkg *types.Package, info *types.Info, qual types.Qualifier) *analysis.SuggestedFix {
	if hasDefaultCase(stmt.Body) {
		return nil
	}

	namedType := namedTypeFromTypeSwitch(stmt, info)
	if namedType == nil {
		return nil
	}

	existingCases := caseTypes(stmt.Body, info)
	// Gather accessible package-level concrete types
	// that implement the switch interface type.
	scope := namedType.Obj().Pkg().Scope()
	var buf bytes.Buffer
	for _, name := range scope.Names() {
		obj := scope.Lookup(name)
		if tname, ok := obj.(*types.TypeName); !ok || tname.IsAlias() {
			continue // not a defined type
		}

		if types.IsInterface(obj.Type()) {
			continue
		}

		samePkg := obj.Pkg() == pkg
		if !samePkg && !obj.Exported() {
			continue // inaccessible
		}

		var key caseType
		if types.AssignableTo(obj.Type(), namedType.Obj().Type()) {
			key.named = obj.Type().(*types.Named)
		} else if ptr := types.NewPointer(obj.Type()); types.AssignableTo(ptr, namedType.Obj().Type()) {
			key.named = obj.Type().(*types.Named)
			key.ptr = true
		}

		if key.named != nil {
			if existingCases[key] {
				continue
			}

			if buf.Len() > 0 {
				buf.WriteString("\t")
			}

			buf.WriteString("case ")
			if key.ptr {
				buf.WriteByte('*')
			}

			if p := key.named.Obj().Pkg(); p != pkg {
				if name := qual(p); name != "" { // not a dot import
					buf.WriteString(name)
					buf.WriteByte('.')
				}
			}
			buf.WriteString(key.named.Obj().Name())
			buf.WriteString(":\n")
		}
	}

	if buf.Len() == 0 {
		return nil
	}

	rBracePos := stmt.Body.Rbrace
	var edits []analysis.TextEdit
	switch assign := stmt.Assign.(type) {
	case *ast.AssignStmt:
		edits = addDefaultCase(f, info, &buf, namedType, assign.Lhs[0], rBracePos)
	case *ast.ExprStmt:
		if assert, ok := assign.X.(*ast.TypeAssertExpr); ok {
			edits = addDefaultCase(f, info, &buf, namedType, assert.X, rBracePos)
		}
	}
	edits = append(edits, analysis.TextEdit{
		Pos:     rBracePos,
		End:     rBracePos,
		NewText: buf.Bytes(),
	})

	return &analysis.SuggestedFix{
		Message:   "Add cases for " + types.TypeString(namedType, qual),
		TextEdits: edits,
	}
}

func suggestedFixSwitch(f *ast.File, stmt *ast.SwitchStmt, pkg *types.Package, info *types.Info, qual types.Qualifier) *analysis.SuggestedFix {
	if hasDefaultCase(stmt.Body) {
		return nil
	}

	namedType, ok := info.TypeOf(stmt.Tag).(*types.Named)
	if !ok {
		return nil
	}

	existingCases := caseConsts(stmt.Body, info)
	// Gather accessible named constants of the same type as the switch value.
	scope := namedType.Obj().Pkg().Scope()
	var buf bytes.Buffer
	for _, name := range scope.Names() {
		obj := scope.Lookup(name)
		if c, ok := obj.(*types.Const); ok &&
			(obj.Pkg() == pkg || obj.Exported()) && // accessible
			types.Identical(obj.Type(), namedType.Obj().Type()) &&
			!existingCases[c] {

			if buf.Len() > 0 {
				buf.WriteString("\t")
			}

			buf.WriteString("case ")
			if c.Pkg() != pkg {
				if name := qual(c.Pkg()); name != "" { // not a dot import
					buf.WriteString(name)
					buf.WriteByte('.')
				}
			}
			buf.WriteString(c.Name())
			buf.WriteString(":\n")
		}
	}

	if buf.Len() == 0 {
		return nil
	}

	rBracePos := stmt.Body.Rbrace
	edits := addDefaultCase(f, info, &buf, namedType, stmt.Tag, rBracePos)
	edits = append(edits, analysis.TextEdit{
		Pos:     rBracePos,
		End:     rBracePos,
		NewText: buf.Bytes(),
	})

	return &analysis.SuggestedFix{
		Message:   "Add cases for " + types.TypeString(namedType, qual),
		TextEdits: edits,
	}
}

// addDefaultCase writes a default switch case to buf containing an
// "unexpected" panic message. It returns edits for any added imports.
func addDefaultCase(f *ast.File, info *types.Info, buf *bytes.Buffer, named *types.Named, expr ast.Expr, rBracePos token.Pos) (importEdits []analysis.TextEdit) {
	var dottedBuf bytes.Buffer
	// writeDotted emits a dotted path a.b.c.
	var writeDotted func(e ast.Expr) bool
	writeDotted = func(e ast.Expr) bool {
		switch e := e.(type) {
		case *ast.SelectorExpr:
			if !writeDotted(e.X) {
				return false
			}
			dottedBuf.WriteByte('.')
			dottedBuf.WriteString(e.Sel.Name)
			return true
		case *ast.Ident:
			dottedBuf.WriteString(e.Name)
			return true
		}
		return false
	}

	buf.WriteString("\tdefault:\n")
	typeName := fmt.Sprintf("%s.%s", named.Obj().Pkg().Name(), named.Obj().Name())
	if writeDotted(expr) {
		// Switch tag expression is a dotted path.
		// It is safe to re-evaluate it in the default case.
		var fmtPrefix string
		fmtPrefix, importEdits = refactor.AddImport(info, f, "fmt", "fmt", "Sprintf", rBracePos)
		format := fmt.Sprintf("unexpected %s: %%#v", typeName)
		fmt.Fprintf(buf, "\t\tpanic(%sSprintf(%q, %s))\n\t", fmtPrefix, format, dottedBuf.String())
		return importEdits
	} else {
		// Emit simpler message, without re-evaluating tag expression.
		fmt.Fprintf(buf, "\t\tpanic(%q)\n\t", "unexpected "+typeName)
		return nil
	}
}

func namedTypeFromTypeSwitch(stmt *ast.TypeSwitchStmt, info *types.Info) *types.Named {
	switch assign := stmt.Assign.(type) {
	case *ast.ExprStmt:
		if typ, ok := assign.X.(*ast.TypeAssertExpr); ok {
			if named, ok := info.TypeOf(typ.X).(*types.Named); ok {
				return named
			}
		}

	case *ast.AssignStmt:
		if typ, ok := assign.Rhs[0].(*ast.TypeAssertExpr); ok {
			if named, ok := info.TypeOf(typ.X).(*types.Named); ok {
				return named
			}
		}
	}

	return nil
}

func hasDefaultCase(body *ast.BlockStmt) bool {
	for _, clause := range body.List {
		if len(clause.(*ast.CaseClause).List) == 0 {
			return true
		}
	}

	return false
}

func caseConsts(body *ast.BlockStmt, info *types.Info) map[*types.Const]bool {
	out := map[*types.Const]bool{}
	for _, stmt := range body.List {
		for _, e := range stmt.(*ast.CaseClause).List {
			if info.Types[e].Value == nil {
				continue // not a constant
			}

			if sel, ok := e.(*ast.SelectorExpr); ok {
				e = sel.Sel // replace pkg.C with C
			}

			if e, ok := e.(*ast.Ident); ok {
				if c, ok := info.Uses[e].(*types.Const); ok {
					out[c] = true
				}
			}
		}
	}

	return out
}

type caseType struct {
	named *types.Named
	ptr   bool
}

func caseTypes(body *ast.BlockStmt, info *types.Info) map[caseType]bool {
	out := map[caseType]bool{}
	for _, stmt := range body.List {
		for _, e := range stmt.(*ast.CaseClause).List {
			if tv, ok := info.Types[e]; ok && tv.IsType() {
				t := tv.Type
				ptr := false
				if p, ok := t.(*types.Pointer); ok {
					t = p.Elem()
					ptr = true
				}

				if named, ok := t.(*types.Named); ok {
					out[caseType{named, ptr}] = true
				}
			}
		}
	}

	return out
}