Fix type any type propagation with patcher relationship

Author: antonmedv

checker/checker.go

@@ -539,9 +539,22 @@ func (v *checker) SliceNode(node *ast.SliceNode) (reflect.Type, info) {
 
 func (v *checker) CallNode(node *ast.CallNode) (reflect.Type, info) {
 	t, i := v.functionReturnType(node)
-	if node.Type() != nil {
+
+	// Check if type was set on node (for example, by patcher)
+	// and use node type instead of function return type.
+	//
+	// If node type is anyType, then we should use function
+	// return type. For example, on error we return anyType
+	// for a call `errCall().Method()` and method will be
+	// evaluated on `anyType.Method()`, so return type will
+	// be anyType `anyType.Method(): anyType`. Patcher can
+	// fix `errCall()` to return proper type, so on second
+	// checker pass we should replace anyType on method node
+	// with new correct function return type.
+	if node.Type() != nil && node.Type() != anyType {
 		return node.Type(), i
 	}
+
 	return t, i
 }
 

patcher/with_context_test.go

@@ -61,3 +61,46 @@ func TestWithContext_with_env_Function(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 42, output)
 }
+
+type TestFoo struct {
+	contextValue int
+}
+
+func (f *TestFoo) GetValue(a int) int64 {
+	return int64(f.contextValue + a)
+}
+
+func TestWithContext_with_env_method_chain(t *testing.T) {
+	env := map[string]any{
+		"ctx": context.TODO(),
+	}
+
+	fn := expr.Function("fn",
+		func(params ...any) (any, error) {
+			ctx := params[0].(context.Context)
+
+			contextValue := ctx.Value("value").(int)
+
+			return &TestFoo{
+				contextValue: contextValue,
+			}, nil
+		},
+		new(func(context.Context) *TestFoo),
+	)
+
+	program, err := expr.Compile(
+		`fn().GetValue(40)`,
+		expr.Env(env),
+		expr.WithContext("ctx"),
+		fn,
+		expr.AsInt64(),
+	)
+	require.NoError(t, err)
+
+	ctx := context.WithValue(context.Background(), "value", 2)
+	env["ctx"] = ctx
+
+	output, err := expr.Run(program, env)
+	require.NoError(t, err)
+	require.Equal(t, int64(42), output)
+}