types_go118.go

//go:build go1.18
// +build go1.18

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/token"
	"go/types"
	"io"
	"sort"
	"strings"

	pkgwalk "github.com/visualfc/gotools/types"
)

type TypeParam = types.TypeParam
type TypeParamList = types.TypeParamList

var TILDE = token.TILDE

func newFuncType(tparams, params, results *ast.FieldList) *ast.FuncType {
	return &ast.FuncType{TypeParams: tparams, Params: params, Results: results}
}

func newTypeSpec(name string, tparams *ast.FieldList) *ast.TypeSpec {
	return &ast.TypeSpec{
		Name:       ast.NewIdent(name),
		TypeParams: tparams,
	}
}

func toTypeParam(pkg *types.Package, t *TypeParam) ast.Expr {
	return toObjectExpr(pkg, t.Obj())
}

func ForSignature(sig *types.Signature) *TypeParamList {
	return sig.TypeParams()
}

func ForFuncType(typ *ast.FuncType) *ast.FieldList {
	return typ.TypeParams
}

// RecvTypeParams returns a nil slice.
func RecvTypeParams(sig *types.Signature) *TypeParamList {
	return sig.RecvTypeParams()
}

func ForNamed(named *types.Named) *TypeParamList {
	return named.TypeParams()
}

func toFieldListX(pkg *types.Package, t *types.TypeParamList) *ast.FieldList {
	if t == nil {
		return nil
	}
	n := t.Len()
	flds := make([]*ast.Field, n)
	for i := 0; i < n; i++ {
		item := t.At(i)
		names := []*ast.Ident{ast.NewIdent(item.Obj().Name())}
		typ := toType(pkg, item.Constraint())
		flds[i] = &ast.Field{Names: names, Type: typ}
	}
	return &ast.FieldList{
		List: flds,
	}
}

func toFuncType(pkg *types.Package, sig *types.Signature) *ast.FuncType {
	params := toFieldList(pkg, sig.Params())
	results := toFieldList(pkg, sig.Results())
	if sig.Variadic() {
		n := len(params)
		if n == 0 {
			panic("TODO: toFuncType error")
		}
		toVariadic(params[n-1])
	}
	return &ast.FuncType{
		TypeParams: toFieldListX(pkg, sig.TypeParams()),
		Params:     &ast.FieldList{List: params},
		Results:    &ast.FieldList{List: results},
	}
}

func toTypeSpec(pkg *types.Package, t *types.TypeName) *ast.TypeSpec {
	var assign token.Pos
	if t.IsAlias() {
		assign = 1
	}
	typ := t.Type()
	ts := &ast.TypeSpec{
		Name:   ast.NewIdent(t.Name()),
		Assign: assign,
		Type:   toType(pkg, typ.Underlying()),
	}
	if named, ok := typ.(*types.Named); ok {
		ts.TypeParams = toFieldListX(pkg, named.TypeParams())
	}
	return ts
}

// converts type expressions like:
// ast.Expr
// *ast.Expr
// $ast$go/ast.Expr
// to a path that can be used to lookup a type related Decl
func get_type_path(e ast.Expr) (r type_path) {
	if e == nil {
		return type_path{"", "", nil}
	}

	switch t := e.(type) {
	case *ast.Ident:
		r.name = t.Name
	case *ast.StarExpr:
		r = get_type_path(t.X)
	case *ast.SelectorExpr:
		if ident, ok := t.X.(*ast.Ident); ok {
			r.pkg = ident.Name
			if r.pkg == "main" {
				r.pkg = ""
			}
		}
		r.name = t.Sel.Name
	case *ast.IndexExpr:
		r = get_type_path(t.X)
		r.targs = []ast.Expr{t.Index}
	case *ast.IndexListExpr:
		r = get_type_path(t.X)
		r.targs = t.Indices
	}
	return
}

func ast_decl_typeparams(decl ast.Decl) *ast.FieldList {
	switch t := decl.(type) {
	case *ast.GenDecl:
		if t.Tok == token.TYPE {
			if len(t.Specs) > 0 {
				if spec, ok := t.Specs[0].(*ast.TypeSpec); ok {
					return spec.TypeParams
				}
			}
		}
	case *ast.FuncDecl:
		return t.Type.TypeParams
	}
	return nil
}

func hasTypeParams(typ types.Type) bool {
	switch t := typ.(type) {
	case *types.Named:
		return t.TypeParams() != nil && (t.Origin() == t)
	case *types.Signature:
		return t.TypeParams() != nil
	}
	return false
}

func funcHasTypeParams(typ *ast.FuncType) bool {
	return typ.TypeParams != nil
}

func toNamedType(pkg *types.Package, t *types.Named) ast.Expr {
	expr := toObjectExpr(pkg, t.Obj())
	if targs := t.TypeArgs(); targs != nil {
		n := targs.Len()
		indices := make([]ast.Expr, n)
		for i := 0; i < n; i++ {
			indices[i] = toType(pkg, targs.At(i))
		}
		if n == 1 {
			expr = &ast.IndexExpr{
				X:     expr,
				Index: indices[0],
			}
		} else {
			expr = &ast.IndexListExpr{
				X:       expr,
				Indices: indices,
			}
		}
	}
	return expr
}

func lookup_types_near_instance(ident *ast.Ident, pos token.Pos, info *types.Info) types.Type {
	var ar []*typ_distance
	for k, v := range info.Instances {
		if ident.Name == k.Name && pos > k.End() {
			ar = append(ar, &typ_distance{pos - k.End(), v.Type})
		}
	}
	switch len(ar) {
	case 0:
		return nil
	case 1:
		return ar[0].typ
	default:
		sort.Slice(ar, func(i, j int) bool {
			return ar[i].pos < ar[j].pos
		})
		return ar[0].typ
	}
	return nil
}

