Add type option

go.mod

@@ -7,18 +7,22 @@ require (
 	github.com/fatih/structtag v1.2.0
 	github.com/iancoleman/strcase v0.2.0 // indirect
 	github.com/lyft/protoc-gen-star v0.5.3 // indirect
+	github.com/stretchr/testify v1.6.1
 	golang.org/x/tools v0.1.7
 	google.golang.org/grpc/cmd/protoc-gen-go-grpc v1.1.0
 	google.golang.org/protobuf v1.27.1
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.0 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/spf13/afero v1.6.0 // indirect
 	golang.org/x/lint v0.0.0-20210508222113-6edffad5e616 // indirect
 	golang.org/x/mod v0.5.0 // indirect
 	golang.org/x/net v0.0.0-20210813160813-60bc85c4be6d // indirect
 	golang.org/x/sys v0.0.0-20210816183151-1e6c022a8912 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
 )

go.sum

@@ -90,6 +90,7 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.27.1 h1:SnqbnDw1V7RiZcXPx5MEeqPv2s79L9i7BJUlG/+RurQ=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=

patch/field_type.go

@@ -0,0 +1,216 @@
+package patch
+
+import (
+	"fmt"
+	"go/ast"
+	"go/token"
+	"go/types"
+	"log"
+	"strings"
+)
+
+func (p *Patcher) patchTypeDef(id *ast.Ident, obj types.Object) {
+	fieldType, ok := p.fieldTypes[obj]
+	if !ok {
+		return
+	}
+
+	castDecl := func(v *ast.Field) bool {
+		switch t := v.Type.(type) {
+		case *ast.Ident:
+			t.Name = fieldType
+			return true
+		case *ast.ArrayType:
+			v.Type = &ast.Ident{
+				Name: fieldType,
+			}
+			return true
+		case *ast.StarExpr:
+			t.X = &ast.Ident{
+				Name: fieldType,
+			}
+			return true
+		default:
+			return false
+		}
+	}
+
+	// Cast Field definition
+	if id.Obj != nil && id.Obj.Decl != nil {
+		v, ok := id.Obj.Decl.(*ast.Field)
+		if !ok {
+			log.Printf("Warning: fieldType declared for non-field object: %v `%s`", obj, fieldType)
+			return
+		}
+		if !castDecl(v) {
+			log.Printf("Warning: unsupported fieldType type: %T `%s`", v.Type, fieldType)
+		}
+		return
+	}
+	switch obj.Type().(type) {
+	// Cast Getter signature
+	case *types.Signature:
+		parent := p.findParentNode(id)
+		n, ok := parent.(*ast.FuncDecl)
+		if !ok {
+			log.Printf("Warning: unexpected type for getter: %v `%T`", obj, parent)
+			break
+		}
+		if l := len(n.Type.Results.List); l != 1 {
+			log.Printf("Warning: unexpected return count for getter: %v `%d`", obj, l)
+			return
+		}
+		if !castDecl(n.Type.Results.List[0]) {
+			log.Printf("Warning: unsupported fieldType type: %T `%s`", n.Type.Results.List[0].Type, fieldType)
+		}
+		return
+	}
+}
+
+func (p *Patcher) patchTypeUsage(id *ast.Ident, obj types.Object) {
+	desiredType, ok := p.fieldTypes[obj]
+	if !ok {
+		return
+	}
+	// we don't want to modify the getter body,
+	// so we check if we are in the getter declaration body ?
+	node := ast.Node(id)
+	for{
+		node = p.findParentNode(node)
+		if node == nil {
+			break
+		}
+		if fn, ok := node.(*ast.FuncDecl); ok {
+			if fn.Name.String() == "Get"+id.Name {
+				return
+			}
+			break
+		}
+	}
+	var originalType string
+	switch t := obj.Type().(type) {
+	case *types.Basic:
+		originalType = t.Name()
+	case *types.Pointer:
+		originalType = t.String()
+		desiredType = "*"+desiredType
+	case *types.Signature:
+		if t.Results().Len() != 1 {
+			return
+		}
+		originalType = t.Results().At(0).Type().String()
+	}
+	cast := func(as string, expr ast.Expr) ast.Expr {
+		if strings.HasPrefix(as, "*") {
+			as = fmt.Sprintf("(%s)", as)
+		}
+		return &ast.CallExpr{
+			Fun: &ast.Ident{
+				Name: as,
+			},
+			Args: []ast.Expr{expr},
+		}
+	}
+
+	expr := p.findParentNode(id)
+
+	if kv, ok := expr.(*ast.KeyValueExpr); ok {
+		if kv.Key == id {
+			kv.Value = cast(desiredType, kv.Value)
+			return
+		}
+		if kv.Value == id {
+			kv.Value = cast(originalType, kv.Value)
+			return
+		}
+		return
+	}
+
+	selectorExpr, ok := expr.(*ast.SelectorExpr)
+	if !ok {
+		return
+	}
+
+	var patch func(node ast.Node)
+	patch = func(node ast.Node) {
+		switch nodeExpr := node.(type) {
+		case *ast.StarExpr:
+			desiredType = strings.TrimPrefix(desiredType, "*")
+			originalType = strings.TrimPrefix(originalType, "*")
+			expr = node
+			patch(p.findParentNode(node))
+			return
+		case *ast.UnaryExpr:
+			if nodeExpr.Op != token.AND {
+				return
+			}
+			desiredType = "*" + desiredType
+			originalType = "*" + originalType
+			expr = node
+			patch(p.findParentNode(node))
+		case *ast.AssignStmt:
+			if len(nodeExpr.Lhs) != len(nodeExpr.Rhs) {
+				return
+			}
+			for i := range nodeExpr.Lhs {
+				if nodeExpr.Lhs[i] == expr {
+					nodeExpr.Rhs[i] = cast(desiredType, nodeExpr.Rhs[i])
+					return
+				}
+			}
+			for i := range nodeExpr.Rhs {
+				if nodeExpr.Rhs[i] == expr {
+					nodeExpr.Rhs[i] = cast(originalType, nodeExpr.Rhs[i])
+					return
+				}
+			}
+		case *ast.CallExpr:
+			for i := range nodeExpr.Args {
+				if nodeExpr.Args[i] == expr {
+					nodeExpr.Args[i] = cast(originalType, nodeExpr.Args[i])
+					return
+				}
+			}
+			parent := p.findParentNode(nodeExpr)
+			assign, isAssign := parent.(*ast.AssignStmt)
+			if nodeExpr.Fun == expr && isAssign {
+				for i := range assign.Rhs {
+					if assign.Rhs[i] == nodeExpr {
+						assign.Rhs[i] = cast(originalType, assign.Rhs[i])
+						return
+					}
+				}
+			}
+			call, isCall := parent.(*ast.CallExpr)
+			if isCall {
+				for i := range call.Args {
+					if call.Args[i] == nodeExpr {
+						call.Args[i] = cast(originalType, call.Args[i])
+						return
+					}
+				}
+			}
+			for i, v := range nodeExpr.Args {
+				if v == expr {
+					nodeExpr.Args[i] = cast(originalType, selectorExpr)
+					return
+				}
+			}
+		case *ast.BinaryExpr:
+			if nodeExpr.X == expr {
+				nodeExpr.X = cast(originalType, nodeExpr.X)
+			}
+			if nodeExpr.Y == expr {
+				nodeExpr.Y = cast(originalType, nodeExpr.Y)
+			}
+		}
+	}
+	patch(p.findParentNode(expr))
+}
+
+func isTypeValid(typeName string) bool {
+	return strings.Contains(typeName, ".") ||
+		strings.Contains(typeName, "/") ||
+		strings.Contains(typeName, "[]") ||
+		strings.Contains(typeName, "*")
+}