feat(injector): Add ArgumentThatImplements method for argument interface matching

This commit introduces the `ArgumentThatImplements` method to the function signature, allowing the retrieval of the first argument that implements a specified interface type. The implementation includes a new helper function, `findImplementingField`, which checks for both exact type matches and interface implementations.

Additionally, the `ResultThatImplements` method has been refactored to utilize the new helper for consistency. The changes enhance the injector's capability to handle argument interface checks, improving the overall functionality and maintainability of the code.

Configuration YAML files and test cases have been added to validate the new functionality across various scenarios, ensuring robust testing of the argument matching feature.

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

Author: kakkoyun

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

@@ -8,6 +8,7 @@ package code
 import (
 	"errors"
 	"fmt"
+	"go/types"
 
 	"github.com/DataDog/orchestrion/internal/injector/aspect/context"
 	"github.com/DataDog/orchestrion/internal/injector/typed"
@@ -28,6 +29,9 @@ type (
 		// ArgumentOfType returns the name of the first argument in this function that has the provided
 		// type, or an empty string if none is found.
 		ArgumentOfType(string) (string, error)
+		// ArgumentThatImplements returns the name of the first argument in this function that implements
+		// the provided interface type, or an empty string if none is found.
+		ArgumentThatImplements(string) (string, error)
 
 		// Result returns the name of the return value at the given index in this function's type,
 		// returning an error if the index is out of bounds.
@@ -111,6 +115,10 @@ func (noFunc) ArgumentOfType(string) (string, error) {
 	return "", errNoFunction
 }
 
+func (noFunc) ArgumentThatImplements(string) (string, error) {
+	return "", errNoFunction
+}
+
 func (noFunc) Result(int) (string, error) {
 	return "", errNoFunction
 }
@@ -144,6 +152,10 @@ func (s signature) ArgumentOfType(name string) (string, error) {
 	return fieldOfType(s.Params, name, "argument")
 }
 
+func (s signature) ArgumentThatImplements(interfaceName string) (string, error) {
+	return findImplementingField(s.context, s.Params, interfaceName, "argument")
+}
+
 func (s signature) Result(index int) (name string, err error) {
 	return fieldAt(s.Results, index, "result")
 }
@@ -152,39 +164,8 @@ func (s signature) ResultOfType(name string) (string, error) {
 	return fieldOfType(s.Results, name, "result")
 }
 
-func (s signature) ResultThatImplements(name string) (string, error) {
-	// Return blank if there are no results.
-	if s.Results == nil {
-		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 {
-		return "", fmt.Errorf("resolving interface type %q: %w", name, err)
-	}
-
-	// Check each result.
-	index := 0
-	for _, field := range s.Results.List {
-		if typed.ExprImplements(s.context, field.Type, iface) {
-			return fieldAt(s.Results, index, "result")
-		}
-
-		count := len(field.Names)
-		if count == 0 {
-			count = 1
-		}
-		index += count
-	}
-
-	// Not found.
-	return "", nil
+func (s signature) ResultThatImplements(interfaceName string) (string, error) {
+	return findImplementingField(s.context, s.Results, interfaceName, "result")
 }
 
 func (s signature) LastResultThatImplements(name string) (string, error) {
@@ -277,11 +258,19 @@ func fieldAt(fields *dst.FieldList, index int, use string) (string, error) {
 		}
 	}
 
-	if idx < index {
-		return "", fmt.Errorf("index out of bounds: %d (only %d items)", index, idx+1)
+	if idx <= index { // Use <= to catch index being exactly the number of items
+		return "", fmt.Errorf("index out of bounds: %d (only %d items)", index, idx)
 	}
 
-	return name, nil
+	// If anonymous, we should have assigned the synthetic name earlier.
+	// If named, it should have been returned immediately.
+	// If we reach here and it was anonymous, we return the generated name.
+	if anonymous {
+		return name, nil
+	}
+
+	// This path should ideally not be reached for named parameters if logic is correct.
+	return "", fmt.Errorf("fieldAt: failed to find field at index %d", index)
 }
 
 func fieldOfType(fields *dst.FieldList, typeName string, use string) (string, error) {
@@ -336,3 +325,42 @@ func (s signature) FinalResultImplements(interfaceName string) (bool, error) {
 	// Check if the last field type implements the interface.
 	return typed.ExprImplements(s.context, lastField.Type, iface), nil
 }
+
+// findImplementingField is a helper to find the first field in a list that matches
+// an interface, either by exact type name or by implementation.
+func findImplementingField(ctx context.AdviceContext, fields *dst.FieldList, interfaceName string, fieldKind string) (string, error) {
+	if fields == nil {
+		return "", nil // No fields, no match.
+	}
+
+	// 1. Check for exact type name match first.
+	if index, found := typed.FindMatchingTypeName(fields, interfaceName); found {
+		return fieldAt(fields, index, fieldKind)
+	}
+
+	// 2. If no exact name match, check for interface implementation.
+	iface, err := typed.ResolveInterfaceTypeByName(interfaceName)
+	if err != nil {
+		// Invalid interface name, cannot proceed with implementation check.
+		return "", fmt.Errorf("resolving interface type %q: %w", interfaceName, err)
+	}
+
+	// Iterate through fields to check for implementation.
+	currentIndex := 0
+	for _, field := range fields.List {
+		actualType := ctx.ResolveType(field.Type)
+		if actualType != nil && types.Implements(actualType, iface) {
+			return fieldAt(fields, currentIndex, fieldKind)
+		}
+
+		// Increment index based on field names.
+		count := len(field.Names)
+		if count == 0 {
+			count = 1
+		}
+		currentIndex += count
+	}
+
+	// 3. No match found.
+	return "", nil
+}

internal/injector/aspect/join/function.go

@@ -391,39 +391,39 @@ func (_ *resultImplements) fileMayMatch(_ *may.FileContext) may.MatchType {
 	return may.Unknown
 }
 
-func (fo *resultImplements) evaluate(info functionInformation) bool {
-	if info.Type.Results == nil || len(info.Type.Results.List) == 0 {
-		// No return values, no match.
+// evaluateFieldListImplements checks if any field in the list matches the interfaceName,
+// either by exact type name or by interface implementation.
+func evaluateFieldListImplements(fields *dst.FieldList, interfaceName string, info functionInformation) bool {
+	if fields == nil || len(fields.List) == 0 {
 		return false
 	}
 
 	// Optimization: First, check for an exact match using the helper.
-	if _, found := typed.FindMatchingTypeName(info.Type.Results, fo.InterfaceName); found {
+	if _, found := typed.FindMatchingTypeName(fields, interfaceName); found {
 		return true // Found direct match
-	} // If not found, fall through to type resolution.
+	}
 
-	// Ensure the type resolver is available.
+	// If no exact match, check implementation (requires type resolver).
 	if info.typeResolver == nil {
-		return false
+		return false // Cannot check implementation without resolver.
 	}
 
-	// Resolve the target interface name (e.g., "io.Reader", "error") to a types.Interface.
-	targetInterface, err := typed.ResolveInterfaceTypeByName(fo.InterfaceName)
+	targetInterface, err := typed.ResolveInterfaceTypeByName(interfaceName)
 	if err != nil {
-		// If the interface name is invalid or cannot be resolved, we cannot match.
-		return false
+		return false // Invalid interface name.
 	}
 
-	for _, field := range info.Type.Results.List {
-		// For each return type (dst.Expr), resolve it to types.Type using the provided resolver
-		// and check if it implements the target interface.
+	for _, field := range fields.List {
 		if typed.ExprImplements(info.typeResolver, field.Type, targetInterface) {
-			return true // Found at least one implementing type, match!
+			return true // Found an implementing type.
 		}
 	}
 
-	// No return type matched.
-	return false
+	return false // No match found.
+}
+
+func (fo *resultImplements) evaluate(info functionInformation) bool {
+	return evaluateFieldListImplements(info.Type.Results, fo.InterfaceName, info)
 }
 
 func (fo *resultImplements) Hash(h *fingerprint.Hasher) error {
@@ -496,6 +496,47 @@ func (fo *finalResultImplements) Hash(h *fingerprint.Hasher) error {
 	return h.Named("final-result-implements", fingerprint.String(fo.InterfaceName))
 }
 
+// argumentImplements matches functions where at least one argument's type
+// implements the specified interface.
+type argumentImplements struct {
+	InterfaceName string
+}
+
+// ArgumentImplements creates a FunctionOption that matches functions where at least one
+// argument implements the named interface.
+func ArgumentImplements(interfaceName string) FunctionOption {
+	return &argumentImplements{InterfaceName: interfaceName}
+}
+
+func (fo *argumentImplements) impliesImported() []string {
+	pkgPath, _ := typed.SplitPackageAndName(fo.InterfaceName)
+	if pkgPath != "" {
+		return []string{pkgPath}
+	}
+	return nil
+}
+
+func (_ *argumentImplements) packageMayMatch(_ *may.PackageContext) may.MatchType {
+	// Cannot reliably determine possibility of match based on package imports
+	// due to structural typing. A type can implement an interface without
+	// importing the interface's package.
+	return may.Unknown
+}
+
+func (_ *argumentImplements) fileMayMatch(_ *may.FileContext) may.MatchType {
+	// Cannot reliably determine possibility of match based on file contents
+	// due to structural typing and type aliases.
+	return may.Unknown
+}
+
+func (fo *argumentImplements) evaluate(info functionInformation) bool {
+	return evaluateFieldListImplements(info.Type.Params, fo.InterfaceName, info)
+}
+
+func (fo *argumentImplements) Hash(h *fingerprint.Hasher) error {
+	return h.Named("argument-implements", fingerprint.String(fo.InterfaceName))
+}
+
 func init() {
 	unmarshalers["function-body"] = func(ctx gocontext.Context, node ast.Node) (Point, error) {
 		up, err := FromYAML(ctx, node)
@@ -621,6 +662,15 @@ func (o *unmarshalFuncDeclOption) UnmarshalYAML(ctx gocontext.Context, node ast.
 		}
 		// NOTE: Validation happens later during type resolution.
 		o.FunctionOption = FinalResultImplements(ifaceName)
+	case "argument-implements":
+		var ifaceName string
+		if err := yaml.NodeToValueContext(ctx, mapping.Values[0].Value, &ifaceName); err != nil {
+			return err
+		}
+		if ifaceName == "" {
+			return fmt.Errorf("line %d: 'argument-implements' cannot be empty", node.GetToken().Position.Line)
+		}
+		o.FunctionOption = ArgumentImplements(ifaceName)
 	default:
 		return fmt.Errorf("unknown FuncDeclOption name: %q", key)
 	}