go/go2go: support instantiating with an inferred local type

Add a graph example.

Change-Id: Iec533487f2757b9fdc883e4a57715ad993e11254

src/cmd/go2go/testdata/go2path/src/graph/graph.go2

@@ -0,0 +1,65 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package graph implements general purpose graph algorithms.
+package graph
+
+import "errors"
+
+// A Graph is a collection of nodes. A node may have an arbitrary number
+// of edges. An edge connects two nodes. Both nodes and edges must be
+// comparable. This is an undirected simple graph.
+type Graph(type Node, Edge G) struct {
+	nodes []Node
+}
+
+// G is the contract that the Graph Node and Edge types must implement.
+contract G(Node, Edge) {
+	Node Edges() []Edge
+	comparable(Node)
+	Edge Nodes() (a, b Node)
+	comparable(Edge)
+}
+
+// New creates a new Graph from a collection of Nodes.
+func New(type Node, Edge G)(nodes []Node) *Graph(Node, Edge) {
+	return &Graph(Node, Edge){nodes: nodes}
+}
+
+type nodePath(type Node, Edge G) struct {
+	node Node
+	path []Edge
+}
+
+// ShortestPath returns the shortest path between two nodes,
+// as an ordered list of edges. If there are multiple shortest paths,
+// which one is returned is unpredictable.
+func (g *Graph(Node, Edge)) ShortestPath(from, to Node) ([]Edge, error) {
+	visited := make(map[Node]bool)
+	visited[from] = true
+	workqueue := [](nodePath(Node, Edge)){nodePath(Node, Edge){from, nil}}
+	for len(workqueue) > 0 {
+		current := workqueue
+		workqueue = nil
+		for _, np := range current {
+			edges := np.node.Edges()
+			for _, edge := range edges {
+				a, b := edge.Nodes()
+				if a == np.node {
+					a = b
+				}
+				if !visited[a] {
+					ve := append([]Edge(nil), np.path...)
+					ve = append(ve, edge)
+					if a == to {
+						return ve, nil
+					}
+					workqueue = append(workqueue, nodePath(Node, Edge){a, ve})
+					visited[a] = true
+				}
+			}
+		}
+	}
+	return nil, errors.New("no path")
+}

src/cmd/go2go/testdata/go2path/src/graph/graphs_test.go2

@@ -0,0 +1,143 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package graph
+
+import (
+	"fmt"
+	"testing"
+
+	"slices"
+)
+
+type direction int
+
+const (
+	north direction = iota
+	ne
+	east
+	se
+	south
+	sw
+	west
+	nw
+	up
+	down
+)
+
+func (dir direction) String() string {
+	strs := map[direction]string {
+		north: "north",
+		ne: "ne",
+		east: "east",
+		se: "se",
+		south: "south",
+		sw: "sw",
+		west: "west",
+		nw: "nw",
+		up: "up",
+		down: "down",
+	}
+	if str, ok := strs[dir]; ok {
+		return str
+	}
+	return fmt.Sprintf("direction %d", dir)
+}
+
+type mazeRoom struct {
+	index int
+	exits [10]int
+}
+
+type mazeEdge struct {
+	from, to int
+	dir   direction
+}
+
+// Edges returns the exits from the room.
+func (m mazeRoom) Edges() []mazeEdge {
+	var r []mazeEdge
+	for i, exit := range m.exits {
+		if exit != 0 {
+			r = append(r, mazeEdge{
+				from: m.index,
+				to: exit,
+				dir: direction(i),
+			})
+		}
+	}
+	return r
+}
+
+// Nodes returns the rooms connected by an edge.
+func (e mazeEdge) Nodes() (mazeRoom, mazeRoom) {
+	m1, ok := zork[e.from]
+	if !ok {
+		panic("bad edge")
+	}
+	m2, ok := zork[e.to]
+	if !ok {
+		panic("bad edge")
+	}
+	return m1, m2
+}
+
+// The first maze in Zork. Room indexes based on original Fortran data file.
+// You are in a maze of twisty little passages, all alike.
+var zork = map[int]mazeRoom{
+	11: { exits: [10]int{north: 11, south: 12, east: 14} }, // west to Troll Room
+	12: { exits: [10]int{south: 11, north: 14, east: 13} },
+	13: { exits: [10]int{west: 12, north: 14, up: 16 } },
+	14: { exits: [10]int{west: 13, north: 11, east: 15 } },
+	15: { exits: [10]int{south: 14 } }, // Dead End
+	16: { exits: [10]int{east: 17, north: 13, sw: 18 } }, // skeleton, etc.
+	17: { exits: [10]int{west: 16 } }, // Dead End
+	18: { exits: [10]int{down: 16, east: 19, west: 18, up: 22 } },
+	19: { exits: [10]int{up: 29, west: 18, ne: 15, east: 20, south: 30 } },
+	20: { exits: [10]int{ne: 19, west: 20, se: 21 } },
+	21: { exits: [10]int{north: 20 } }, // Dead End
+	22: { exits: [10]int{north: 18, east: 24, down: 23, south: 28, west: 26, nw: 22 } },
+	23: { exits: [10]int{east: 22, west: 28, up: 24 } },
+	24: { exits: [10]int{ne: 25, down: 23, nw: 28, sw: 26 } },
+	25: { exits: [10]int{sw: 24 } }, // Grating room (up to Clearing)
+	26: { exits: [10]int{west: 16, sw: 24, east: 28, up: 22, north: 27 } },
+	27: { exits: [10]int{south: 26 } }, // Dead End
+	28: { exits: [10]int{east: 22, down: 26, south: 23, west: 24 } },
+	29: { exits: [10]int{west: 30, nw: 29, ne: 19, south: 19 } },
+	30: { exits: [10]int{west: 29, south: 19 } }, // ne to Cyclops Room
+}
+
+func TestShortestPath(t *testing.T) {
+	// The Zork maze is not a proper undirected simple graph,
+	// as there are some one way paths (e.g., 19 -> 15),
+	// but for this test that doesn't matter.
+
+	// Set the index field in the map. Simpler than doing it in the
+	// composite literal.
+	for k := range zork {
+		r := zork[k]
+		r.index = k
+		zork[k] = r
+	}
+
+	var nodes []mazeRoom
+	for idx, room := range zork {
+		mridx := room
+		mridx.index = idx
+		nodes = append(nodes, mridx)
+	}
+	g := New(mazeRoom, mazeEdge)(nodes)
+	path, err := g.ShortestPath(zork[11], zork[30])
+	if err != nil {
+		t.Fatal(err)
+	}
+	var steps []direction
+	for _, edge := range path {
+		steps = append(steps, edge.dir)
+	}
+	want := []direction{east, west, up, sw, east, south}
+	if !slices.Equal(steps, want) {
+		t.Errorf("ShortestPath returned %v, want %v", steps, want)
+	}
+}

