flatten const expr

flatten_const_callexpr.go

package internal

import (
	"go/ast"
	"go/types"

	"github.com/microsoft/terraform-azurerm-provider-linter/internal/log"
	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/go/ast/astutil"
)

// NOTE: Since some replaced nodes might derived from external package, for which we loaded explicitly (which breaks
// `go/analysis` API design a bit), the position information of them is not recorded in the passed in `pass.Fset`.
// So the resulted node might have some child nodes not be able to retrieve their position based on `pass.Fset`.
// USE IT WITH CARE!!!

// FlattenConstantCallExpr will try its best to return new node with all constant `*ast.CallExpr` flattened.
// The const callexpr here means the function has exactly one statement in its body which returns exactly one expression.
// Especially, if the callexpr return a function type, then we will no longer flatten that function.
// In case there is a program error, e.g. unhandled ast type, it wil just abort the program.
// In case there is a non-const call expr, it will just ignore that node.
func FlattenConstantCallExpr(node ast.Node, pass *analysis.Pass) ast.Node {
	return astutil.Apply(node, buildApplyFunc(pass), nil)
}

func buildApplyFunc(pass *analysis.Pass) astutil.ApplyFunc {
	return func(cursor *astutil.Cursor) bool {

		node := cursor.Node()

		//ast.Print(pass.Fset, node)

		// Stop traverse if current node is a function declaration. Since we do not
		// flatten the const callexpr inside a function declaration.
		switch tnode := node.(type) {
		case *ast.FuncLit:
			return false
		case *ast.Ident:
			if obj := pass.TypesInfo.ObjectOf(tnode); obj != nil {
				if _, ok := obj.Type().(*types.Signature); ok {
					return false
				}
			}
		case *ast.SelectorExpr:
			if methodObj := pass.TypesInfo.Uses[tnode.Sel]; methodObj != nil {
				if _, ok := methodObj.Type().(*types.Signature); ok {
					return false
				}
			}
		}

		// Skip flattening for non call expressions, but go on traverse
		call, ok := node.(*ast.CallExpr)
		if !ok {
			return true
		}

		var body *ast.BlockStmt
		switch f := call.Fun.(type) {

		case *ast.FuncLit:
			body = f.Body

		case *ast.Ident:
			if f.Obj == nil {
				// Already hit the final function, no need to go on traverse.
				return false
			}
			funcdecl, ok := f.Obj.Decl.(*ast.FuncDecl)
			if !ok {
				log.PosFatalf(pass.Fset, f, "ident is not func declare")
			}
			body = funcdecl.Body

		case *ast.SelectorExpr:
			methodObj := pass.TypesInfo.Uses[f.Sel]
			if methodObj == nil {
				log.PosFatalf(pass.Fset, f, "failed ot look up the ident in selector expr")
			}

			node, err := TypeObject2AstNode4IdentDecl(pass, methodObj)
			if err != nil {
				log.PosFatalf(pass.Fset, methodObj, "failed to find ast node from type obejct")
			}

			funcdecl, ok := node.(*ast.FuncDecl)
			if !ok {
				log.PosFatalf(pass.Fset, node, "Sel is not ast.FuncDecl")
			}
			body = funcdecl.Body

		default:
			log.PosFatalf(pass.Fset, f, "unexpected type of CallExpr")
		}

		if body.List == nil {
			// stop traverse
			return false
		}
		if len(body.List) != 1 {
			// stop traverse
			return false
		}
		stmt := body.List[0]
		retstmt, ok := stmt.(*ast.ReturnStmt)
		if !ok {
			// stop traverse
			return false
		}
		if len(retstmt.Results) != 1 {
			// stop traverse
			return false
		}

		// Replaced node will not be applied any more, hence we need to
		// apply them explicitly
		cursor.Replace(FlattenConstantCallExpr(retstmt.Results[0], pass))
		return false
	}
}

lookup_object.go

package internal

import (
	"errors"
	"fmt"
	"go/ast"
	"go/types"

	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/packages"
)

type ExternalPackage int

const (
	ExtPkgTerraformSdkHelperSchema ExternalPackage = iota
	ExtPkgTerraformSdkHelperValidation
	ExtPkgAzureProviderAzurermHelpersValidate
	ExtPkgAzureProviderAzurermHelpersAzure
	ExtPkgFmt
)