func DefaultPkgConfig() *pkgwalk.PkgConfig {
	conf := &pkgwalk.PkgConfig{IgnoreFuncBodies: false, AllowBinary: true, WithTestFiles: true}
	conf.Info = &types.Info{
		Uses:       make(map[*ast.Ident]types.Object),
		Defs:       make(map[*ast.Ident]types.Object),
		Selections: make(map[*ast.SelectorExpr]*types.Selection),
		Types:      make(map[ast.Expr]types.TypeAndValue),
		Scopes:     make(map[ast.Node]*types.Scope),
		Implicits:  make(map[ast.Node]types.Object),
		Instances:  make(map[*ast.Ident]types.Instance),
	}
	conf.XInfo = &types.Info{
		Uses:       make(map[*ast.Ident]types.Object),
		Defs:       make(map[*ast.Ident]types.Object),
		Selections: make(map[*ast.SelectorExpr]*types.Selection),
		Types:      make(map[ast.Expr]types.TypeAndValue),
		Scopes:     make(map[ast.Node]*types.Scope),
		Implicits:  make(map[ast.Node]types.Object),
		Instances:  make(map[*ast.Ident]types.Instance),
	}
	return conf
}

func pretty_print_type_expr(out io.Writer, e ast.Expr, canonical_aliases map[string]string) {
	switch t := e.(type) {
	case *ast.StarExpr:
		fmt.Fprintf(out, "*")
		pretty_print_type_expr(out, t.X, canonical_aliases)
	case *ast.Ident:
		if strings.HasPrefix(t.Name, "$") {
			// beautify anonymous types
			switch t.Name[1] {
			case 's':
				fmt.Fprintf(out, "struct")
			case 'i':
				// ok, in most cases anonymous interface is an
				// empty interface, I'll just pretend that
				// it's always true
				fmt.Fprintf(out, "interface{}")
			}
		} else if !*g_debug && strings.HasPrefix(t.Name, "!") {
			// these are full package names for disambiguating and pretty
			// printing packages within packages, e.g.
			// !go/ast!ast vs. !github.com/nsf/my/ast!ast
			// another ugly hack, if people are punished in hell for ugly hacks
			// I'm screwed...
			emarkIdx := strings.LastIndex(t.Name, "!")
			path := t.Name[1:emarkIdx]
			alias := canonical_aliases[path]
			if alias == "" {
				alias = t.Name[emarkIdx+1:]
			}
			fmt.Fprintf(out, alias)
		} else {
			fmt.Fprintf(out, t.Name)
		}
	case *ast.ArrayType:
		al := ""
		if t.Len != nil {
			al = get_array_len(t.Len)
		}
		if al != "" {
			fmt.Fprintf(out, "[%s]", al)
		} else {
			fmt.Fprintf(out, "[]")
		}
		pretty_print_type_expr(out, t.Elt, canonical_aliases)
	case *ast.SelectorExpr:
		pretty_print_type_expr(out, t.X, canonical_aliases)
		fmt.Fprintf(out, ".%s", t.Sel.Name)
	case *ast.FuncType:
		fmt.Fprintf(out, "func(")
		pretty_print_func_field_list(out, t.Params, canonical_aliases)
		fmt.Fprintf(out, ")")

		buf := bytes.NewBuffer(make([]byte, 0, 256))
		nresults := pretty_print_func_field_list(buf, t.Results, canonical_aliases)
		if nresults > 0 {
			results := buf.String()
			if strings.IndexAny(results, ", ") != -1 {
				results = "(" + results + ")"
			}
			fmt.Fprintf(out, " %s", results)
		}
	case *ast.MapType:
		fmt.Fprintf(out, "map[")
		pretty_print_type_expr(out, t.Key, canonical_aliases)
		fmt.Fprintf(out, "]")
		pretty_print_type_expr(out, t.Value, canonical_aliases)
	case *ast.InterfaceType:
		fmt.Fprintf(out, "interface{}")
	case *ast.Ellipsis:
		fmt.Fprintf(out, "...")
		pretty_print_type_expr(out, t.Elt, canonical_aliases)
	case *ast.StructType:
		fmt.Fprintf(out, "struct")
	case *ast.ChanType:
		switch t.Dir {
		case ast.RECV:
			fmt.Fprintf(out, "<-chan ")
		case ast.SEND:
			fmt.Fprintf(out, "chan<- ")
		case ast.SEND | ast.RECV:
			fmt.Fprintf(out, "chan ")
		}
		pretty_print_type_expr(out, t.Value, canonical_aliases)
	case *ast.ParenExpr:
		fmt.Fprintf(out, "(")
		pretty_print_type_expr(out, t.X, canonical_aliases)
		fmt.Fprintf(out, ")")
	case *ast.IndexExpr:
		pretty_print_type_expr(out, t.X, canonical_aliases)
		fmt.Fprintf(out, "[")
		pretty_print_type_expr(out, t.Index, canonical_aliases)
		fmt.Fprintf(out, "]")
	case *ast.IndexListExpr:
		pretty_print_type_expr(out, t.X, canonical_aliases)
		fmt.Fprintf(out, "[")
		for i, index := range t.Indices {
			if i != 0 {
				fmt.Fprintf(out, ", ")
			}
			pretty_print_type_expr(out, index, canonical_aliases)
		}
		fmt.Fprintf(out, "]")
	case *ast.BadExpr:
		// TODO: probably I should check that in a separate function
		// and simply discard declarations with BadExpr as a part of their
		// type
	default:
		// the element has some weird type, just ignore it
	}
}

func funHasTypeArgs(fun ast.Expr) bool {
	switch fun.(type) {
	case *ast.IndexExpr:
		return true
	case *ast.IndexListExpr:
		return true
	}
	return false
}