go/internal/lsp/source/code_lens.go

// Copyright 2020 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 source

import (
	"context"
	"go/ast"
	"go/token"
	"go/types"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/tools/internal/lsp/protocol"
)

// CodeLens computes code lens for Go source code.
func CodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]protocol.CodeLens, error) {
	f, _, m, _, err := snapshot.View().Session().Cache().ParseGoHandle(fh, ParseFull).Parse(ctx)
	if err != nil {
		return nil, err
	}

	var codeLens []protocol.CodeLens

	if snapshot.View().Options().EnabledCodeLens[CommandGenerate] {
		ggcl, err := goGenerateCodeLens(ctx, snapshot, fh, f, m)
		if err != nil {
			return nil, err
		}
		codeLens = append(codeLens, ggcl...)
	}

	if snapshot.View().Options().EnabledCodeLens[CommandTest] {
		rtcl, err := runTestCodeLens(ctx, snapshot, fh, f, m)
		if err != nil {
			return nil, err
		}
		codeLens = append(codeLens, rtcl...)
	}

	return codeLens, nil
}

var testMatcher = regexp.MustCompile("^Test[^a-z]")
var benchMatcher = regexp.MustCompile("^Benchmark[^a-z]")

func runTestCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle, f *ast.File, m *protocol.ColumnMapper) ([]protocol.CodeLens, error) {
	codeLens := make([]protocol.CodeLens, 0)

	pkg, _, err := getParsedFile(ctx, snapshot, fh, WidestPackageHandle)
	if err != nil {
		return nil, err
	}

	if !strings.HasSuffix(fh.Identity().URI.Filename(), "_test.go") {
		return nil, nil
	}

	for _, d := range f.Decls {
		fn, ok := d.(*ast.FuncDecl)
		if !ok {
			continue
		}

		if isTestFunc(fn, pkg) {
			fset := snapshot.View().Session().Cache().FileSet()
			rng, err := newMappedRange(fset, m, d.Pos(), d.Pos()).Range()
			if err != nil {
				return nil, err
			}

			uri := fh.Identity().URI
			codeLens = append(codeLens, protocol.CodeLens{
				Range: rng,
				Command: protocol.Command{
					Title:     "run test",
					Command:   "test",
					Arguments: []interface{}{fn.Name.Name, uri},
				},
			})
		}
	}

	return codeLens, nil
}

func isTestFunc(fn *ast.FuncDecl, pkg Package) bool {
	typesInfo := pkg.GetTypesInfo()
	if typesInfo == nil {
		return false
	}

	sig, ok := typesInfo.ObjectOf(fn.Name).Type().(*types.Signature)
	if !ok {
		return false
	}

	// test funcs should have a single parameter, so we can exit early if that's not the case.
	if sig.Params().Len() != 1 {
		return false
	}

	firstParam, ok := sig.Params().At(0).Type().(*types.Pointer)
	if !ok {
		return false
	}

	firstParamElem, ok := firstParam.Elem().(*types.Named)
	if !ok {
		return false
	}

	firstParamObj := firstParamElem.Obj()
	if firstParamObj.Pkg().Path() != "testing" {
		return false
	}

	firstParamName := firstParamObj.Id()
	return (firstParamName == "T" && testMatcher.MatchString(fn.Name.Name)) ||
		(firstParamName == "B" && benchMatcher.MatchString(fn.Name.Name))
}

func goGenerateCodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle, f *ast.File, m *protocol.ColumnMapper) ([]protocol.CodeLens, error) {
	const ggDirective = "//go:generate"
	for _, c := range f.Comments {
		for _, l := range c.List {
			if !strings.HasPrefix(l.Text, ggDirective) {
				continue
			}
			fset := snapshot.View().Session().Cache().FileSet()
			rng, err := newMappedRange(fset, m, l.Pos(), l.Pos()+token.Pos(len(ggDirective))).Range()
			if err != nil {
				return nil, err
			}
			dir := filepath.Dir(fh.Identity().URI.Filename())
			return []protocol.CodeLens{
				{
					Range: rng,
					Command: protocol.Command{
						Title:     "run go generate",
						Command:   CommandGenerate,
						Arguments: []interface{}{dir, false},
					},
				},
				{
					Range: rng,
					Command: protocol.Command{
						Title:     "run go generate ./...",
						Command:   CommandGenerate,
						Arguments: []interface{}{dir, true},
					},
				},
			}, nil

		}
	}
	return nil, nil
}