methods.go

// Copyright 2019 The Hugo Authors. All rights reserved.
// Some functions in this file (see comments) is based on the Go source code,
// copyright The Go Authors and  governed by a BSD-style license.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package codegen contains helpers for code generation.
package codegen

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path"
	"path/filepath"
	"reflect"
	"regexp"
	"slices"
	"sort"
	"strings"
	"sync"
)

// Make room for insertions
const weightWidth = 1000

// NewInspector creates a new Inspector given a source root.
func NewInspector(root string) *Inspector {
	return &Inspector{ProjectRootDir: root}
}

// Inspector provides methods to help code generation. It uses a combination
// of reflection and source code AST to do the heavy lifting.
type Inspector struct {
	ProjectRootDir string

	init sync.Once

	// Determines method order. Go's reflect sorts lexicographically, so
	// we must parse the source to preserve this order.
	methodWeight map[string]map[string]int
}

// MethodsFromTypes create a method set from the include slice, excluding any
// method in exclude.
func (c *Inspector) MethodsFromTypes(include []reflect.Type, exclude []reflect.Type) Methods {
	c.parseSource()

	var methods Methods

	excludes := make(map[string]bool)

	if len(exclude) > 0 {
		for _, m := range c.MethodsFromTypes(exclude, nil) {
			excludes[m.Name] = true
		}
	}

	// There may be overlapping interfaces in types. Do a simple check for now.
	seen := make(map[string]bool)

	nameAndPackage := func(t reflect.Type) (string, string) {
		var name, pkg string

		isPointer := t.Kind() == reflect.Ptr

		if isPointer {
			t = t.Elem()
		}

		pkgPrefix := ""
		if pkgPath := t.PkgPath(); pkgPath != "" {
			pkgPath = strings.TrimSuffix(pkgPath, "/")
			_, shortPath := path.Split(pkgPath)
			pkgPrefix = shortPath + "."
			pkg = pkgPath
		}

		name = t.Name()
		if name == "" {
			// interface{}
			name = t.String()
		}

		if isPointer {
			pkgPrefix = "*" + pkgPrefix
		}

		name = pkgPrefix + name

		return name, pkg
	}

	for _, t := range include {
		for i := range t.NumMethod() {

			m := t.Method(i)
			if excludes[m.Name] || seen[m.Name] {
				continue
			}

			seen[m.Name] = true

			if m.PkgPath != "" {
				// Not exported
				continue
			}

			numIn := m.Type.NumIn()

			ownerName, _ := nameAndPackage(t)

			method := Method{Owner: t, OwnerName: ownerName, Name: m.Name}

			for i := range numIn {
				in := m.Type.In(i)

				name, pkg := nameAndPackage(in)

				if pkg != "" {
					method.Imports = append(method.Imports, pkg)
				}

				method.In = append(method.In, name)
			}

			numOut := m.Type.NumOut()

			if numOut > 0 {
				for i := range numOut {
					out := m.Type.Out(i)
					name, pkg := nameAndPackage(out)

					if pkg != "" {
						method.Imports = append(method.Imports, pkg)
					}

					method.Out = append(method.Out, name)
				}
			}

			methods = append(methods, method)
		}
	}

	sort.SliceStable(methods, func(i, j int) bool {
		mi, mj := methods[i], methods[j]

		wi := c.methodWeight[mi.OwnerName][mi.Name]
		wj := c.methodWeight[mj.OwnerName][mj.Name]

		if wi == wj {
			return mi.Name < mj.Name
		}

		return wi < wj
	})

	return methods
}

func (c *Inspector) parseSource() {
	c.init.Do(func() {
		if !strings.Contains(c.ProjectRootDir, "hugo") {
			panic("dir must be set to the Hugo root")
		}

		c.methodWeight = make(map[string]map[string]int)
		dirExcludes := regexp.MustCompile("docs|examples")
		fileExcludes := regexp.MustCompile("autogen")
		var filenames []string

		filepath.Walk(c.ProjectRootDir, func(path string, info os.FileInfo, err error) error {
			if info.IsDir() {
				if dirExcludes.MatchString(info.Name()) {
					return filepath.SkipDir
				}
			}

			if !strings.HasSuffix(path, ".go") || fileExcludes.MatchString(path) {
				return nil
			}

			filenames = append(filenames, path)

			return nil
		})

		for _, filename := range filenames {

			pkg := c.packageFromPath(filename)

			fset := token.NewFileSet()
			node, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
			if err != nil {
				panic(err)
			}

			ast.Inspect(node, func(n ast.Node) bool {
				switch t := n.(type) {
				case *ast.TypeSpec:
					if t.Name.IsExported() {
						switch it := t.Type.(type) {
						case *ast.InterfaceType:
							iface := pkg + "." + t.Name.Name
							methodNames := collectMethodsRecursive(pkg, it.Methods.List)
							weights := make(map[string]int)
							weight := weightWidth
							for _, name := range methodNames {
								weights[name] = weight
								weight += weightWidth
							}
							c.methodWeight[iface] = weights
						}
					}
				}
				return true
			})

		}

		// Complement
		for _, v1 := range c.methodWeight {
			for k2, w := range v1 {
				if v, found := c.methodWeight[k2]; found {
					for k3, v3 := range v {
						v1[k3] = (v3 / weightWidth) + w
					}
				}
			}
		}
	})
}