src/go/go2go/go2go.go

@@ -99,9 +99,9 @@ func rewriteFilesInPath(importer *Importer, importPath, dir string, go2files []s
 		tpkgs = append(tpkgs, pkgfiles)
 	}
 
-	for _, tpkg := range tpkgs {
-		for i, pkgfile := range tpkg {
-			if err := rewriteFile(dir, fset, importer, importPath, pkgfile.name, pkgfile.ast, i == 0); err != nil {
+	for i, tpkg := range tpkgs {
+		for j, pkgfile := range tpkg {
+			if err := rewriteFile(dir, fset, importer, importPath, rpkgs[i], pkgfile.name, pkgfile.ast, j == 0); err != nil {
 				return nil, err
 			}
 		}
@@ -124,11 +124,12 @@ func RewriteBuffer(importer *Importer, filename string, file []byte) ([]byte, er
 		Importer: importer,
 		Error:    merr.add,
 	}
-	if _, err := conf.Check(pf.Name.Name, fset, []*ast.File{pf}, importer.info); err != nil {
+	tpkg, err := conf.Check(pf.Name.Name, fset, []*ast.File{pf}, importer.info)
+	if err != nil {
 		return nil, fmt.Errorf("type checking failed for %s\n%v", pf.Name.Name, merr)
 	}
 	importer.addIDs(pf)
-	if err := rewriteAST(fset, importer, "", pf, true); err != nil {
+	if err := rewriteAST(fset, importer, "", tpkg, pf, true); err != nil {
 		return nil, err
 	}
 	var buf bytes.Buffer

src/go/go2go/rewrite.go

@@ -56,6 +56,7 @@ func isParameterizedTypeDecl(s ast.Spec) bool {
 type translator struct {
 	fset               *token.FileSet
 	importer           *Importer
+	tpkg               *types.Package
 	types              map[ast.Expr]types.Type
 	instantiations     map[string][]*instantiation
 	newDecls           []ast.Decl
@@ -80,8 +81,8 @@ type typeInstantiation struct {
 }
 
 // rewrite rewrites the contents of one file.
-func rewriteFile(dir string, fset *token.FileSet, importer *Importer, importPath, filename string, file *ast.File, addImportableName bool) (err error) {
-	if err := rewriteAST(fset, importer, importPath, file, addImportableName); err != nil {
+func rewriteFile(dir string, fset *token.FileSet, importer *Importer, importPath string, tpkg *types.Package, filename string, file *ast.File, addImportableName bool) (err error) {
+	if err := rewriteAST(fset, importer, importPath, tpkg, file, addImportableName); err != nil {
 		return err
 	}
 
@@ -109,10 +110,11 @@ func rewriteFile(dir string, fset *token.FileSet, importer *Importer, importPath
 }
 
 // rewriteAST rewrites the AST for a file.
-func rewriteAST(fset *token.FileSet, importer *Importer, importPath string, file *ast.File, addImportableName bool) (err error) {
+func rewriteAST(fset *token.FileSet, importer *Importer, importPath string, tpkg *types.Package, file *ast.File, addImportableName bool) (err error) {
 	t := translator{
 		fset:               fset,
 		importer:           importer,
+		tpkg:               tpkg,
 		types:              make(map[ast.Expr]types.Type),
 		instantiations:     make(map[string][]*instantiation),
 		typeInstantiations: make(map[types.Type][]*typeInstantiation),
@@ -726,11 +728,19 @@ func (t *translator) instantiationTypes(call *ast.CallExpr) (argList []ast.Expr,
 	} else {
 		for _, typ := range inferred.Targs {
 			arg := ast.NewIdent(typ.String())
-			if named, ok := typ.(*types.Named); ok && len(named.TArgs()) > 0 {
-				var narg *ast.Ident
-				typ, narg = t.lookupInstantiatedType(named)
-				if narg != nil {
-					arg = ast.NewIdent(narg.Name)
+			if named, ok := typ.(*types.Named); ok {
+				if len(named.TArgs()) > 0 {
+					var narg *ast.Ident
+					typ, narg = t.lookupInstantiatedType(named)
+					if narg != nil {
+						arg = ast.NewIdent(narg.Name)
+					}
+				}
+				if named.Obj().Pkg() == t.tpkg {
+					fields := strings.Split(arg.Name, ".")
+					if len(fields) > 1 {
+						arg = ast.NewIdent(fields[1])
+					}
 				}
 			}
 			typeList = append(typeList, typ)