refactor(injector): Refactor TypeName handling and migrate to typed package

This commit refactors the handling of TypeName across the injector package, transitioning from the join package to a new typed package. The changes include:

- Replacing instances of join.TypeName with typed.TypeName in various files, ensuring consistent type handling.
- Updating related functions and methods to accommodate the new TypeName structure, including adjustments to import path retrieval and pointer handling.
- Adding comprehensive tests for the new TypeName implementation to validate its functionality and error handling.

This refactor enhances code clarity and maintainability by centralizing type name logic within the typed package.

Signed-off-by: Kemal Akkoyun <kemal.akkoyun@datadoghq.com>

Author: kakkoyun

internal/injector/aspect/advice/call.go

@@ -14,20 +14,20 @@ import (
 	"github.com/DataDog/orchestrion/internal/fingerprint"
 	"github.com/DataDog/orchestrion/internal/injector/aspect/advice/code"
 	"github.com/DataDog/orchestrion/internal/injector/aspect/context"
-	"github.com/DataDog/orchestrion/internal/injector/aspect/join"
+	"github.com/DataDog/orchestrion/internal/injector/typed"
 	"github.com/DataDog/orchestrion/internal/yaml"
 	"github.com/dave/dst"
 	"github.com/goccy/go-yaml/ast"
 )
 
 type appendArgs struct {
-	TypeName  join.TypeName
+	TypeName  typed.TypeName
 	Templates []*code.Template
 }
 
 // AppendArgs appends arguments of a given type to the end of a function call. All arguments must be
 // of the same type, as they may be appended at the tail end of a variadic call.
