Support recv generic funcs

Author: txabman42

demo/basic/main.go

@@ -31,6 +31,15 @@ type MyStruct struct{}
 
 func (m *MyStruct) Example() { println("MyStruct.Example") }
 
+type GenStruct[T any] struct {
+	Value T
+}
+
+func (m *GenStruct[T]) GenericRecvExample(t T) T {
+	fmt.Printf("%s%s\n", m.Value, t)
+	return t
+}
+
 func GenericExample[K comparable, V any](key K, value V) V {
 	println("Hello, Generic World!", key, value)
 	return value
@@ -59,11 +68,14 @@ func main() {
 	// Call the Example function to trigger the instrumentation
 	Example()
 	m := &MyStruct{}
-	GenericExample(1, 2)
 	// Add a new field to the struct
 	m.NewField = "abc"
 	m.Example()
 
+	GenericExample(1, 2)
+	g := &GenStruct[string]{Value: "Hello"}
+	_ = g.GenericRecvExample(", Generic Recv World!")
+
 	// Call real module function
 	println(rate.Every(time.Duration(1)))
 }

pkg/instrumentation/helloworld/helloworld_hook.go

@@ -71,6 +71,14 @@ func MyHook1After(ictx inst.HookContext) {
 	println("After MyStruct.Example()")
 }
 
+func MyHookRecvBefore(ictx inst.HookContext, recv, _ interface{}) {
+	println("GenericRecvExample before hook")
+}
+
+func MyHookRecvAfter(ictx inst.HookContext, _ interface{}) {
+	println("GenericRecvExample after hook")
+}
+
 func MyHookGenericBefore(ictx inst.HookContext, _, _ interface{}) {
 	println("GenericExample before hook")
 }

test/integration/basic_test.go

@@ -27,6 +27,10 @@ func TestBasic(t *testing.T) {
 		"Hello, Generic World! 1 2",
 		"GenericExample after hook",
 		"traceID: 123, spanID: 456",
+		"GenericRecvExample before hook",
+		"Hello, Generic Recv World!",
+		"GenericRecvExample after hook",
+		"traceID: 123, spanID: 456",
 		"[MyHook]",
 		"=setupOpenTelemetry=",
 		"RawCode",

tool/data/helloworld.yaml

@@ -64,3 +64,11 @@ hook_generic:
   before: MyHookGenericBefore
   after: MyHookGenericAfter
   path: "github.com/open-telemetry/opentelemetry-go-compile-instrumentation/pkg/instrumentation/helloworld"
+
+hook_generic_recv:
+  target: main
+  func: GenericRecvExample
+  recv: "*GenStruct"
+  before: MyHookRecvBefore
+  after: MyHookRecvAfter
+  path: "github.com/open-telemetry/opentelemetry-go-compile-instrumentation/pkg/instrumentation/helloworld"

tool/internal/ast/shared.go

@@ -44,6 +44,48 @@ func FindFuncDeclWithoutRecv(root *dst.File, funcName string) *dst.FuncDecl {
 	return decls[0]
 }
 
+// extractBaseReceiverType extracts the base type name from a receiver expression,
+// handling both generic and non-generic types.
+// For example:
+// - *MyStruct -> *MyStruct
+// - MyStruct -> MyStruct
+// - *GenStruct[T] -> *GenStruct
+// - GenStruct[T] -> GenStruct
+func extractBaseReceiverType(recvTypeExpr dst.Expr) string {
+	switch expr := recvTypeExpr.(type) {
+	case *dst.StarExpr: // func (*Recv)T or func (*Recv[T])T
+		// Check if X is an Ident (non-generic) or IndexExpr/IndexListExpr (generic)
+		switch x := expr.X.(type) {
+		case *dst.Ident:
+			// Non-generic pointer receiver: *MyStruct
+			return "*" + x.Name
+		case *dst.IndexExpr:
+			// Generic pointer receiver with single type param: *GenStruct[T]
+			if baseIdent, ok := x.X.(*dst.Ident); ok {
+				return "*" + baseIdent.Name
+			}
+		case *dst.IndexListExpr:
+			// Generic pointer receiver with multiple type params: *GenStruct[T, U]
+			if baseIdent, ok := x.X.(*dst.Ident); ok {
+				return "*" + baseIdent.Name
+			}
+		}
+	case *dst.Ident: // func (Recv)T
+		return expr.Name
+	case *dst.IndexExpr:
+		// Generic value receiver with single type param: GenStruct[T]
+		if baseIdent, ok := expr.X.(*dst.Ident); ok {
+			return baseIdent.Name
+		}
+	case *dst.IndexListExpr:
+		// Generic value receiver with multiple type params: GenStruct[T, U]
+		if baseIdent, ok := expr.X.(*dst.Ident); ok {
+			return baseIdent.Name
+		}
+	}
+	return ""
+}
+
 func FindFuncDecl(root *dst.File, funcName, recv string) *dst.FuncDecl {
 	decls := findFuncDecls(root, func(funcDecl *dst.FuncDecl) bool {
 		// Receiver type is ignored, match func name only
@@ -59,26 +101,15 @@ func FindFuncDecl(root *dst.File, funcName, recv string) *dst.FuncDecl {
 
 		// Receiver type is specified, and target function has receiver
 		// Match both func name and receiver type
-		switch recvTypeExpr := funcDecl.Recv.List[0].Type.(type) {
-		case *dst.StarExpr: // func (*Recv)T
-			tn, ok := recvTypeExpr.X.(*dst.Ident)
-			if !ok {
-				// This is a generic type, we don't support it yet
-				return false
-			}
-			t := "*" + tn.Name
-			return t == recv && name == funcName
-		case *dst.Ident: // func (Recv)T
-			t := recvTypeExpr.Name
-			return t == recv && name == funcName
-		case *dst.IndexExpr:
-			// This is a generic type, we don't support it yet
-			return false
-		default:
+		recvTypeExpr := funcDecl.Recv.List[0].Type
+		baseType := extractBaseReceiverType(recvTypeExpr)
+
+		if baseType == "" {
 			msg := fmt.Sprintf("unexpected receiver type: %T", recvTypeExpr)
 			util.Unimplemented(msg)
 		}
-		return false
+
+		return baseType == recv && name == funcName
 	})
 
 	if len(decls) == 0 {

tool/internal/instrument/trampoline.go

@@ -473,8 +473,9 @@ func (ip *InstrumentPhase) buildTrampolineTypes() {
 		}
 	}
 	addHookContext(afterHookFunc.Type.Params)
-	beforeHookFunc.Type.TypeParams = ast.CloneTypeParams(ip.targetFunc.Type.TypeParams)
-	afterHookFunc.Type.TypeParams = ast.CloneTypeParams(ip.targetFunc.Type.TypeParams)
+	trampolineTypeParams := combineTypeParams(ip.targetFunc)
+	beforeHookFunc.Type.TypeParams = ast.CloneTypeParams(trampolineTypeParams)
+	afterHookFunc.Type.TypeParams = ast.CloneTypeParams(trampolineTypeParams)
 }
 
 func assignString(assignStmt *dst.AssignStmt, val string) bool {
@@ -640,6 +641,108 @@ func setReturnValClause(idx int, t dst.Expr) *dst.CaseClause {
 	return setValue(trampolineReturnValsIdentifier, idx, t)
 }
 
+// extractReceiverTypeParams extracts type parameters from a receiver type expression
+// For example: *GenStruct[T] or GenStruct[T, U] -> FieldList with T and U as type parameters
+func extractReceiverTypeParams(recvType dst.Expr) *dst.FieldList {
+	switch t := recvType.(type) {
+	case *dst.StarExpr:
+		// *GenStruct[T] - recurse into X
+		return extractReceiverTypeParams(t.X)
+	case *dst.IndexExpr:
+		// GenStruct[T] - single type parameter
+		if ident, ok := t.Index.(*dst.Ident); ok {
+			return &dst.FieldList{
+				List: []*dst.Field{{
+					Names: []*dst.Ident{ident},
+					Type:  ast.Ident("any"), // Type constraint for the parameter
+				}},
+			}
+		}
+	case *dst.IndexListExpr:
+		// GenStruct[T, U, ...] - multiple type parameters
+		fields := make([]*dst.Field, 0, len(t.Indices))
+		for _, idx := range t.Indices {
+			if ident, ok := idx.(*dst.Ident); ok {
+				fields = append(fields, &dst.Field{
+					Names: []*dst.Ident{ident},
+					Type:  ast.Ident("any"), // Type constraint for the parameter
+				})
+			}
+		}
+		if len(fields) > 0 {
+			return &dst.FieldList{List: fields}
+		}
+	}
+	return nil
+}
+
+// combineTypeParams combines type parameters from the receiver and function type parameters.
+// For methods on generic types, it extracts type parameters from the receiver and merges
+// them with the function's type parameters.
+// Receiver type parameters come first, followed by function type parameters.
+//
+// Example:
+//
+//	Original: func (c *Container[K]) Transform[V any]() V
+//	Result: [K, V]
+//
+//	Generated trampolines:
+//	  func OtelBeforeTrampoline_Container_Transform[K comparable, V any](
+//	      hookContext *HookContext,
+//	      recv0 *Container[K],  // ← Uses K
+//	  ) { ... }
+//
+//	  func OtelAfterTrampoline_Container_Transform[K comparable, V any](
+//	      hookContext *HookContext,
+//	      arg0 *V,  // ← Uses V (return type)
+//	  ) { ... }
+func combineTypeParams(targetFunc *dst.FuncDecl) *dst.FieldList {
+	var trampolineTypeParams *dst.FieldList
+	if ast.HasReceiver(targetFunc) {
+		receiverTypeParams := extractReceiverTypeParams(targetFunc.Recv.List[0].Type)
+		if receiverTypeParams != nil {
+			trampolineTypeParams = receiverTypeParams
+		}
+	}
+	if targetFunc.Type.TypeParams != nil {
+		if trampolineTypeParams == nil {
+			trampolineTypeParams = targetFunc.Type.TypeParams
+		} else {
+			combined := &dst.FieldList{List: make([]*dst.Field, 0)}
+			combined.List = append(combined.List, trampolineTypeParams.List...)
+			combined.List = append(combined.List, targetFunc.Type.TypeParams.List...)
+			trampolineTypeParams = combined
+		}
+	}
+	return trampolineTypeParams
+}
+
+// replaceGenericInstantiations replaces generic type instantiations with interface{}
+// For example: *GenStruct[T] -> *interface{}, []GenStruct[T, U] -> []interface{}
+func replaceGenericInstantiations(t dst.Expr) dst.Expr {
+	switch tType := t.(type) {
+	case *dst.StarExpr:
+		// *GenStruct[T] -> *interface{}
+		return ast.DereferenceOf(replaceGenericInstantiations(tType.X))
+	case *dst.ArrayType:
+		// []GenStruct[T] -> []interface{}
+		return ast.ArrayType(replaceGenericInstantiations(tType.Elt))
+	case *dst.MapType:
+		// map[K]GenStruct[T] -> map[interface{}]interface{}
+		return &dst.MapType{
+			Key:   replaceGenericInstantiations(tType.Key),
+			Value: replaceGenericInstantiations(tType.Value),
+		}
+	case *dst.IndexExpr:
+		// GenStruct[T] -> interface{}
+		return ast.InterfaceType()
+	case *dst.IndexListExpr:
+		// GenStruct[T, U] -> interface{}
+		return ast.InterfaceType()
+	}
+	return t
+}
+
 // desugarType desugars parameter type to its original type, if parameter
 // is type of ...T, it will be converted to []T
 func desugarType(param *dst.Field) dst.Expr {
@@ -677,10 +780,22 @@ func (ip *InstrumentPhase) rewriteHookContext() {
 	methodGetParamBody := findSwitchBlock(methodGetParam, 0)
 	methodSetRetValBody := findSwitchBlock(methodSetRetVal, 1)
 	methodGetRetValBody := findSwitchBlock(methodGetRetVal, 0)
+
+	combinedTypeParams := combineTypeParams(ip.targetFunc)
+
+	ip.rewriteHookContextParams(methodSetParamBody, methodGetParamBody, combinedTypeParams)
+	ip.rewriteHookContextResults(methodSetRetValBody, methodGetRetValBody, combinedTypeParams)
+}
+
+func (ip *InstrumentPhase) rewriteHookContextParams(
+	methodSetParamBody, methodGetParamBody *dst.BlockStmt,
+	combinedTypeParams *dst.FieldList,
+) {
 	idx := 0
 	if ast.HasReceiver(ip.targetFunc) {
 		splitRecv := ast.SplitMultiNameFields(ip.targetFunc.Recv)
-		recvType := replaceTypeParamsWithAny(splitRecv.List[0].Type, ip.targetFunc.Type.TypeParams)
+		recvType := replaceGenericInstantiations(splitRecv.List[0].Type)
+		recvType = replaceTypeParamsWithAny(recvType, combinedTypeParams)
 		clause := setParamClause(idx, recvType)
 		methodSetParamBody.List = append(methodSetParamBody.List, clause)
 		clause = getParamClause(idx, recvType)
@@ -689,19 +804,24 @@ func (ip *InstrumentPhase) rewriteHookContext() {
 	}
 	splitParams := ast.SplitMultiNameFields(ip.targetFunc.Type.Params)
 	for _, param := range splitParams.List {
-		paramType := replaceTypeParamsWithAny(desugarType(param), ip.targetFunc.Type.TypeParams)
+		paramType := replaceTypeParamsWithAny(desugarType(param), combinedTypeParams)
 		clause := setParamClause(idx, paramType)
 		methodSetParamBody.List = append(methodSetParamBody.List, clause)
 		clause = getParamClause(idx, paramType)
 		methodGetParamBody.List = append(methodGetParamBody.List, clause)
 		idx++
 	}
-	// Rewrite GetReturnVal and SetReturnVal methods
+}
+
+func (ip *InstrumentPhase) rewriteHookContextResults(
+	methodSetRetValBody, methodGetRetValBody *dst.BlockStmt,
+	combinedTypeParams *dst.FieldList,
+) {
 	if ip.targetFunc.Type.Results != nil {
-		idx = 0
+		idx := 0
 		splitResults := ast.SplitMultiNameFields(ip.targetFunc.Type.Results)
 		for _, retval := range splitResults.List {
-			retType := replaceTypeParamsWithAny(desugarType(retval), ip.targetFunc.Type.TypeParams)
+			retType := replaceTypeParamsWithAny(desugarType(retval), combinedTypeParams)
 			clause := getReturnValClause(idx, retType)
 			methodGetRetValBody.List = append(methodGetRetValBody.List, clause)
 			clause = setReturnValClause(idx, retType)
@@ -752,6 +872,13 @@ func replaceTypeParamsWithAny(t dst.Expr, typeParams *dst.FieldList) dst.Expr {
 			Key:   replaceTypeParamsWithAny(tType.Key, typeParams),
 			Value: replaceTypeParamsWithAny(tType.Value, typeParams),
 		}
+	case *dst.IndexExpr:
+		// GenStruct[T] -> interface{} (for generic receiver methods)
+		// The hook function expects interface{} for generic types
+		return ast.InterfaceType()
+	case *dst.IndexListExpr:
+		// GenStruct[T, U] -> interface{} (for generic receiver methods with multiple type params)
+		return ast.InterfaceType()
 	}
 	return t
 }