vrl/stdlib/
filter.rs

1use crate::compiler::prelude::*;
2use std::collections::BTreeMap;
3
4fn filter<T>(value: Value, ctx: &mut Context, runner: &closure::Runner<T>) -> Resolved
5where
6    T: Fn(&mut Context) -> Resolved,
7{
8    match value {
9        Value::Object(object) => object
10            .into_iter()
11            .filter_map(
12                |(key, value)| match runner.run_key_value(ctx, &key, &value) {
13                    Ok(v) => v
14                        .as_boolean()
15                        .expect("compiler guarantees boolean return type")
16                        .then_some(Ok((key, value))),
17                    Err(err) => Some(Err(err)),
18                },
19            )
20            .collect::<ExpressionResult<BTreeMap<_, _>>>()
21            .map(Into::into),
22
23        Value::Array(array) => array
24            .into_iter()
25            .enumerate()
26            .filter_map(
27                |(index, value)| match runner.run_index_value(ctx, index, &value) {
28                    Ok(v) => v
29                        .as_boolean()
30                        .expect("compiler guarantees boolean return type")
31                        .then_some(Ok(value)),
32                    Err(err) => Some(Err(err)),
33                },
34            )
35            .collect::<ExpressionResult<Vec<_>>>()
36            .map(Into::into),
37
38        _ => Err("function requires collection types as input".into()),
39    }
40}
41
42#[derive(Clone, Copy, Debug)]
43pub struct Filter;
44
45impl Function for Filter {
46    fn identifier(&self) -> &'static str {
47        "filter"
48    }
49
50    fn parameters(&self) -> &'static [Parameter] {
51        &[Parameter {
52            keyword: "value",
53            kind: kind::OBJECT | kind::ARRAY,
54            required: true,
55        }]
56    }
57
58    fn examples(&self) -> &'static [Example] {
59        &[
60            Example {
61                title: "filter object",
62                source: r#"filter({ "a": 1, "b": 2 }) -> |key, _value| { key == "a" }"#,
63                result: Ok(r#"{ "a": 1 }"#),
64            },
65            Example {
66                title: "filter array",
67                source: "filter([1, 2]) -> |_index, value| { value < 2 }",
68                result: Ok("[1]"),
69            },
70        ]
71    }
72
73    fn compile(
74        &self,
75        _state: &state::TypeState,
76        _ctx: &mut FunctionCompileContext,
77        arguments: ArgumentList,
78    ) -> Compiled {
79        let value = arguments.required("value");
80        let closure = arguments.required_closure()?;
81
82        Ok(FilterFn { value, closure }.as_expr())
83    }
84
85    fn closure(&self) -> Option<closure::Definition> {
86        use closure::{Definition, Input, Output, Variable, VariableKind};
87
88        Some(Definition {
89            inputs: vec![Input {
90                parameter_keyword: "value",
91                kind: Kind::object(Collection::any()).or_array(Collection::any()),
92                variables: vec![
93                    Variable {
94                        kind: VariableKind::TargetInnerKey,
95                    },
96                    Variable {
97                        kind: VariableKind::TargetInnerValue,
98                    },
99                ],
100                output: Output::Kind(Kind::boolean()),
101                example: Example {
102                    title: "filter array",
103                    source: "filter([1, 2]) -> |index, _value| { index == 0 }",
104                    result: Ok("[1]"),
105                },
106            }],
107            is_iterator: true,
108        })
109    }
110}
111
112#[derive(Debug, Clone)]
113struct FilterFn {
114    value: Box<dyn Expression>,
115    closure: Closure,
116}
117
118impl FunctionExpression for FilterFn {
119    fn resolve(&self, ctx: &mut Context) -> ExpressionResult<Value> {
120        let value = self.value.resolve(ctx)?;
121        let Closure {
122            variables,
123            block,
124            block_type_def: _,
125        } = &self.closure;
126        let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));
127
128        filter(value, ctx, &runner)
129    }
130
131    fn type_def(&self, ctx: &state::TypeState) -> TypeDef {
132        let mut type_def = self.value.type_def(ctx);
133
134        // Erase any type information from the array or object, as we can't know
135        // which elements are removed at runtime.
136        if type_def.contains_array() {
137            type_def.kind_mut().add_array(Collection::any());
138        }
139
140        if type_def.contains_object() {
141            type_def.kind_mut().add_object(Collection::any());
142        }
143
144        type_def
145    }
146}