pkg: teach gen.go about go/types

This cuts about 30 lines of code, since logic like toValue is obsolete.
We also no longer need to deal with irrelevant bits of the syntax tree.

I'm leaving a TODO about polishing up the use of strings in the future.
For now, this is enough to unblock the use of more complex constants,
since the previous code with toValue was too simple.

Signed-off-by: Daniel Martí <[email protected]>
Change-Id: Idb6edb844359a9936a719f42ed788e6704f961f8
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/557328
TryBot-Result: CUEcueckoo <[email protected]>
Reviewed-by: Roger Peppe <[email protected]>

Author: mvdan

pkg/gen.go

@@ -28,17 +28,16 @@ import (
 	_ "embed"
 	"flag"
 	"fmt"
-	"go/ast"
 	"go/constant"
 	"go/format"
-	"go/parser"
-	"go/printer"
 	"go/token"
+	"go/types"
 	"log"
 	"math/big"
 	"os"
 	"path"
 	"path/filepath"
+	"sort"
 	"strings"
 	"text/template"
 
@@ -107,7 +106,7 @@ func main() {
 		packagesList = append(packagesList, path.Join(pkgParent, pkg))
 	}
 
-	cfg := &packages.Config{Mode: packages.NeedName | packages.NeedFiles}
+	cfg := &packages.Config{Mode: packages.NeedName | packages.NeedFiles | packages.NeedTypes}
 	pkgs, err := packages.Load(cfg, packagesList...)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "load: %v\n", err)
@@ -127,7 +126,6 @@ type generator struct {
 	dir        string
 	w          *bytes.Buffer
 	cuePkgPath string
-	fset       *token.FileSet
 	first      bool
 }
 
