add cache and reverse pointer to Universe to Type

Author: ziyue-101

v2/parser/parse.go

@@ -26,6 +26,7 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 
 	"golang.org/x/tools/go/packages"
@@ -637,6 +638,13 @@ func (p *Parser) convertSignature(u types.Universe, t *gotypes.Signature) *types
 
 // walkType adds the type, and any necessary child types.
 func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type) *types.Type {
+	var out *types.Type
+	defer func() {
+		out.Universe = &u
+		out.PointerTypeCache = &types.PointerTypeCache{
+			Mu: sync.Mutex{},
+		}
+	}()
 	// Most of the cases are underlying types of the named type.
 	name := goNameToName(in.String())
 	if useName != nil {
@@ -645,13 +653,13 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 
 	// Handle alias types conditionally on go1.22+.
 	// Inline this once the minimum supported version is go1.22
-	if out := p.walkAliasType(u, in); out != nil {
+	if out = p.walkAliasType(u, in); out != nil {
 		return out
 	}
 
 	switch t := in.(type) {
 	case *gotypes.Struct:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -670,7 +678,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		}
 		return out
 	case *gotypes.Map:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -680,7 +688,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		out.Key = p.walkType(u, nil, t.Key())
 		return out
 	case *gotypes.Pointer:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -689,7 +697,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		out.Elem = p.walkType(u, nil, t.Elem())
 		return out
 	case *gotypes.Slice:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -698,7 +706,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		out.Elem = p.walkType(u, nil, t.Elem())
 		return out
 	case *gotypes.Array:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -708,7 +716,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		out.Len = in.(*gotypes.Array).Len()
 		return out
 	case *gotypes.Chan:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -719,7 +727,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		// cannot be properly written.
 		return out
 	case *gotypes.Basic:
-		out := u.Type(types.Name{
+		out = u.Type(types.Name{
 			Package: "", // This is a magic package name in the Universe.
 			Name:    t.Name(),
 		})
@@ -730,7 +738,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		out.Kind = types.Unsupported
 		return out
 	case *gotypes.Signature:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -739,7 +747,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		out.Signature = p.convertSignature(u, t)
 		return out
 	case *gotypes.Interface:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out
@@ -758,7 +766,6 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 		}
 		return out
 	case *gotypes.Named:
-		var out *types.Type
 		switch t.Underlying().(type) {
 		case *gotypes.Named, *gotypes.Basic, *gotypes.Map, *gotypes.Slice:
 			name := goNameToName(t.String())
@@ -785,7 +792,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 				name.Name = fmt.Sprintf("%s[%s]", strings.SplitN(name.Name, "[", 2)[0], strings.Join(tpNames, ","))
 			}
 
-			if out := u.Type(name); out.Kind != types.Unknown {
+			if out = u.Type(name); out.Kind != types.Unknown {
 				out.GoType = in
 				return out // short circuit if we've already made this.
 			}
@@ -797,7 +804,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 			// "feature" for users. This flattens those types
 			// together.
 			name := goNameToName(t.String())
-			if out := u.Type(name); out.Kind != types.Unknown {
+			if out = u.Type(name); out.Kind != types.Unknown {
 				return out // short circuit if we've already made this.
 			}
 			out = p.walkType(u, &name, t.Underlying())
@@ -827,7 +834,7 @@ func (p *Parser) walkType(u types.Universe, useName *types.Name, in gotypes.Type
 			Kind: types.TypeParam,
 		}
 	default:
-		out := u.Type(name)
+		out = u.Type(name)
 		out.GoType = in
 		if out.Kind != types.Unknown {
 			return out

v2/parser/parse_test.go

@@ -548,6 +548,15 @@ func TestAddOnePkgToUniverse(t *testing.T) {
 		if len(ud.Types) == 0 {
 			t.Errorf("expected non-zero types in direct package")
 		} else {
+			// verify that Universe and PointerTypeCache has been populated
+			for _, tp := range ud.Types {
+				if tp.Universe == nil {
+					t.Errorf("expected no nil Universe in type %v", tp)
+				}
+				if tp.PointerTypeCache == nil {
+					t.Errorf("expected no nil PointerTypeCache in type %v", tp)
+				}
+			}
 			type testcase struct {
 				kind       types.Kind
 				elem       string // just the type name

v2/types/types.go

@@ -19,6 +19,7 @@ package types
 import (
 	gotypes "go/types"
 	"strings"
+	"sync"
 )
 
 // Ref makes a reference to the given type. It can only be used for e.g.
@@ -289,6 +290,13 @@ func (u Universe) Package(packagePath string) *Package {
 	return p
 }
 
+type PointerTypeCache struct {
+	// mutex to protect concurrent update to the cache
+	Mu sync.Mutex
+	// A pointer Type for the Type
+	PointerType *Type
+}
+
 // Type represents a subset of possible go types.
 type Type struct {
 	// There are two general categories of types, those explicitly named
@@ -364,6 +372,12 @@ type Type struct {
 
 	// The underlying Go type.
 	GoType gotypes.Type
+
+	// Revevrse reference to Universe
+	Universe *Universe
+
+	// A pointer type cache
+	PointerTypeCache *PointerTypeCache
 }
 
 // String returns the name of the type.
@@ -556,13 +570,41 @@ var (
 )
 
 func PointerTo(t *Type) *Type {
-	return &Type{
+	if t.Universe != nil {
+		for _, pkg := range *t.Universe {
+			for _, tp := range pkg.Types {
+				existingTypeName := tp.Name.String()
+				if strings.HasPrefix(existingTypeName, "*") {
+					if existingTypeName[1:] == t.Name.String() {
+						return tp
+					}
+				}
+			}
+		}
+	}
+	cache := t.PointerTypeCache
+	if cache != nil && cache.PointerType != nil {
+		return cache.PointerType
+	}
+	out := &Type{
 		Name: Name{
 			Name: "*" + t.Name.String(),
 		},
 		Kind: Pointer,
 		Elem: t,
 	}
+	// PointerTypeCache should not be empty at this point. This is a sanity processing.
+	if cache == nil {
+		cache = &PointerTypeCache{
+			Mu: sync.Mutex{},
+		}
+	}
+	if cache.PointerType == nil {
+		cache.Mu.Lock()
+		cache.PointerType = out
+		cache.Mu.Unlock()
+	}
+	return out
 }
 
 func IsInteger(t *Type) bool {

v2/types/types_test.go

@@ -17,6 +17,8 @@ limitations under the License.
 package types
 
 import (
+	"reflect"
+	"sync"
 	"testing"
 )
 
@@ -37,6 +39,86 @@ func TestGetBuiltin(t *testing.T) {
 	}
 }
 
+func TestPointerTo(t *testing.T) {
+	u := Universe{
+		"pkgname": &Package{
+			Types: map[string]*Type{},
+		},
+	}
+	pkgInUniverse := u["pkgname"]
+	type1 := &Type{Name: Name{Name: "structname"}, Kind: Struct}
+	type2 := &Type{Name: Name{Name: "secondstructname"}, Kind: Struct}
+	type3 := &Type{Name: Name{Name: "thirdstructname"}, Kind: Struct}
+
+	setupUniverseAndPointerTypeCache := func(tp *Type, u *Universe, pkg *Package, cached bool, existPointerTypeInUniverse bool) {
+		tp.Universe = u
+		tp.PointerTypeCache = &PointerTypeCache{
+			Mu: sync.Mutex{},
+		}
+		if cached {
+			tp.PointerTypeCache.PointerType = &Type{
+				Name: Name{
+					Name: "*" + tp.Name.String(),
+				},
+				Kind: Pointer,
+				Elem: tp,
+			}
+		}
+		if existPointerTypeInUniverse {
+			pkg.Types["*"+tp.Name.String()] = &Type{
+				Name: Name{
+					Name: "*" + tp.Name.String(),
+				},
+				Kind: Pointer,
+				Elem: tp,
+			}
+		}
+	}
+
+	setupUniverseAndPointerTypeCache(type1, &u, pkgInUniverse, false, true)
+	setupUniverseAndPointerTypeCache(type2, &u, pkgInUniverse, true, false)
+	setupUniverseAndPointerTypeCache(type3, &u, pkgInUniverse, false, false)
+
+	testCases := []struct {
+		name           string
+		typ            *Type
+		expected       *Type
+		expectCreation bool
+	}{
+		{
+			name:           "universe has the pointer type",
+			typ:            type1,
+			expected:       pkgInUniverse.Types["*"+type1.Name.String()],
+			expectCreation: false,
+		},
+		{
+			name:           "cache has the pointer type",
+			typ:            type2,
+			expected:       type2.PointerTypeCache.PointerType,
+			expectCreation: false,
+		},
+		{
+			name: "cache doesn't have the pointer type",
+			typ:  type3,
+			expected: &Type{
+				Name: Name{Name: "*thirdstructname"},
+				Kind: Pointer,
+				Elem: type3,
+			},
+			expectCreation: true,
+		},
+	}
+	for _, tc := range testCases {
+		tp := PointerTo(tc.typ)
+		if tc.expectCreation && !reflect.DeepEqual(tp, tc.expected) {
+			t.Errorf("PointerTo failed, expected %v, got : %v", tc.expected, tp)
+		}
+		if !tc.expectCreation && tp != tc.expected {
+			t.Errorf("PointerTo should return an existing pointer type, expected %v, got : %v", tc.expected, tp)
+		}
+	}
+}
+
 func TestGetMarker(t *testing.T) {
 	u := Universe{}
 	n := Name{Package: "path/to/package", Name: "Foo"}