func (c *Inspector) packageFromPath(p string) string {
	p = filepath.ToSlash(p)
	base := path.Base(p)
	if !strings.Contains(base, ".") {
		return base
	}
	return path.Base(strings.TrimSuffix(p, base))
}

// Method holds enough information about it to recreate it.
type Method struct {
	// The interface we extracted this method from.
	Owner reflect.Type

	// String version of the above, on the form PACKAGE.NAME, e.g.
	// page.Page
	OwnerName string

	// Method name.
	Name string

	// Imports needed to satisfy the method signature.
	Imports []string

	// Argument types, including any package prefix, e.g. string, int, interface{},
	// net.Url
	In []string

	// Return types.
	Out []string
}

// Declaration creates a method declaration (without any body) for the given receiver.
func (m Method) Declaration(receiver string) string {
	return fmt.Sprintf("func (%s %s) %s%s %s", receiverShort(receiver), receiver, m.Name, m.inStr(), m.outStr())
}

// DeclarationNamed creates a method declaration (without any body) for the given receiver
// with named return values.
func (m Method) DeclarationNamed(receiver string) string {
	return fmt.Sprintf("func (%s %s) %s%s %s", receiverShort(receiver), receiver, m.Name, m.inStr(), m.outStrNamed())
}

// Delegate creates a delegate call string.
func (m Method) Delegate(receiver, delegate string) string {
	ret := ""
	if len(m.Out) > 0 {
		ret = "return "
	}
	return fmt.Sprintf("%s%s.%s.%s%s", ret, receiverShort(receiver), delegate, m.Name, m.inOutStr())
}

func (m Method) String() string {
	return m.Name + m.inStr() + " " + m.outStr() + "\n"
}

func (m Method) inOutStr() string {
	if len(m.In) == 0 {
		return "()"
	}

	args := make([]string, len(m.In))
	for i := range args {
		args[i] = fmt.Sprintf("arg%d", i)
	}
	return "(" + strings.Join(args, ", ") + ")"
}

func (m Method) inStr() string {
	if len(m.In) == 0 {
		return "()"
	}

	args := make([]string, len(m.In))
	for i := range args {
		args[i] = fmt.Sprintf("arg%d %s", i, m.In[i])
	}
	return "(" + strings.Join(args, ", ") + ")"
}

func (m Method) outStr() string {
	if len(m.Out) == 0 {
		return ""
	}
	if len(m.Out) == 1 {
		return m.Out[0]
	}

	return "(" + strings.Join(m.Out, ", ") + ")"
}

func (m Method) outStrNamed() string {
	if len(m.Out) == 0 {
		return ""
	}

	outs := make([]string, len(m.Out))
	for i := range outs {
		outs[i] = fmt.Sprintf("o%d %s", i, m.Out[i])
	}

	return "(" + strings.Join(outs, ", ") + ")"
}

// Methods represents a list of methods for one or more interfaces.
// The order matches the defined order in their source file(s).
type Methods []Method

// Imports returns a sorted list of package imports needed to satisfy the
// signatures of all methods.
func (m Methods) Imports() []string {
	var pkgImports []string
	for _, method := range m {
		pkgImports = append(pkgImports, method.Imports...)
	}
	if len(pkgImports) > 0 {
		pkgImports = uniqueNonEmptyStrings(pkgImports)
		sort.Strings(pkgImports)
	}
	return pkgImports
}