@@ -142,7 +140,6 @@ func generate(pkg *packages.Package) error {
 		dir:        pkgDir,
 		cuePkgPath: cuePkg,
 		w:          &bytes.Buffer{},
-		fset:       token.NewFileSet(),
 	}
 
 	params := headerParams{
@@ -175,13 +172,8 @@ func generate(pkg *packages.Package) error {
 	if !skipRegister {
 		fmt.Fprintf(g.w, "var p = &pkg.Package{\nNative: []*pkg.Builtin{")
 		g.first = true
-		for _, filename := range pkg.GoFiles {
-			if filename == genFile {
-				continue
-			}
-			if err := g.processGo(filename); err != nil {
-				return err
-			}
+		if err := g.processGo(pkg); err != nil {
+			return err
 		}
 		fmt.Fprintf(g.w, "},\n")
 		if err := g.processCUE(); err != nil {
@@ -192,6 +184,7 @@ func generate(pkg *packages.Package) error {
 
 	b, err := format.Source(g.w.Bytes())
 	if err != nil {
+		fmt.Printf("go/format error on %s: %v\n", pkg.PkgPath, err)
 		b = g.w.Bytes() // write the unformatted source
 	}
 
@@ -241,115 +234,90 @@ func (g *generator) processCUE() error {
 	return nil
 }
 
-func (g *generator) processGo(filename string) error {
-	f, err := parser.ParseFile(g.fset, filename, nil, parser.ParseComments)
-	if err != nil {
-		return err
-	}
+func (g *generator) processGo(pkg *packages.Package) error {
+	// We sort the objects by their original source code position.
+	// Otherwise go/types defaults to sorting by name strings.
+	// We could remove this code if we were fine with sorting by name.
+	scope := pkg.Types.Scope()
+	type objWithPos struct {
+		obj types.Object
+		pos token.Position
+	}
+	var objs []objWithPos
+	for _, name := range scope.Names() {
+		obj := scope.Lookup(name)
+		objs = append(objs, objWithPos{obj, pkg.Fset.Position(obj.Pos())})
+	}
+	sort.Slice(objs, func(i, j int) bool {
+		obj1, obj2 := objs[i], objs[j]
+		if obj1.pos.Filename == obj2.pos.Filename {
+			return obj1.pos.Line < obj2.pos.Line
+		}
+		return obj1.pos.Filename < obj2.pos.Filename
+	})
 
-	for _, d := range f.Decls {
-		switch x := d.(type) {
-		case *ast.GenDecl:
-			switch x.Tok {
-			case token.CONST:
-				for _, spec := range x.Specs {
-					spec := spec.(*ast.ValueSpec)
-					if ast.IsExported(spec.Names[0].Name) {
-						g.genConst(spec)
-					}
-				}
-			case token.VAR:
-				continue
-			case token.TYPE:
-				// TODO: support type declarations.
-				continue
-			case token.IMPORT:
-				continue
+	for _, obj := range objs {
+		obj := obj.obj // no longer need the token.Position
+		if !obj.Exported() {
+			continue
+		}
+		// TODO: support type declarations.
+		switch obj := obj.(type) {
+		case *types.Const:
+			var value string
+			switch v := obj.Val(); v.Kind() {
+			case constant.Bool, constant.Int, constant.String:
+				// TODO: convert octal numbers
+				value = v.ExactString()
+			case constant.Float:
+				var rat big.Rat
+				rat.SetString(v.ExactString())
+				var float big.Float
+				float.SetRat(&rat)
+				value = float.Text('g', -1)
 			default:
-				panic(fmt.Errorf("gen %s: unexpected spec of type %s", filename, x.Tok))
+				fmt.Printf("Dropped entry %s.%s (%T: %v)\n", g.cuePkgPath, obj.Name(), v.Kind(), v.ExactString())
+				continue
 			}
-		case *ast.FuncDecl:
-			g.genFunc(x)
+			g.sep()
+			fmt.Fprintf(g.w, "{\nName: %q,\n Const: %q,\n}", obj.Name(), value)
+		case *types.Func:
+			g.genFunc(obj)
 		}
 	}
 	return nil
 }
 
-func (g *generator) genConst(spec *ast.ValueSpec) {
-	name := spec.Names[0].Name
-	value := ""
-	switch v := g.toValue(spec.Values[0]); v.Kind() {
-	case constant.Bool, constant.Int, constant.String:
-		// TODO: convert octal numbers
-		value = v.ExactString()
-	case constant.Float:
-		var rat big.Rat
-		rat.SetString(v.ExactString())
-		var float big.Float
-		float.SetRat(&rat)
-		value = float.Text('g', -1)
-	default:
-		fmt.Printf("Dropped entry %s.%s (%T: %v)\n", g.cuePkgPath, name, v.Kind(), v.ExactString())
-		return
-	}
-	g.sep()
-	fmt.Fprintf(g.w, "{\nName: %q,\n Const: %q,\n}", name, value)
-}
-
-func (g *generator) toValue(x ast.Expr) constant.Value {
-	switch x := x.(type) {
-	case *ast.BasicLit:
-		return constant.MakeFromLiteral(x.Value, x.Kind, 0)
-	case *ast.BinaryExpr:
-		return constant.BinaryOp(g.toValue(x.X), x.Op, g.toValue(x.Y))
-	case *ast.UnaryExpr:
-		return constant.UnaryOp(x.Op, g.toValue(x.X), 0)
-	default:
-		panic(fmt.Errorf("%s: unsupported expression type %T: %#v", g.cuePkgPath, x, x))
-	}
-}
+var errorType = types.Universe.Lookup("error").Type()
 
-func (g *generator) genFunc(x *ast.FuncDecl) {
-	if x.Body == nil || !ast.IsExported(x.Name.Name) {
+func (g *generator) genFunc(fn *types.Func) {
+	sign := fn.Type().(*types.Signature)
+	if sign.Recv() != nil {
 		return
 	}
-	types := []string{}
-	if x.Type.Results != nil {
-		for _, f := range x.Type.Results.List {
-			if len(f.Names) > 0 {
-				for range f.Names {
-					types = append(types, g.goKind(f.Type))
-				}
-			} else {
-				types = append(types, g.goKind(f.Type))
-			}
-		}
-	}
-	if x.Recv != nil {
-		return
-	}
-	if n := len(types); n != 1 && (n != 2 || types[1] != "error") {
-		fmt.Printf("Dropped func %s.%s: must have one return value or a value and an error %v\n", g.cuePkgPath, x.Name.Name, types)
+	params := sign.Params()
+	results := sign.Results()
+	if results == nil || (results.Len() != 1 && results.At(1).Type() != errorType) {
+		fmt.Printf("Dropped func %s.%s: must have one return value or a value and an error %v\n", g.cuePkgPath, fn.Name(), sign)
 		return
 	}
 
 	g.sep()
 	fmt.Fprintf(g.w, "{\n")
 	defer fmt.Fprintf(g.w, "}")
 
-	fmt.Fprintf(g.w, "Name: %q,\n", x.Name.Name)
+	fmt.Fprintf(g.w, "Name: %q,\n", fn.Name())
 
 	args := []string{}
 	vals := []string{}
 	kind := []string{}
-	for _, f := range x.Type.Params.List {
-		for _, name := range f.Names {
-			typ := strings.Title(g.goKind(f.Type))
-			argKind := g.goToCUE(f.Type)
-			vals = append(vals, fmt.Sprintf("c.%s(%d)", typ, len(args)))
-			args = append(args, name.Name)
-			kind = append(kind, argKind)
-		}
+	for i := 0; i < params.Len(); i++ {
+		param := params.At(i)
+		typ := strings.Title(g.goKind(param.Type()))
+		argKind := g.goToCUE(param.Type())
+		vals = append(vals, fmt.Sprintf("c.%s(%d)", typ, len(args)))
+		args = append(args, param.Name())
+		kind = append(kind, argKind)
 	}
 
 	fmt.Fprintf(g.w, "Params: []pkg.Param{\n")
@@ -358,8 +326,7 @@ func (g *generator) genFunc(x *ast.FuncDecl) {
 	}
 	fmt.Fprintf(g.w, "\n},\n")
 
-	expr := x.Type.Results.List[0].Type
-	fmt.Fprintf(g.w, "Result: %s,\n", g.goToCUE(expr))
+	fmt.Fprintf(g.w, "Result: %s,\n", g.goToCUE(results.At(0).Type()))
 
 	argList := strings.Join(args, ", ")
 	valList := strings.Join(vals, ", ")
@@ -376,45 +343,45 @@ func (g *generator) genFunc(x *ast.FuncDecl) {
 	}
 	fmt.Fprintln(g.w, "if c.Do() {")
 	defer fmt.Fprintln(g.w, "}")
-	if len(types) == 1 {
-		fmt.Fprintf(g.w, "c.Ret = %s(%s)", x.Name.Name, argList)
+	if results.Len() == 1 {
+		fmt.Fprintf(g.w, "c.Ret = %s(%s)", fn.Name(), argList)
 	} else {
-		fmt.Fprintf(g.w, "c.Ret, c.Err = %s(%s)", x.Name.Name, argList)
+		fmt.Fprintf(g.w, "c.Ret, c.Err = %s(%s)", fn.Name(), argList)
 	}
 }
 
-func (g *generator) goKind(expr ast.Expr) string {
-	if star, isStar := expr.(*ast.StarExpr); isStar {
-		expr = star.X
+// TODO(mvdan): goKind and goToCUE still use a lot of strings; simplify.
+
+func (g *generator) goKind(typ types.Type) string {
+	if ptr, ok := typ.(*types.Pointer); ok {
+		typ = ptr.Elem()
 	}
-	w := &bytes.Buffer{}
-	printer.Fprint(w, g.fset, expr)
-	switch str := w.String(); str {
-	case "big.Int":
+	switch str := types.TypeString(typ, nil); str {
+	case "math/big.Int":
 		return "bigInt"
-	case "big.Float":
+	case "math/big.Float":
 		return "bigFloat"
-	case "big.Rat":
+	case "math/big.Rat":
 		return "bigRat"
-	case "adt.Bottom":
+	case "cuelang.org/go/internal/core/adt.Bottom":
 		return "error"
-	case "internal.Decimal":
+	case "github.com/cockroachdb/apd/v3.Decimal":
 		return "decimal"
-	case "pkg.List":
+	case "cuelang.org/go/internal/pkg.List":
 		return "cueList"
-	case "pkg.Struct":
+	case "cuelang.org/go/internal/pkg.Struct":
 		return "struct"
-	case "[]*internal.Decimal":
+	case "[]*github.com/cockroachdb/apd/v3.Decimal":
 		return "decimalList"
-	case "cue.Value":
+	case "cuelang.org/go/cue.Value":
 		return "value"
-	case "cue.List":
+	case "cuelang.org/go/cue.List":
 		return "list"
 	case "[]string":
 		return "stringList"
 	case "[]byte":
 		return "bytes"
-	case "[]cue.Value":
+	case "[]cuelang.org/go/cue.Value":
 		return "list"
 	case "io.Reader":
 		return "reader"
@@ -425,9 +392,9 @@ func (g *generator) goKind(expr ast.Expr) string {
 	}
 }
 
-func (g *generator) goToCUE(expr ast.Expr) (cueKind string) {
+func (g *generator) goToCUE(typ types.Type) (cueKind string) {
 	// TODO: detect list and structs types for return values.
-	switch k := g.goKind(expr); k {
+	switch k := g.goKind(typ); k {
 	case "error":
 		cueKind += "adt.BottomKind"
 	case "bool":