Fix type checker to visit all arguments of func even in args count mi…

…smatch

checker/checker.go

@@ -906,28 +906,38 @@ func (v *checker) checkArguments(
 		fnInOffset = 1
 	}
 
+	var err *file.Error
 	if fn.IsVariadic() {
 		if len(arguments) < fnNumIn-1 {
-			return anyType, &file.Error{
+			err = &file.Error{
 				Location: node.Location(),
 				Message:  fmt.Sprintf("not enough arguments to call %v", name),
 			}
 		}
 	} else {
 		if len(arguments) > fnNumIn {
-			return anyType, &file.Error{
+			err = &file.Error{
 				Location: node.Location(),
 				Message:  fmt.Sprintf("too many arguments to call %v", name),
 			}
 		}
 		if len(arguments) < fnNumIn {
-			return anyType, &file.Error{
+			err = &file.Error{
 				Location: node.Location(),
 				Message:  fmt.Sprintf("not enough arguments to call %v", name),
 			}
 		}
 	}
 
+	if err != nil {
+		// If we have an error, we should still visit all arguments to
+		// type check them, as a patch can fix the error later.
+		for _, arg := range arguments {
+			_, _ = v.visit(arg)
+		}
+		return fn.Out(0), err
+	}
+
 	for i, arg := range arguments {
 		t, _ := v.visit(arg)
 

checker/checker_test.go

@@ -170,9 +170,9 @@ type mock.Foo has no field bar (1:4)
  | Foo['bar']
  | ...^
 
-Foo.Method(Not)
+Foo.Method(42)
 too many arguments to call Method (1:5)
- | Foo.Method(Not)
+ | Foo.Method(42)
  | ....^
 
 Foo.Bar()
@@ -210,9 +210,9 @@ array elements can only be selected using an integer (got string) (1:12)
  | ArrayOfFoo.Not
  | ...........^
 
-FuncParam(Not)
+FuncParam(true)
 not enough arguments to call FuncParam (1:1)
- | FuncParam(Not)
+ | FuncParam(true)
  | ^
 
 MapOfFoo['str'].Not

patcher/with_context.go

@@ -13,10 +13,12 @@ type WithContext struct {
 
 // Visit adds WithContext.Name argument to all functions calls with a context.Context argument.
 func (w WithContext) Visit(node *ast.Node) {
-	switch (*node).(type) {
+	switch call := (*node).(type) {
 	case *ast.CallNode:
-		call := (*node).(*ast.CallNode)
 		fn := call.Callee.Type()
+		if fn == nil {
+			return
+		}
 		if fn.Kind() != reflect.Func {
 			return
 		}

patcher/with_context_test.go

@@ -104,3 +104,27 @@ func TestWithContext_with_env_method_chain(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, int64(42), output)
 }
+
+func TestWithContext_issue529(t *testing.T) {
+	env := map[string]any{
+		"ctx": context.Background(),
+		"foo": func(ctx context.Context, n int) int {
+			if ctx == nil {
+				panic("wanted a context")
+			}
+			return n + 1
+		},
+	}
+	options := []expr.Option{
+		expr.Env(env),
+		expr.WithContext("ctx"),
+	}
+
+	code := "foo(0) | foo()"
+	program, err := expr.Compile(code, options...)
+	require.NoError(t, err)
+
+	out, err := expr.Run(program, env)
+	require.NoError(t, err)
+	require.Equal(t, 2, out)
+}