var externalPackagePathMap = map[ExternalPackage]string{
	ExtPkgTerraformSdkHelperSchema:            "github.com/hashicorp/terraform-plugin-sdk/helper/schema",
	ExtPkgTerraformSdkHelperValidation:        "github.com/hashicorp/terraform-plugin-sdk/helper/validation",
	ExtPkgAzureProviderAzurermHelpersValidate: "github.com/terraform-providers/terraform-provider-azurerm/azurerm/helpers/validate",
	ExtPkgAzureProviderAzurermHelpersAzure:    "github.com/terraform-providers/terraform-provider-azurerm/azurerm/helpers/azure",
	ExtPkgFmt:                                 "fmt",
}

var pkgPool = map[string]*packages.Package{}

func IsExtPkgImported(pass *analysis.Pass, epkg ExternalPackage) bool {
	for _, importpkg := range pass.Pkg.Imports() {
		if importpkg.Path() == externalPackagePathMap[epkg] {
			return true
		}
	}
	return false
}

func LookupObjectInPackage(pass *analysis.Pass, pkg ExternalPackage, t string) types.Object {
	path := externalPackagePathMap[pkg]
	return lookupObjectInPackagePath(pass, path, t)
}

func lookupObjectInPackagePath(pass *analysis.Pass, path string, t string) types.Object {
	for _, importpkg := range pass.Pkg.Imports() {
		if importpkg.Path() == path {
			obj := importpkg.Scope().Lookup(t)
			if obj == nil {
				// this indicates a program error
				panic(fmt.Errorf(`Object "%s" not defined in %s`, t, path))
			}
			return obj
		}
	}
	return nil
}

// LookupObjectByForceLoad load the package indicated by types.Object and return the AST
// and the updated ast.Object of it (since the passed in object may not reflect the position
// in the returned AST).
// Since this API force load the package and it's transitive dependencies, there is performance
// penalty. So it should be considered as a last resort.
func LookupObjectByForceLoad(obj types.Object) (*ast.File, *ast.Object) {
	p := obj.Pkg().Path()
	var pkg *packages.Package
	if pkgPool[p] != nil {
		pkg = pkgPool[p]
	} else {
		cfg := &packages.Config{Mode: packages.NeedSyntax}
		pkgs, err := packages.Load(cfg, obj.Pkg().Path())
		if err != nil {
			return nil, nil
		}
		if packages.PrintErrors(pkgs) > 0 {
			return nil, nil
		}

		pkg = pkgs[0]
		pkgPool[p] = pkg
	}
	for _, f := range pkg.Syntax {
		obj := f.Scope.Lookup(obj.Name())
		if obj != nil {
			return f, obj
		}
	}
	return nil, nil
}

// TypeObject2AstNode4IdentDecl convert a `types.Object` to its corresponding `ast.Node`, and
// return the parent `ast.Node` which declares that node.
// It will first try to lookup inside file set of the `pass.Fset`, if it doesn't exist there,
// it means the object is defined in external package depended on of the package under lint.
// In that case we will force load the object by "go/packages", which introduces performance cost.
func TypeObject2AstNode4IdentDecl(pass *analysis.Pass, obj types.Object) (ast.Node, error) {
	var path []ast.Node
	f := LookupInternalPoserDefFile(pass, obj)
	if f != nil {
		path, _ = astutil.PathEnclosingInterval(f, obj.Pos(), obj.Pos())
	} else {
		// The files loaded in pass doesn't contain the node, it mostly indicates the node
		// is from some external package. In which case we need to load the ast manually.
		f, newobj := LookupObjectByForceLoad(obj)
		if f == nil {
			return nil, errors.New("no file defines this Object")
		}
		path, _ = astutil.PathEnclosingInterval(f, newobj.Pos(), newobj.Pos())
	}

	// the 1st node is the object's declaring identifier,
	// by walking up one level, we get the enclosing declaration
	if len(path) < 2 {
		return nil, errors.New("enclosing interval has less than 2 levels of ast.Node(s)")
	}
	declNode := path[1]
	return declNode, nil
}