// ToMarshalJSON creates a MarshalJSON method for these methods. Any method name
// matching any of the regexps in excludes will be ignored.
func (m Methods) ToMarshalJSON(receiver, pkgPath string, excludes ...string) (string, []string) {
	var sb strings.Builder

	r := receiverShort(receiver)
	what := firstToUpper(trimAsterisk(receiver))
	pgkName := path.Base(pkgPath)

	fmt.Fprintf(&sb, "func Marshal%sToJSON(%s %s) ([]byte, error) {\n", what, r, receiver)

	var methods Methods
	excludeRes := make([]*regexp.Regexp, len(excludes))

	for i, exclude := range excludes {
		excludeRes[i] = regexp.MustCompile(exclude)
	}

	for _, method := range m {
		// Exclude methods with arguments and incompatible return values
		if len(method.In) > 0 || len(method.Out) == 0 || len(method.Out) > 2 {
			continue
		}

		if len(method.Out) == 2 {
			if method.Out[1] != "error" {
				continue
			}
		}

		for _, re := range excludeRes {
			if re.MatchString(method.Name) {
				continue
			}
		}

		methods = append(methods, method)
	}

	for _, method := range methods {
		varn := varName(method.Name)
		if len(method.Out) == 1 {
			fmt.Fprintf(&sb, "\t%s := %s.%s()\n", varn, r, method.Name)
		} else {
			fmt.Fprintf(&sb, "\t%s, err := %s.%s()\n", varn, r, method.Name)
			fmt.Fprint(&sb, "\tif err != nil {\n\t\treturn nil, err\n\t}\n")
		}
	}

	fmt.Fprint(&sb, "\n\ts := struct {\n")

	for _, method := range methods {
		fmt.Fprintf(&sb, "\t\t%s %s\n", method.Name, typeName(method.Out[0], pgkName))
	}

	fmt.Fprint(&sb, "\n\t}{\n")

	for _, method := range methods {
		varn := varName(method.Name)
		fmt.Fprintf(&sb, "\t\t%s: %s,\n", method.Name, varn)
	}

	fmt.Fprint(&sb, "\n\t}\n\n")
	fmt.Fprint(&sb, "\treturn json.Marshal(&s)\n}")

	pkgImports := append(methods.Imports(), "encoding/json")

	if pkgPath != "" {
		// Exclude self
		for i, pkgImp := range pkgImports {
			if pkgImp == pkgPath {
				pkgImports = slices.Delete(pkgImports, i, i+1)
			}
		}
	}

	return sb.String(), pkgImports
}

func collectMethodsRecursive(pkg string, f []*ast.Field) []string {
	var methodNames []string
	for _, m := range f {
		if m.Names != nil {
			methodNames = append(methodNames, m.Names[0].Name)
			continue
		}

		if ident, ok := m.Type.(*ast.Ident); ok && ident.Obj != nil {
			switch tt := ident.Obj.Decl.(*ast.TypeSpec).Type.(type) {
			case *ast.InterfaceType:
				// Embedded interface
				methodNames = append(
					methodNames,
					collectMethodsRecursive(
						pkg,
						tt.Methods.List)...)
			}
		} else {
			// Embedded, but in a different file/package. Return the
			// package.Name and deal with that later.
			name := packageName(m.Type)
			if !strings.Contains(name, ".") {
				// Assume current package
				name = pkg + "." + name
			}
			methodNames = append(methodNames, name)
		}
	}

	return methodNames
}

func firstToLower(name string) string {
	return strings.ToLower(name[:1]) + name[1:]
}

func firstToUpper(name string) string {
	return strings.ToUpper(name[:1]) + name[1:]
}

func packageName(e ast.Expr) string {
	switch tp := e.(type) {
	case *ast.Ident:
		return tp.Name
	case *ast.SelectorExpr:
		return fmt.Sprintf("%s.%s", packageName(tp.X), packageName(tp.Sel))
	}
	return ""
}

func receiverShort(receiver string) string {
	return strings.ToLower(trimAsterisk(receiver))[:1]
}

func trimAsterisk(name string) string {
	return strings.TrimPrefix(name, "*")
}

func typeName(name, pkg string) string {
	return strings.TrimPrefix(name, pkg+".")
}

func uniqueNonEmptyStrings(s []string) []string {
	var unique []string
	set := map[string]any{}
	for _, val := range s {
		if val == "" {
			continue
		}
		if _, ok := set[val]; !ok {
			unique = append(unique, val)
			set[val] = val
		}
	}
	return unique
}

func varName(name string) string {
	name = firstToLower(name)

	// Adjust some reserved keywords, see https://golang.org/ref/spec#Keywords
	switch name {
	case "type":
		name = "typ"
	case "package":
		name = "pkg"
		// Not reserved, but syntax highlighters has it as a keyword.
	case "len":
		name = "length"
	}

	return name
}