Enhance pluginrpc-gen parser

Now handles `package.Type` and `*package.Type`
Fixes parsing issues with slice and map types.

Signed-off-by: Brian Goff <cpuguy83@gmail.com>

plugins/pluginrpc-gen/parser.go

@@ -6,7 +6,9 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
+	"path"
 	"reflect"
+	"strings"
 )
 
 var errBadReturn = errors.New("found return arg with no name: all args must be named")
@@ -25,6 +27,7 @@ func (e errUnexpectedType) Error() string {
 type ParsedPkg struct {
 	Name      string
 	Functions []function
+	Imports   []importSpec
 }
 
 type function struct {
@@ -35,14 +38,29 @@ type function struct {
 }
 
 type arg struct {
-	Name    string
-	ArgType string
+	Name            string
+	ArgType         string
+	PackageSelector string
 }
 
 func (a *arg) String() string {
 	return a.Name + " " + a.ArgType
 }
 
+type importSpec struct {
+	Name string
+	Path string
+}
+
+func (s *importSpec) String() string {
+	var ss string
+	if len(s.Name) != 0 {
+		ss += s.Name
+	}
+	ss += s.Path
+	return ss
+}
+
 // Parse parses the given file for an interface definition with the given name.
 func Parse(filePath string, objName string) (*ParsedPkg, error) {
 	fs := token.NewFileSet()
@@ -73,6 +91,44 @@ func Parse(filePath string, objName string) (*ParsedPkg, error) {
 		return nil, err
 	}
 
+	// figure out what imports will be needed
+	imports := make(map[string]importSpec)
+	for _, f := range p.Functions {
+		args := append(f.Args, f.Returns...)
+		for _, arg := range args {
+			if len(arg.PackageSelector) == 0 {
+				continue
+			}
+
+			for _, i := range pkg.Imports {
+				if i.Name != nil {
+					if i.Name.Name != arg.PackageSelector {
+						continue
+					}
+					imports[i.Path.Value] = importSpec{Name: arg.PackageSelector, Path: i.Path.Value}
+					break
+				}
+
+				_, name := path.Split(i.Path.Value)
+				splitName := strings.Split(name, "-")
+				if len(splitName) > 1 {
+					name = splitName[len(splitName)-1]
+				}
+				// import paths have quotes already added in, so need to remove them for name comparison
+				name = strings.TrimPrefix(name, `"`)
+				name = strings.TrimSuffix(name, `"`)
+				if name == arg.PackageSelector {
+					imports[i.Path.Value] = importSpec{Path: i.Path.Value}
+					break
+				}
+			}
+		}
+	}
+
+	for _, spec := range imports {
+		p.Imports = append(p.Imports, spec)
+	}
+
 	return p, nil
 }
 
@@ -142,22 +198,66 @@ func parseArgs(fields []*ast.Field) ([]arg, error) {
 			return nil, errBadReturn
 		}
 		for _, name := range f.Names {
-			var typeName string
-			switch argType := f.Type.(type) {
-			case *ast.Ident:
-				typeName = argType.Name
-			case *ast.StarExpr:
-				i, ok := argType.X.(*ast.Ident)
-				if !ok {
-					return nil, errUnexpectedType{"*ast.Ident", f.Type}
-				}
-				typeName = "*" + i.Name
-			default:
-				return nil, errUnexpectedType{"*ast.Ident or *ast.StarExpr", f.Type}
+			p, err := parseExpr(f.Type)
+			if err != nil {
+				return nil, err
 			}
-
-			args = append(args, arg{name.Name, typeName})
+			args = append(args, arg{name.Name, p.value, p.pkg})
 		}
 	}
 	return args, nil
 }
+
+type parsedExpr struct {
+	value string
+	pkg   string
+}
+
+func parseExpr(e ast.Expr) (parsedExpr, error) {
+	var parsed parsedExpr
+	switch i := e.(type) {
+	case *ast.Ident:
+		parsed.value += i.Name
+	case *ast.StarExpr:
+		p, err := parseExpr(i.X)
+		if err != nil {
+			return parsed, err
+		}
+		parsed.value += "*"
+		parsed.value += p.value
+		parsed.pkg = p.pkg
+	case *ast.SelectorExpr:
+		p, err := parseExpr(i.X)
+		if err != nil {
+			return parsed, err
+		}
+		parsed.pkg = p.value
+		parsed.value += p.value + "."
+		parsed.value += i.Sel.Name
+	case *ast.MapType:
+		parsed.value += "map["
+		p, err := parseExpr(i.Key)
+		if err != nil {
+			return parsed, err
+		}
+		parsed.value += p.value
+		parsed.value += "]"
+		p, err = parseExpr(i.Value)
+		if err != nil {
+			return parsed, err
+		}
+		parsed.value += p.value
+		parsed.pkg = p.pkg
+	case *ast.ArrayType:
+		parsed.value += "[]"
+		p, err := parseExpr(i.Elt)
+		if err != nil {
+			return parsed, err
+		}
+		parsed.value += p.value
+		parsed.pkg = p.pkg
+	default:
+		return parsed, errUnexpectedType{"*ast.Ident or *ast.StarExpr", i}
+	}
+	return parsed, nil
+}