fix(type system): Fix function closures not correctly updating TypeState

changelog.d/1399.fix.md

@@ -0,0 +1,3 @@
+Function closures now correctly update type state of the program.
+
+authors: zettroke

src/compiler/compiler.rs

@@ -84,6 +84,19 @@ impl CompilerError {
 }
 
 impl<'a> Compiler<'a> {
+    pub(crate) fn new(fns: &'a [Box<dyn Function>], config: CompileConfig) -> Self {
+        Self {
+            fns,
+            diagnostics: vec![],
+            fallible: false,
+            abortable: false,
+            external_queries: vec![],
+            external_assignments: vec![],
+            skip_missing_query_target: vec![],
+            fallible_expression_error: None,
+            config,
+        }
+    }
     /// Compiles a given source into the final [`Program`].
     ///
     /// # Arguments
@@ -106,17 +119,7 @@ impl<'a> Compiler<'a> {
         let initial_state = state.clone();
         let mut state = state.clone();
 
-        let mut compiler = Self {
-            fns,
-            diagnostics: vec![],
-            fallible: false,
-            abortable: false,
-            external_queries: vec![],
-            external_assignments: vec![],
-            skip_missing_query_target: vec![],
-            fallible_expression_error: None,
-            config,
-        };
+        let mut compiler = Compiler::new(fns, config);
         let expressions = compiler.compile_root_exprs(ast, &mut state);
 
         let (errors, warnings): (Vec<_>, Vec<_>) =
@@ -272,7 +275,7 @@ impl<'a> Compiler<'a> {
         Some(Group::new(expr))
     }
 
-    fn compile_root_exprs(
+    pub(crate) fn compile_root_exprs(
         &mut self,
         nodes: impl IntoIterator<Item = Node<ast::RootExpr>>,
         state: &mut TypeState,

src/compiler/expression/function_call.rs

@@ -500,30 +500,38 @@ impl<'a> Builder<'a> {
         state: &mut TypeState,
     ) -> Result<(Option<Closure>, bool), FunctionCallError> {
         // Check if we have a closure we need to compile.
-        if let Some((variables, input)) = self.closure.clone() {
+        if let Some((variables, input)) = &self.closure {
             // TODO: This assumes the closure will run exactly once, which is incorrect.
             // see: https://github.com/vectordotdev/vector/issues/13782
 
-            let block = closure_block.expect("closure must contain block");
-
+            let mut variables_types = vec![];
             // At this point, we've compiled the block, so we can remove the
             // closure variables from the compiler's local environment.
-            for ident in &variables {
-                match locals.remove_variable(ident) {
-                    Some(details) => state.local.insert_variable(ident.clone(), details),
-                    None => {
-                        state.local.remove_variable(ident);
-                    }
+            for ident in variables {
+                let variable_details = state
+                    .local
+                    .remove_variable(ident)
+                    .expect("Closure variable must be present");
+                variables_types.push(variable_details);
+
+                // If outer scope has this variable, restore its state
+                if let Some(details) = locals.remove_variable(ident) {
+                    state.local.insert_variable(ident.clone(), details);
                 }
             }
 
-            let (block_span, (block, block_type_def)) = block.take();
+            let (block_span, (block, block_type_def)) = closure_block
+                .ok_or(FunctionCallError::MissingClosure {
+                    call_span: Span::default(), // TODO can we provide a better span?
+                    example: None,
+                })?
+                .take();
 
             let closure_fallible = block_type_def.is_fallible();
 
             // Check the type definition of the resulting block.This needs to match
             // whatever is configured by the closure input type.
-            let expected_kind = input.output.into_kind();
+            let expected_kind = input.clone().output.into_kind();
             let found_kind = block_type_def
                 .kind()
                 .union(block_type_def.returns().clone());
@@ -536,7 +544,7 @@ impl<'a> Builder<'a> {
                 });
             }
 
-            let fnclosure = Closure::new(variables, block, block_type_def);
+            let fnclosure = Closure::new(variables.clone(), variables_types, block, block_type_def);
             self.list.set_closure(fnclosure.clone());
 
             // closure = Some(fnclosure);
@@ -699,6 +707,28 @@ impl Expression for FunctionCall {
 
         let mut expr_result = self.expr.apply_type_info(&mut state);
 
+        // Closure can change state of locals in our `state`, so we need to update it.
+        if let Some(closure) = &self.closure {
+            // To get correct `type_info()` from closure we need to add closure arguments into current state
+            let mut closure_state = state.clone();
+            dbg!(&closure.variables_types);
+            for (ident, details) in closure
+                .variables
+                .iter()
+                .cloned()
+                .zip(closure.variables_types.iter().cloned())
+            {
+                closure_state.local.insert_variable(ident, details);
+            }
+            let mut closure_info = closure.block.type_info(&closure_state);
+            // Interaction with closure arguments can't affect parent state, so remove them before merge
+            for ident in &closure.variables {
+                closure_info.state.local.remove_variable(ident);
+            }
+
+            state = state.merge(closure_info.state);
+        }
+
         // If one of the arguments only partially matches the function type
         // definition, then we mark the entire function as fallible.
         //
@@ -1247,9 +1277,10 @@ impl DiagnosticMessage for FunctionCallError {
 
 #[cfg(test)]
 mod tests {
-    use crate::compiler::{Category, FunctionExpression, value::kind};
-
     use super::*;
+    use crate::compiler::{Category, Compiler, FunctionExpression, value::kind};
+    use crate::parser::parse;
+    use crate::stdlib::ForEach;
 
     #[derive(Clone, Debug)]
     struct Fn;
@@ -1436,4 +1467,27 @@ mod tests {
 
         assert_eq!(Ok(expected), params);
     }
+
+    #[test]
+    fn closure_type_state() {
+        let program = parse(
+            r#"
+            v = ""
+
+            for_each({}) -> |key, value| {
+                v = 0
+            }
+        "#,
+        )
+        .unwrap();
+
+        let fns = vec![Box::new(ForEach) as Box<dyn Function>];
+        let mut compiler = Compiler::new(&fns, CompileConfig::default());
+
+        let mut state = TypeState::default();
+        compiler.compile_root_exprs(program, &mut state);
+        let var = state.local.variable(&Ident::new("v")).unwrap();
+
+        assert_eq!(var.type_def.kind(), &Kind::bytes().or_integer());
+    }
 }

src/compiler/function.rs

@@ -14,6 +14,7 @@ use super::{
     CompileConfig, Span, TypeDef,
     expression::{Block, Container, Expr, Expression, container::Variant},
     state::TypeState,
+    type_def::Details,
     value::{Kind, kind},
 };
 
@@ -627,15 +628,22 @@ mod test_impls {
 #[derive(Debug, Clone, PartialEq)]
 pub struct Closure {
     pub variables: Vec<Ident>,
+    pub variables_types: Vec<Details>,
     pub block: Block,
     pub block_type_def: TypeDef,
 }
 
 impl Closure {
     #[must_use]
-    pub fn new<T: Into<Ident>>(variables: Vec<T>, block: Block, block_type_def: TypeDef) -> Self {
+    pub fn new(
+        variables: Vec<Ident>,
+        variables_types: Vec<Details>,
+        block: Block,
+        block_type_def: TypeDef,
+    ) -> Self {
         Self {
-            variables: variables.into_iter().map(Into::into).collect(),
+            variables,
+            variables_types,
             block,
             block_type_def,
         }

src/compiler/program.rs

@@ -48,7 +48,7 @@ pub struct ProgramInfo {
     /// Returns whether the compiled program can fail at runtime.
     ///
     /// A program can only fail at runtime if the fallible-function-call
-    /// (`foo!()`) is used within the source.
+    /// (`foo!()`) is used within the source.vrl
     pub fallible: bool,
 
     /// Returns whether the compiled program can be aborted at runtime.

src/compiler/type_def.rs

@@ -538,7 +538,7 @@ impl From<TypeDef> for Kind {
 }
 
 #[derive(Debug, Clone, PartialEq)]
-pub(crate) struct Details {
+pub struct Details {
     pub(crate) type_def: TypeDef,
     pub(crate) value: Option<Value>,
 }

src/stdlib/filter.rs

@@ -161,6 +161,7 @@ impl FunctionExpression for FilterFn {
             variables,
             block,
             block_type_def: _,
+            ..
         } = &self.closure;
         let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));
 

src/stdlib/for_each.rs

@@ -155,6 +155,7 @@ impl FunctionExpression for ForEachFn {
             variables,
             block,
             block_type_def: _,
+            ..
         } = &self.closure;
         let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));
 

src/stdlib/map_keys.rs

@@ -191,6 +191,7 @@ impl FunctionExpression for MapKeysFn {
             variables,
             block,
             block_type_def: _,
+            ..
         } = &self.closure;
         let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));
 

src/stdlib/map_values.rs

@@ -180,6 +180,7 @@ impl FunctionExpression for MapValuesFn {
             variables,
             block,
             block_type_def: _,
+            ..
         } = &self.closure;
         let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));