fix parse array with the external const correctly (golang#569)

Fixes golang#561 - Panic on array parsing if length is externally defined It also fixes the following minor problems in TestParseArrayWithConstLength, a test function related to this feature. It doesn't check all outputs to be asserted. It doesn't initialize some of the fields in mockgen.fileParser, even though they are initialized in the production code. (Actually, it caused an unnecessary test error by test data I added)
ybbus · Jun 11, 2021 · aba2ff9 · aba2ff9
1 parent bb196fc
commit aba2ff9
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 7 deletions.
diff --git a/mockgen/internal/tests/const_array_length/input.go b/mockgen/internal/tests/const_array_length/input.go
@@ -1,10 +1,13 @@
 package const_length
 
+import "math"
+
 //go:generate mockgen -package const_length -destination mock.go -source input.go
 
 const C = 2
 
 type I interface {
 	Foo() [C]int
 	Bar() [2]int
+	Baz() [math.MaxInt8]int
 }
diff --git a/mockgen/internal/tests/const_array_length/mock.go b/mockgen/internal/tests/const_array_length/mock.go
diff --git a/mockgen/parse.go b/mockgen/parse.go
@@ -22,8 +22,10 @@ import (
 	"fmt"
 	"go/ast"
 	"go/build"
+	"go/importer"
 	"go/parser"
 	"go/token"
+	"go/types"
 	"io/ioutil"
 	"log"
 	"path"
@@ -409,8 +411,19 @@ func (p *fileParser) parseType(pkg string, typ ast.Expr) (model.Type, error) {
 			case (*ast.BasicLit):
 				value = val.Value
 			case (*ast.Ident):
-				// when the length is a const
+				// when the length is a const defined locally
 				value = val.Obj.Decl.(*ast.ValueSpec).Values[0].(*ast.BasicLit).Value
+			case (*ast.SelectorExpr):
+				// when the length is a const defined in an external package
+				usedPkg, err := importer.Default().Import(fmt.Sprintf("%s", val.X))
+				if err != nil {
+					return nil, p.errorf(v.Len.Pos(), "unknown package in array length: %v", err)
+				}
+				ev, err := types.Eval(token.NewFileSet(), usedPkg, token.NoPos, val.Sel.Name)
+				if err != nil {
+					return nil, p.errorf(v.Len.Pos(), "unknown constant in array length: %v", err)
+				}
+				value = ev.Value.String()
 			}
 
 			x, err := strconv.Atoi(value)

diff --git a/mockgen/parse_test.go b/mockgen/parse_test.go
@@ -116,24 +116,31 @@ func Benchmark_parseFile(b *testing.B) {
 
 func TestParseArrayWithConstLength(t *testing.T) {
 	fs := token.NewFileSet()
+	srcDir := "internal/tests/const_array_length/input.go"
 
-	file, err := parser.ParseFile(fs, "internal/tests/const_array_length/input.go", nil, 0)
+	file, err := parser.ParseFile(fs, srcDir, nil, 0)
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
 
 	p := fileParser{
-		fileSet: fs,
+		fileSet:            fs,
+		imports:            make(map[string]importedPackage),
+		importedInterfaces: make(map[string]map[string]*ast.InterfaceType),
+		auxInterfaces:      make(map[string]map[string]*ast.InterfaceType),
+		srcDir:             srcDir,
 	}
 
 	pkg, err := p.parseFile("", file)
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
 
-	expect := "[2]int"
-	got := pkg.Interfaces[0].Methods[0].Out[0].Type.String(nil, "")
-	if got != expect {
-		t.Fatalf("got %v; expected %v", got, expect)
+	expects := []string{"[2]int", "[2]int", "[127]int"}
+	for i, e := range expects {
+		got := pkg.Interfaces[0].Methods[i].Out[0].Type.String(nil, "")
+		if got != e {
+			t.Fatalf("got %v; expected %v", got, e)
+		}
 	}
 }