cue/ast/astutil: refactor to prepare for Sanitize func

- Factor out code
- Resolve imports in Resolve func

The latter is now possible because the package name, and
thus identifier, is now uniquely determined from the import
path.

Change-Id: I648c02eb43d0b104bf4bc3e70a3405f0a605e34f
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/5962
Reviewed-by: Marcel van Lohuizen <mpvl@google.com>
diff --git a/cue/ast.go b/cue/ast.go
index 9fd68e7..1738463 100644
--- a/cue/ast.go
+++ b/cue/ast.go
@@ -600,6 +600,14 @@
 		}
 
 		f := v.label(name, true)
+		if _, ok := n.Node.(*ast.ImportSpec); ok {
+			n2 := v.mapScope(n.Node)
+			ref := &nodeRef{baseValue: newExpr(n), node: n2, label: f}
+			ret = ref
+			break
+		}
+
+		// TODO: probably unused. Verify and remove.
 		if n.Scope == nil {
 			// Package or direct ancestor node.
 			n2 := v.mapScope(n.Node)
diff --git a/cue/ast/astutil/apply.go b/cue/ast/astutil/apply.go
index fe36a9c..fa0572d 100644
--- a/cue/ast/astutil/apply.go
+++ b/cue/ast/astutil/apply.go
@@ -18,13 +18,9 @@
 	"encoding/hex"
 	"fmt"
 	"hash/fnv"
-	"path"
 	"reflect"
-	"strconv"
-	"strings"
 
 	"cuelang.org/go/cue/ast"
-	"cuelang.org/go/cue/token"
 )
 
 // A Cursor describes a node encountered during Apply.
@@ -130,10 +126,7 @@
 		return nil
 	}
 
-	name := path.Base(importPath)
-	if p := strings.LastIndexByte(name, ':'); p > 0 {
-		name = name[p+1:]
-	}
+	name := importPathName(importPath)
 
 	// TODO: come up with something much better.
 	// For instance, hoist the uniquer form cue/export.go to
@@ -141,51 +134,10 @@
 	hash := fnv.New32()
 	name += hex.EncodeToString(hash.Sum([]byte(importPath)))[:6]
 
-	quoted := strconv.Quote(importPath)
-
-	var imports *ast.ImportDecl
-	var spec *ast.ImportSpec
-	decls := info.current.decls
-	i := 0
-outer:
-	for ; i < len(decls); i++ {
-		d := decls[i]
-		switch t := d.(type) {
-		default:
-			break outer
-
-		case *ast.Package:
-		case *ast.CommentGroup:
-		case *ast.ImportDecl:
-			imports = t
-			for _, s := range t.Specs {
-				if s.Path.Value != quoted ||
-					s.Name == nil ||
-					s.Name.Name != name {
-					continue
-				}
-				spec = s
-				break
-			}
-		}
-	}
-
-	if spec == nil {
-		// Import not found, add one.
-		if imports == nil {
-			imports = &ast.ImportDecl{}
-			a := append(append(decls[:i], imports), decls[i:]...)
-			decls = a
-			info.current.decls = decls
-		}
-		path := ast.NewLit(token.STRING, quoted)
-		spec = &ast.ImportSpec{
-			Name: ast.NewIdent(name),
-			Path: path,
-		}
-		imports.Specs = append(imports.Specs, spec)
-		ast.SetRelPos(imports.Specs[0], token.NoRelPos)
-	}
+	spec := insertImport(&info.current.decls, &ast.ImportSpec{
+		Name: ast.NewIdent(name),
+		Path: ast.NewString(importPath),
+	})
 
 	ident := &ast.Ident{Node: spec} // Name is set later.
 	info.importPatch = append(info.importPatch, ident)
diff --git a/cue/ast/astutil/resolve.go b/cue/ast/astutil/resolve.go
index 84fcb90..9ddd8f7 100644
--- a/cue/ast/astutil/resolve.go
+++ b/cue/ast/astutil/resolve.go
@@ -27,17 +27,59 @@
 // An ErrFunc processes errors.
 type ErrFunc func(pos token.Pos, msg string, args ...interface{})
 