-func AppendArgs(typeName join.TypeName, templates ...*code.Template) *appendArgs {
+func AppendArgs(typeName typed.TypeName, templates ...*code.Template) *appendArgs {
 	return &appendArgs{typeName, templates}
 }
 
@@ -92,7 +92,7 @@ func (a *appendArgs) Apply(ctx context.AdviceContext) (bool, error) {
 		Ellipsis: true,
 	}
 
-	if importPath := a.TypeName.ImportPath(); importPath != "" {
+	if importPath := a.TypeName.ImportPath; importPath != "" {
 		ctx.AddImport(importPath, inferPkgName(importPath))
 	}
 
@@ -101,7 +101,7 @@ func (a *appendArgs) Apply(ctx context.AdviceContext) (bool, error) {
 
 func (a *appendArgs) AddedImports() []string {
 	imports := make([]string, 0, len(a.Templates)+1)
-	if argTypeImportPath := a.TypeName.ImportPath(); argTypeImportPath != "" {
+	if argTypeImportPath := a.TypeName.ImportPath; argTypeImportPath != "" {
 		imports = append(imports, argTypeImportPath)
 	}
 	for _, t := range a.Templates {
@@ -168,7 +168,7 @@ func init() {
 			return nil, err
 		}
 
-		tn, err := join.NewTypeName(args.TypeName)
+		tn, err := typed.NewTypeName(args.TypeName)
 		if err != nil {
 			return nil, err
 		}

internal/injector/aspect/advice/call_test.go

@@ -11,31 +11,31 @@ import (
 	"github.com/DataDog/orchestrion/internal/injector/aspect/advice"
 	"github.com/DataDog/orchestrion/internal/injector/aspect/advice/code"
 	"github.com/DataDog/orchestrion/internal/injector/aspect/context"
-	"github.com/DataDog/orchestrion/internal/injector/aspect/join"
+	"github.com/DataDog/orchestrion/internal/injector/typed"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestAppendArgs(t *testing.T) {
 	t.Run("AddedImports", func(t *testing.T) {
 		type testCase struct {
-			argType         join.TypeName
+			argType         typed.TypeName
 			args            []*code.Template
 			expectedImports []string
 		}
 
 		testCases := map[string]testCase{
 			"imports-none": {
-				argType: join.MustTypeName("any"),
+				argType: typed.MustTypeName("any"),
 				args:    []*code.Template{code.MustTemplate("true", nil, context.GoLangVersion{})},
 			},
 			"imports-from-arg-type": {
-				argType:         join.MustTypeName("*net/http.Request"),
+				argType:         typed.MustTypeName("*net/http.Request"),
 				args:            []*code.Template{code.MustTemplate("true", nil, context.GoLangVersion{})},
 				expectedImports: []string{"net/http"},
 			},
 			"imports-from-templates": {
-				argType: join.MustTypeName("any"),
+				argType: typed.MustTypeName("any"),
 				args: []*code.Template{
 					code.MustTemplate("imp.Value", map[string]string{"imp": "github.com/namespace/foo"}, context.GoLangVersion{}),
 					code.MustTemplate("imp.Value", map[string]string{"imp": "github.com/namespace/bar"}, context.GoLangVersion{}),

internal/injector/aspect/advice/code/dot_function.go

@@ -8,14 +8,10 @@ package code
 import (
 	"errors"
 	"fmt"
-	"go/importer"
-	"go/types"
-	"strings"
 
 	"github.com/dave/dst"
 
 	"github.com/DataDog/orchestrion/internal/injector/aspect/context"
-	"github.com/DataDog/orchestrion/internal/injector/aspect/join"
 	"github.com/DataDog/orchestrion/internal/injector/typed"
 )
 
@@ -162,6 +158,11 @@ func (s signature) ResultThatImplements(name string) (string, error) {
 		return "", nil
 	}
 
+	// Optimization: First, check for an exact match using the helper.
+	if index, found := typed.FindMatchingTypeName(s.Results, name); found {
+		return fieldAt(s.Results, index, "result")
+	} // If not found, fall through to type resolution.
+
 	// Resolve the interface type.
 	iface, err := typed.ResolveInterfaceTypeByName(name)
 	if err != nil {
@@ -192,18 +193,38 @@ func (s signature) LastResultThatImplements(name string) (string, error) {
 		return "", nil
 	}
 
+	// Optimization: First, check for an exact match using TypeName parsing, finding the last one.
+	lastMatchIndex := -1
+	if tn, err := typed.NewTypeName(name); err == nil {
+		currentIndex := 0
+		for _, field := range s.Results.List {
+			if tn.Matches(field.Type) {
+				lastMatchIndex = currentIndex // Update last found index
+			}
+			// Increment index by the number of names in the field (or 1 if unnamed).
+			count := len(field.Names)
+			if count == 0 {
+				count = 1
+			}
+			currentIndex += count
+		}
+	}
+	// If we found a match via TypeName, return it.
+	if lastMatchIndex != -1 {
+		return fieldAt(s.Results, lastMatchIndex, "result")
+	} // If parsing failed or no match, fall through to type resolution.
+
 	// Resolve the interface type.
 	iface, err := typed.ResolveInterfaceTypeByName(name)
 	if err != nil {
+		// Propagate error if interface resolution fails
 		return "", fmt.Errorf("resolving interface type %q: %w", name, err)
 	}
 
-	// First, we need to build a map of result fields to their indices
-	// that takes into account named and unnamed parameters.
-	var (
-		fieldIndices = make(map[*dst.Field]int)
-		index        = 0
-	)
+	// Fallback: Check using ExprImplements, iterating backward.
+	// Need field indices map again for this path.
+	fieldIndices := make(map[*dst.Field]int)
+	index := 0
 	for _, field := range s.Results.List {
 		fieldIndices[field] = index
 		count := len(field.Names)
@@ -213,7 +234,6 @@ func (s signature) LastResultThatImplements(name string) (string, error) {
 		index += count
 	}
 
-	// Loop backward through the results list.
 	for i := len(s.Results.List) - 1; i >= 0; i-- {
 		field := s.Results.List[i]
 		if typed.ExprImplements(s.context, field.Type, iface) {
@@ -265,7 +285,7 @@ func fieldAt(fields *dst.FieldList, index int, use string) (string, error) {
 }
 
 func fieldOfType(fields *dst.FieldList, typeName string, use string) (string, error) {
-	tn, err := join.NewTypeName(typeName)
+	tn, err := typed.NewTypeName(typeName)
 	if err != nil {
 		return "", err
 	}
@@ -292,118 +312,27 @@ func fieldOfType(fields *dst.FieldList, typeName string, use string) (string, er
 	return "", nil
 }
 
-// exprImplements checks if an expression's type implements an interface.
-func exprImplements(ctx context.AdviceContext, expr dst.Expr, iface *types.Interface) bool {
-	actualType := ctx.ResolveType(expr)
-	if actualType == nil {
-		return false
-	}
-
-	return typeImplements(actualType, iface)
-}
-
-// typeImplements checks if a type implements an interface (including pointer receivers).
-func typeImplements(t types.Type, iface *types.Interface) bool {
-	if t == nil || iface == nil {
-		return false
-	}
-
-	// Direct implementation check.
-	if types.Implements(t, iface) {
-		return true
-	}
-
-	return false
-}
-
-// resolveInterfaceTypeByName takes an interface name as a string and resolves it to an interface type.
-// It supports built-in interfaces (e.g. "error"), package qualified interfaces (e.g. "io.Reader"),
-// and third-party package interfaces (e.g. "example.com/pkg.Interface").
-func resolveInterfaceTypeByName(name string) (*types.Interface, error) {
-	// Handle built-in types.
-	if obj := types.Universe.Lookup(name); obj != nil {
-		typeObj, ok := obj.(*types.TypeName)
-		if !ok {
-			return nil, fmt.Errorf("object %s is not a type name but a %T", name, obj)
-		}
-
-		typ := typeObj.Type()
-		if !types.IsInterface(typ) {
-			return nil, fmt.Errorf("type %s is not an interface", name)
-		}
-
-		t, ok := typ.Underlying().(*types.Interface)
-		if !ok {
-			return nil, fmt.Errorf("type %s is not an interface", name)
-		}
-
-		return t, nil
-	}
-
-	// Handle package-qualified types (e.g., "io.Writer").
-	pkgName, typeName := splitPackageAndName(name)
-	if pkgName == "" {
-		return nil, fmt.Errorf("invalid type name: %s", name)
-	}
-
-	// Import the package
-	imp := importer.Default()
-	pkg, err := imp.Import(pkgName)
-	if err != nil {
-		return nil, fmt.Errorf("failed to import package %q: %w", pkgName, err)
-	}
-
-	// Look up the type in the package's scope
-	obj := pkg.Scope().Lookup(typeName)
-	if obj == nil {
-		return nil, fmt.Errorf("type %q not found in package %q", typeName, pkgName)
-	}
-
-	typeObj, ok := obj.(*types.TypeName)
-	if !ok {
-		return nil, fmt.Errorf("object %s is not a type name but a %T", name, obj)
-	}
-
-	typ := typeObj.Type()
-	if !types.IsInterface(typ) {
-		return nil, fmt.Errorf("type %s is not an interface", name)
-	}
-
-	t, ok := typ.Underlying().(*types.Interface)
-	if !ok {
-		return nil, fmt.Errorf("type %s is not an interface", name)
-	}
-
-	return t, nil
-}
-
-// splitPackageAndName splits a fully qualified type name like "io.Reader" or "example.com/pkg.Type"
-// into its package path and local name.
-// Returns ("", "error") for built-in "error".
-// Returns ("", "MyType") for unqualified "MyType".
-func splitPackageAndName(fullName string) (pkgPath string, localName string) {
-	if !strings.Contains(fullName, ".") {
-		// Assume built-in type (like "error") or unqualified local type.
-		return "", fullName
-	}
-	lastDot := strings.LastIndex(fullName, ".")
-	pkgPath = fullName[:lastDot]
-	localName = fullName[lastDot+1:]
-	return pkgPath, localName
-}
-
 // FinalResultImplements returns whether the final result implements the provided interface type.
 func (s signature) FinalResultImplements(interfaceName string) (bool, error) {
 	if s.Results == nil || len(s.Results.List) == 0 {
 		return false, nil
 	}
 
-	iface, err := resolveInterfaceTypeByName(interfaceName)
+	lastField := s.Results.List[len(s.Results.List)-1]
+
+	// Optimization: First, check for an exact match using TypeName parsing.
+	// Note: Not using FindMatchingTypeName as we only need to check the last field.
+	if tn, err := typed.NewTypeName(interfaceName); err == nil {
+		if tn.Matches(lastField.Type) {
+			return true, nil
+		}
+	} // If parsing failed or no match, fall through to type resolution.
+
+	iface, err := typed.ResolveInterfaceTypeByName(interfaceName)
 	if err != nil {
 		return false, fmt.Errorf("resolving interface type %q: %w", interfaceName, err)
 	}
 
 	// Check if the last field type implements the interface.
-	lastField := s.Results.List[len(s.Results.List)-1]
-	return exprImplements(s.context, lastField.Type, iface), nil
+	return typed.ExprImplements(s.context, lastField.Type, iface), nil
 }

internal/injector/aspect/advice/struct.go

@@ -11,19 +11,19 @@ import (
 
 	"github.com/DataDog/orchestrion/internal/fingerprint"
 	"github.com/DataDog/orchestrion/internal/injector/aspect/context"
-	"github.com/DataDog/orchestrion/internal/injector/aspect/join"
+	"github.com/DataDog/orchestrion/internal/injector/typed"
 	"github.com/DataDog/orchestrion/internal/yaml"
 	"github.com/dave/dst"
 	"github.com/goccy/go-yaml/ast"
 )
 
 type addStructField struct {
 	Name     string
-	TypeName join.TypeName
+	TypeName typed.TypeName
 }
 
 // AddStructField adds a new synthetic field at the tail end of a struct declaration.
-func AddStructField(fieldName string, fieldType join.TypeName) *addStructField {
+func AddStructField(fieldName string, fieldType typed.TypeName) *addStructField {
 	return &addStructField{fieldName, fieldType}
 }
 
@@ -47,7 +47,7 @@ func (a *addStructField) Apply(ctx context.AdviceContext) (bool, error) {
 		Type:  a.TypeName.AsNode(),
 	})
 
-	if importPath := a.TypeName.ImportPath(); importPath != "" {
+	if importPath := a.TypeName.ImportPath; importPath != "" {
 		// If the type name is qualified, we may need to import the package, too.
 		_ = ctx.AddImport(importPath, inferPkgName(importPath))
 	}
@@ -60,7 +60,7 @@ func (a *addStructField) Hash(h *fingerprint.Hasher) error {
 }
 
 func (a *addStructField) AddedImports() []string {
-	if path := a.TypeName.ImportPath(); path != "" {
+	if path := a.TypeName.ImportPath; path != "" {
 		return []string{path}
 	}
 	return nil
@@ -76,7 +76,7 @@ func init() {
 		if err := yaml.NodeToValueContext(ctx, node, &spec); err != nil {
 			return nil, err
 		}
-		tn, err := join.NewTypeName(spec.Type)
+		tn, err := typed.NewTypeName(spec.Type)
 		if err != nil {
 			return nil, err
 		}