addlogcall.go

// Copyright 2020 cloudeng llc. All rights reserved.
// Use of this source code is governed by the Apache-2.0
// license that can be found in the LICENSE file.

package annotators

import (
	"context"
	"fmt"
	"go/ast"
	"go/token"
	"go/types"
	"strings"

	"cloudeng.io/errors"
	"cloudeng.io/go/cmd/goannotate/annotators/functions"
	"cloudeng.io/go/cmd/goannotate/annotators/internal"
	"cloudeng.io/go/locate"
	"cloudeng.io/go/locate/locateutil"
	"cloudeng.io/text/edit"
	"golang.org/x/tools/go/packages"
	"gopkg.in/yaml.v2"
)

// AddLogCall represents an annotator for adding a function call that logs
// the entry and exit to every function and method that is matched by the
// locator.
type AddLogCall struct {
	EssentialOptions `yaml:",inline"`
	LocateOptions    `yaml:",inline"`

	AtLeastStatements   int            `yaml:"atLeastStatements" annotator:"the number of statements that must be present in a function in order for it to be annotated."`
	NoAnnotationComment string         `yaml:"noAnnotationComment" annotator:"do not annotate functions that contain this comment"`
	CallGenerator       functions.Spec `yaml:"callGenerator" annotator:"the spec for the function call to be generated"`
}

func init() {
	Register(&AddLogCall{})
}

// New implements annotators.Annotator.
func (lc *AddLogCall) New(name string) Annotation {
	n := &AddLogCall{}
	n.Name = name
	return n
}

// UnmarshalYAML implements annotators.Annotation.
func (lc *AddLogCall) UnmarshalYAML(buf []byte) error {
	return yaml.Unmarshal(buf, lc)
}

// AddLogCallDescription documents AddLogCall.
const AddLogCallDescription = `
AddLogCall is an annotator to add function calls that are intended to log entry and exit from functions. The calls will be added as the first statement in the specified function.
`

// Describe implements annotators.Annotation.
func (lc *AddLogCall) Describe() string {
	out := &strings.Builder{}
	out.WriteString(internal.MustDescribe(lc, AddLogCallDescription))
	out.WriteString("\n")
	out.WriteString("    Available Call Generators:\n\n")
	for _, cg := range functions.CallGenerators() {
		out.WriteString("    " + cg + "\n")
	}
	out.WriteString("\n")
	for _, cg := range functions.CallGenerators() {
		out.WriteString(internal.Indent(functions.Lookup(cg).Describe(), 4))
		out.WriteString("\n")
	}
	return out.String()
}

// Do implements annotators.Annotation.
func (lc *AddLogCall) Do(ctx context.Context, root string, pkgs []string) error {
	callgen := functions.Lookup(lc.CallGenerator.Type)
	if callgen == nil {
		return fmt.Errorf("failed to find function call generator for %v", lc.CallGenerator.Type)
	}
	locator := locate.New(
		concurrencyOpt(lc.Concurrency),
		locate.Trace(Verbosef),
		locate.IgnoreMissingFuctionsEtc(),
		locate.IncludeMethods(lc.IncludeMethods),
	)
	locator.AddInterfaces(lc.Interfaces...)
	locator.AddFunctions(lc.Functions...)
	if len(pkgs) == 0 {
		pkgs = lc.Packages
	}
	locator.AddPackages(pkgs...)
	Verbosef("locating functions to be annotated with a logcall...")
	if err := locator.Do(ctx); err != nil {
		return fmt.Errorf("failed to locate functions and/or interface implementations: %v", err)
	}

	commentMaps := locator.MakeCommentMaps()
	comment := fmt.Sprintf("DO NOT EDIT, AUTO GENERATED BY %s#%s", lc.Type, lc.Name)

	dirty := map[string]bool{}
	edits := map[string][]edit.Delta{}
	errs := &errors.M{}
	locator.WalkFunctions(func(fullname string,
		pkg *packages.Package,
		file *ast.File,
		fn *types.Func,
		decl *ast.FuncDecl,
		implements []string) {
		if locateutil.FunctionStatements(decl) < lc.AtLeastStatements {
			return
		}
		if len(lc.NoAnnotationComment) > 0 {
			cmap := commentMaps[file]
			if locateutil.FunctionHasComment(decl, cmap, lc.NoAnnotationComment) {
				return
			}
		}

		invovation, err := callgen.Generate(pkg.Fset, fn, decl)
		if err != nil {
			errs.Append(err)
			return
		}

		if lc.alreadyAnnotated(fn, decl, commentMaps[file], comment) {
			Verbosef("%v: already annotated\n", fullname)
			return
		}
		lbrace := pkg.Fset.PositionFor(decl.Body.Lbrace, false)
		delta := edit.InsertString(lbrace.Offset+1, invovation+" // "+comment)
		edits[lbrace.Filename] = append(edits[lbrace.Filename], delta)
		dirty[lbrace.Filename] = true
		Verbosef("function: %v @ %v\n", fullname, lbrace)
	})

	importPath := callgen.Import()
	importStatement := "\n" + `import "` + importPath + `"` + "\n"

	locator.WalkFiles(func(filename string,
		pkg *packages.Package,
		comments ast.CommentMap,
		file *ast.File,
		mask locate.HitMask) {
		if !dirty[filename] || ((mask | locate.HasFunction) == 0) {
			return
		}
		if locateutil.IsImportedByFile(file, importPath) {
			Verbosef("%v: %v: already imported\n", filename, importPath)
			return
		}
		_, end := locateutil.ImportBlock(file)
		if end == token.NoPos {
			end = file.Name.End()
		}
		pos := pkg.Fset.PositionFor(end, false)
		delta := edit.InsertString(pos.Offset, importStatement)
		edits[pos.Filename] = append(edits[pos.Filename], delta)
		Verbosef("import: %v @ %v\n", importPath, pos)
	})

	if err := errs.Err(); err != nil {
		return err
	}
	return applyEdits(ctx, computeOutputs(root, edits), edits)
}

func (lc *AddLogCall) alreadyAnnotated(fn *types.Func, decl *ast.FuncDecl, cmap ast.CommentMap, comment string) bool {
	if locateutil.FunctionStatements(decl) == 0 {
		return false
	}
	deferStmt, ok := decl.Body.List[0].(*ast.DeferStmt)
	if !ok {
		return false
	}
	comments := cmap[deferStmt]
	for _, c := range comments {
		if c := c.Text(); strings.HasPrefix(c, comment) {
			return true
		}
	}
	return false
}