+// TODO: future development
+//
+// Resolution currently assigns values along the table below. This is based on
+// Go's resolver and is not quite convenient for CUE's purposes. For one, CUE
+// allows manually setting resolution and than call astutil.Sanitize to
+// normalize the ast.File. Manually assigning resolutions according to the
+// below table is rather tedious though.
+//
+// Instead of using the Scope and Node fields in identifiers, we suggest the
+// following assignments:
+//
+//    Reference Node // an Decl or Clause
+//    Ident *Ident   // The identifier in References (optional)
+//
+// References always refers to the direct element in the scope in which the
+// identifier occurs, not the final value, so: *Field, *LetClause, *ForClause,
+// etc. In case Ident is defined, it must be the same pointer as the
+// referencing identifier. In case it is not defined, the Name of the
+// referencing identifier can be used to locate the proper identifier in the
+// referenced node.
+//
+// The Scope field in the original design then loses its function.
+//
+// Type of reference      Scope          Node
+// Let Clause             File/Struct    LetClause
+// Alias declaration      File/Struct    Alias (deprecated)
+// Illegal Reference      File/Struct
+// Fields
+//   X in X: y            File/Struct    Expr (y)
+//   X in X=x: y          File/Struct    Field
+//   X in X="\(x)": y     File/Struct    Field
+//   X in [X=x]: y        Field          Expr (x)
+//   X in X=[x]: y        Field          Field
+//
+// for k, v in            Field          ForClause
+//
+// Template               Field          Template
+// Fields inside lambda
+//    Label               Field          Expr
+//    Value               Field          Field
+// Pkg                    nil            ImportSpec
+
 // Resolve resolves all identifiers in a file. Unresolved identifiers are
 // recorded in Unresolved. It will not overwrite already resolved values.
 func Resolve(f *ast.File, errFn ErrFunc) {
-	walk(&scope{errFn: errFn}, f)
+	walk(&scope{errFn: errFn, identFn: resolveIdent}, f)
 }
 
 // Resolve resolves all identifiers in an expression.
 // It will not overwrite already resolved values.
 func ResolveExpr(e ast.Expr, errFn ErrFunc) {
 	f := &ast.File{}
-	walk(&scope{file: f, errFn: errFn}, e)
+	walk(&scope{file: f, errFn: errFn, identFn: resolveIdent}, e)
 }
 
 // A Scope maintains the set of named language entities declared
@@ -48,20 +90,29 @@
 	file    *ast.File
 	outer   *scope
 	node    ast.Node
-	index   map[string]ast.Node
+	index   map[string]entry
 	inField bool
 
-	errFn func(p token.Pos, msg string, args ...interface{})
+	identFn func(s *scope, n *ast.Ident) bool
+	nameFn  func(name string)
+	errFn   func(p token.Pos, msg string, args ...interface{})
+}
+
+type entry struct {
+	node ast.Node
+	link ast.Node // Alias, LetClause, or Field
 }
 
 func newScope(f *ast.File, outer *scope, node ast.Node, decls []ast.Decl) *scope {
 	const n = 4 // initial scope capacity
 	s := &scope{
-		file:  f,
-		outer: outer,
-		node:  node,
-		index: make(map[string]ast.Node, n),
-		errFn: outer.errFn,
+		file:    f,
+		outer:   outer,
+		node:    node,
+		index:   make(map[string]entry, n),
+		identFn: outer.identFn,
+		nameFn:  outer.nameFn,
+		errFn:   outer.errFn,
 	}
 	for _, d := range decls {
 		switch x := d.(type) {
@@ -72,7 +123,7 @@
 				// TODO(legacy): use name := a.Ident.Name once quoted
 				// identifiers are no longer supported.
 				if name, _, _ := ast.LabelName(a.Ident); name != "" {
-					s.insert(name, x)
+					s.insert(name, x, a)
 				}
 				label, _ = a.Expr.(ast.Label)
 			}
@@ -85,26 +136,30 @@
 					break
 				}
 				if a, ok := y.Elts[0].(*ast.Alias); ok {
-					s.insert(a.Ident.Name, x)
+					s.insert(a.Ident.Name, x, a)
 				}
 			}
 
 			// default:
 			name, isIdent, _ := ast.LabelName(label)
 			if isIdent {
-				s.insert(name, x.Value)
+				s.insert(name, x.Value, x)
 			}
 		case *ast.LetClause:
 			name, isIdent, _ := ast.LabelName(x.Ident)
 			if isIdent {
-				s.insert(name, x)
+				s.insert(name, x, x)
 			}
 		case *ast.Alias:
 			name, isIdent, _ := ast.LabelName(x.Ident)
 			if isIdent {
-				s.insert(name, x)
+				s.insert(name, x, x)
 			}
-			// Handle imports
+		case *ast.ImportDecl:
+			for _, spec := range x.Specs {
+				info, _ := ParseImportSpec(spec)
+				s.insert(info.Ident, spec, spec)
+			}
 		}
 	}
 	return s
@@ -115,32 +170,48 @@
 		return true
 	}
 	switch n.(type) {
-	case *ast.LetClause:
-		return true
-
-	case *ast.Alias:
-		return true
-
-	case *ast.Field:
+	case *ast.LetClause, *ast.Alias, *ast.Field:
 		return true
 	}
 	return false
 }
 
-func (s *scope) insert(name string, n ast.Node) {
+func (s *scope) mustBeUnique(n ast.Node) bool {
+	if _, ok := s.node.(*ast.Field); ok {
+		return true
+	}
+	switch n.(type) {
+	// TODO: add *ast.ImportSpec when some implementations are moved over to
+	// Sanitize.
+	case *ast.ImportSpec, *ast.LetClause, *ast.Alias, *ast.Field:
+		return true
+	}
+	return false
+}
+
+func (s *scope) insert(name string, n, link ast.Node) {
 	if name == "" {
 		return
 	}
+	if s.nameFn != nil {
+		s.nameFn(name)
+	}
 	// TODO: record both positions.
-	if outer, _, existing := s.lookup(name); existing != nil {
-		isAlias1 := s.isLet(n)
-		isAlias2 := outer.isLet(existing)
-		if isAlias1 != isAlias2 {
+	if outer, _, existing := s.lookup(name); existing.node != nil {
+		if s.isLet(n) != outer.isLet(existing.node) {
 			s.errFn(n.Pos(), "cannot have both alias and field with name %q in same scope", name)
 			return
-		} else if isAlias1 || isAlias2 {
+		} else if s.mustBeUnique(n) || outer.mustBeUnique(existing.node) {
 			if outer == s {
-				s.errFn(n.Pos(), "alias %q redeclared in same scope", name)
+				if _, ok := existing.node.(*ast.ImportSpec); ok {
+					return
+					// TODO:
+					s.errFn(n.Pos(), "conflicting declaration %s\n"+
+						"\tprevious declaration at %s",
+						name, existing.node.Pos())
+				} else {
+					s.errFn(n.Pos(), "alias %q redeclared in same scope", name)
+				}
 				return
 			}
 			// TODO: Should we disallow shadowing of aliases?
@@ -149,37 +220,37 @@
 			// s.errFn(n.Pos(), "alias %q already declared in enclosing scope", name)
 		}
 	}
-	s.index[name] = n
+	s.index[name] = entry{node: n, link: link}
 }
 
-func (s *scope) resolveScope(name string, node ast.Node) (scope ast.Node, ok bool) {
+func (s *scope) resolveScope(name string, node ast.Node) (scope ast.Node, e entry, ok bool) {
 	last := s
 	for s != nil {
-		if n, ok := s.index[name]; ok && node == n {
-			if last.node == n {
-				return nil, true
+		if n, ok := s.index[name]; ok && node == n.node {
+			if last.node == n.node {
+				return nil, n, true
 			}
-			return s.node, true
+			return s.node, n, true
 		}
 		s, last = s.outer, s
 	}
-	return nil, false
+	return nil, entry{}, false
 }
 
-func (s *scope) lookup(name string) (p *scope, obj, node ast.Node) {
+func (s *scope) lookup(name string) (p *scope, obj ast.Node, node entry) {
 	// TODO(#152): consider returning nil for obj if it is a reference to root.
 	// last := s
 	for s != nil {
 		if n, ok := s.index[name]; ok {
-			// if last.node == n {
-			// 	return nil, n
-			// }
+			if _, ok := n.node.(*ast.ImportSpec); ok {
+				return s, nil, n
+			}
 			return s, s.node, n
 		}
 		// s, last = s.outer, s
 		s = s.outer
 	}
-	return nil, nil, nil
+	return nil, nil, entry{}
 }
 
 func (s *scope) After(n ast.Node) {}
@@ -220,7 +291,7 @@
 			s = newScope(s.file, s, x, nil)
 			if alias != nil {
 				if name, _, _ := ast.LabelName(alias.Ident); name != "" {
-					s.insert(name, x)
+					s.insert(name, x, alias)
 				}
 			}
 
@@ -237,7 +308,7 @@
 				// to detect illegal usage, though.
 				name, err := ast.ParseIdent(a.Ident)
 				if err == nil {
-					s.insert(name, a.Expr)
+					s.insert(name, a.Expr, a)
 				}
 			}
 
@@ -257,7 +328,7 @@
 			s = newScope(s.file, s, x, nil)
 			name, err := ast.ParseIdent(label.Ident)
 			if err == nil {
-				s.insert(name, x.Label) // Field used for entire lambda.
+				s.insert(name, x.Label, x) // Field used for entire lambda.
 			}
 		}
 
@@ -305,35 +376,41 @@
 		return nil
 
 	case *ast.Ident:
-		name, ok, _ := ast.LabelName(x)
-		if !ok {
-			// TODO: generate error
-			break
+		if s.identFn(s, x) {
+			return nil
 		}
-		if _, obj, node := s.lookup(name); node != nil {
-			switch {
-			case x.Node == nil:
-				x.Node = node
-				x.Scope = obj
-
-			case x.Node == node:
-				x.Scope = obj
-
-			default: // x.Node != node
-				scope, ok := s.resolveScope(name, x.Node)
-				if !ok {
-					s.file.Unresolved = append(s.file.Unresolved, x)
-				}
-				x.Scope = scope
-			}
-		} else {
-			s.file.Unresolved = append(s.file.Unresolved, x)
-		}
-		return nil
 	}
 	return s
 }
 
+func resolveIdent(s *scope, x *ast.Ident) bool {
+	name, ok, _ := ast.LabelName(x)
+	if !ok {
+		// TODO: generate error
+		return false
+	}
+	if _, obj, node := s.lookup(name); node.node != nil {
+		switch {
+		case x.Node == nil:
+			x.Node = node.node
+			x.Scope = obj
+
+		case x.Node == node.node:
+			x.Scope = obj
+
+		default: // x.Node != node
+			scope, _, ok := s.resolveScope(name, x.Node)
+			if !ok {
+				s.file.Unresolved = append(s.file.Unresolved, x)
+			}
+			x.Scope = scope
+		}
+	} else {
+		s.file.Unresolved = append(s.file.Unresolved, x)
+	}
+	return true
+}
+
 func scopeClauses(s *scope, clauses []ast.Clause) *scope {
 	for _, c := range clauses {
 		if f, ok := c.(*ast.ForClause); ok { // TODO(let): support let clause
@@ -342,12 +419,12 @@
 			if f.Key != nil {
 				name, err := ast.ParseIdent(f.Key)
 				if err == nil {
-					s.insert(name, f.Key)
+					s.insert(name, f.Key, f)
 				}
 			}
 			name, err := ast.ParseIdent(f.Value)
 			if err == nil {
-				s.insert(name, f.Value)
+				s.insert(name, f.Value, f)
 			}
 		} else {
 			walk(s, c)
diff --git a/cue/ast/astutil/util.go b/cue/ast/astutil/util.go
index cbd0629..e3439b8 100644
--- a/cue/ast/astutil/util.go
+++ b/cue/ast/astutil/util.go
@@ -14,7 +14,63 @@
 
 package astutil
 
-import "cuelang.org/go/cue/ast"
+import (
+	"path"
+	"strconv"
+	"strings"
+
+	"cuelang.org/go/cue/ast"
+	"cuelang.org/go/cue/token"
+)
+
+// importPathName derives the name from the given import path.
+//
+// Examples:
+//      string           string
+//      foo.com/bar      bar
+//      foo.com/bar:baz  baz
+//
+func importPathName(id string) string {
+	name := path.Base(id)
+	if p := strings.LastIndexByte(name, ':'); p > 0 {
+		name = name[p+1:]
+	}
+	return name
+}
+
+// ImportInfo describes the information contained in an ImportSpec.
+type ImportInfo struct {
+	Ident   string // identifier used to refer to the import
+	PkgName string // name of the package
+	ID      string // full import path, including the name
+	Dir     string // import path, excluding the name
+}
+
+// ParseImportSpec returns the name and full path of an ImportSpec.
+func ParseImportSpec(spec *ast.ImportSpec) (info ImportInfo, err error) {
+	str, err := strconv.Unquote(spec.Path.Value)
+	if err != nil {
+		return info, err
+	}
+
+	info.ID = str
+
+	if p := strings.LastIndexByte(str, ':'); p > 0 {
+		info.Dir = str[:p]
+		info.PkgName = str[p+1:]
+	} else {
+		info.Dir = str
+		info.PkgName = path.Base(str)
+	}
+
+	if spec.Name != nil {
+		info.Ident = spec.Name.Name
+	} else {
+		info.Ident = info.PkgName
+	}
+
+	return info, nil
+}
 
 // CopyComments associates comments of one node with another.
 // It may change the relative position of comments.
@@ -43,3 +99,54 @@
 	ast.SetPos(to, from.Pos())
 	return to
 }
+
+// insertImport looks up an existing import with the given name and path or will
+// add spec if it doesn't exist. It returns a spec in decls matching spec.
+func insertImport(decls *[]ast.Decl, spec *ast.ImportSpec) *ast.ImportSpec {
+	x, _ := ParseImportSpec(spec)
+
+	a := *decls
+
+	var imports *ast.ImportDecl
+	var orig *ast.ImportSpec
+	i := 0
+outer:
+	for ; i < len(a); i++ {
+		d := a[i]
+		switch t := d.(type) {
+		default:
+			break outer
+
+		case *ast.Package:
+		case *ast.CommentGroup:
+		case *ast.ImportDecl:
+			imports = t
+			for _, s := range t.Specs {
+				y, _ := ParseImportSpec(s)
+				if y.ID != x.ID {
+					continue
+				}
+				orig = s
+				if x.Ident == "" || y.Ident == x.Ident {
+					return s
+				}
+			}
+		}
+	}
+
+	// Import not found, add one.
+	if imports == nil {
+		imports = &ast.ImportDecl{}
+		preamble := append(a[:i:i], imports)
+		a = append(preamble, a[i:]...)
+		*decls = a
+	}
+
+	if orig != nil {
+		CopyComments(spec, orig)
+	}
+	imports.Specs = append(imports.Specs, spec)
+	ast.SetRelPos(imports.Specs[0], token.NoRelPos)
+
+	return spec
+}
diff --git a/cue/ast/astutil/util_test.go b/cue/ast/astutil/util_test.go
new file mode 100644
index 0000000..c66f092
--- /dev/null
+++ b/cue/ast/astutil/util_test.go
@@ -0,0 +1,74 @@
+// Copyright 2020 CUE Authors
+//
+// 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 astutil
+
+import (
+	"path"
+	"testing"
+
+	"cuelang.org/go/cue/ast"
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestImportInfo(t *testing.T) {
+	testCases := []struct {
+		name string
+		path string
+		want ImportInfo
+	}{
+		{"", "a.b/bar", ImportInfo{
+			Ident:   "bar",
+			PkgName: "bar",
+			ID:      "a.b/bar",
+			Dir:     "a.b/bar",
+		}},
+		{"foo", "a.b/bar", ImportInfo{
+			Ident:   "foo",
+			PkgName: "bar",
+			ID:      "a.b/bar",
+			Dir:     "a.b/bar",
+		}},
+		{"", "a.b/bar:foo", ImportInfo{
+			Ident:   "foo",
+			PkgName: "foo",
+			ID:      "a.b/bar:foo",
+			Dir:     "a.b/bar",
+		}},
+		{"", "strings", ImportInfo{
+			Ident:   "strings",
+			PkgName: "strings",
+			ID:      "strings",
+			Dir:     "strings",
+		}},
+	}
+	for _, tc := range testCases {
+		t.Run(path.Join(tc.name, tc.path), func(t *testing.T) {
+			var ident *ast.Ident
+			if tc.name != "" {
+				ident = ast.NewIdent(tc.name)
+			}
+			got, err := ParseImportSpec(&ast.ImportSpec{
+				Name: ident,
+				Path: ast.NewString(tc.path),
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !cmp.Equal(got, tc.want) {
+				t.Error(cmp.Diff(got, tc.want))
+			}
+		})
+	}
+}
diff --git a/cue/ast_test.go b/cue/ast_test.go
index 966135e3..40f5fae 100644
--- a/cue/ast_test.go
+++ b/cue/ast_test.go
@@ -613,13 +613,30 @@
 	}{{
 		file: &ast.File{Decls: []ast.Decl{
 			&ast.ImportDecl{Specs: []*ast.ImportSpec{spec}},
+			&ast.EmbedDecl{
+				Expr: ast.NewStruct(
+					&ast.Field{
+						Label: mustParseExpr(`list`).(*ast.Ident),
+						Value: ast.NewCall(
+							ast.NewSel(
+								&ast.Ident{Name: "list", Node: spec},
+								"Min")),
+					})},
+		}},
+		want: "import \"list\", let LIST=list, {list: LIST.Min()}",
+	}, {
+		file: &ast.File{Decls: []ast.Decl{
+			&ast.ImportDecl{Specs: []*ast.ImportSpec{spec}},
 			&ast.Field{
-				Label: mustParseExpr(`list`).(*ast.Ident),
-				Value: ast.NewCall(
-					ast.NewSel(&ast.Ident{Name: "list", Node: spec}, "Min")),
+				Label: ast.NewIdent("a"),
+				Value: ast.NewStruct(&ast.Field{
+					Label: mustParseExpr(`list`).(*ast.Ident),
+					Value: ast.NewCall(
+						ast.NewSel(&ast.Ident{Name: "list", Node: spec}, "Min")),
+				}),
 			},
 		}},
-		want: "import listx \"list\", list: listx.Min()",
+		want: "import \"list\", let LIST=list, a: {list: LIST.Min()}",
 	}}
 	for _, tc := range testCases {
 		t.Run("", func(t *testing.T) {
diff --git a/cue/format/testdata/imports.golden b/cue/format/testdata/imports.golden
index abd15e4..5d67085 100644
--- a/cue/format/testdata/imports.golden
+++ b/cue/format/testdata/imports.golden
@@ -7,7 +7,7 @@
 )
 
 import (
-	"time"
+	time1 "time"
 
 	// comment f2
 	f2 "cuelang.org/go/foo"
@@ -15,10 +15,10 @@
 )
 
 import (
-	"time"
+	time2 "time"
 
-	same "cuelang.org/go/foo" // comment 2
-	same "cuelang.org/go/foo" // comment 1
+	same "cuelang.org/go/foo"  // comment 1
+	same2 "cuelang.org/go/foo" // comment 2
 )
 
 a: time.time
diff --git a/cue/format/testdata/imports.input b/cue/format/testdata/imports.input
index 3d526ca..a617796 100644
--- a/cue/format/testdata/imports.input
+++ b/cue/format/testdata/imports.input
@@ -7,7 +7,7 @@
 )
 
 import (
-    "time"
+    time1 "time"
 
     // comment f2
     f2 "cuelang.org/go/foo"
@@ -15,10 +15,10 @@
 )
 
 import (
-    "time"
+    time2 "time"
 
-    same "cuelang.org/go/foo" // comment 2
     same "cuelang.org/go/foo" // comment 1
+    same2 "cuelang.org/go/foo" // comment 2
 )