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 usage(&self) -> &'static str {
51        indoc! {"
52            Filter elements from a collection.
53
54            This function currently *does not* support recursive iteration.
55
56            The function uses the function closure syntax to allow reading
57            the key-value or index-value combination for each item in the
58            collection.
59
60            The same scoping rules apply to closure blocks as they do for
61            regular blocks. This means that any variable defined in parent scopes
62            is accessible, and mutations to those variables are preserved,
63            but any new variables instantiated in the closure block are
64            unavailable outside of the block.
65
66            See the examples below to learn about the closure syntax.
67        "}
68    }
69
70    fn category(&self) -> &'static str {
71        Category::Enumerate.as_ref()
72    }
73
74    fn return_kind(&self) -> u16 {
75        kind::ARRAY | kind::OBJECT
76    }
77
78    fn parameters(&self) -> &'static [Parameter] {
79        const PARAMETERS: &[Parameter] = &[Parameter::required(
80            "value",
81            kind::OBJECT | kind::ARRAY,
82            "The array or object to filter.",
83        )];
84        PARAMETERS
85    }
86
87    fn examples(&self) -> &'static [Example] {
88        &[
89            example! {
90                title: "Filter elements",
91                source: indoc! {r#"
92                    . = { "tags": ["foo", "bar", "foo", "baz"] }
93                    filter(array(.tags)) -> |_index, value| {
94                        value != "foo"
95                    }
96                "#},
97                result: Ok(r#"["bar", "baz"]"#),
98            },
99            example! {
100                title: "Filter object",
101                source: r#"filter({ "a": 1, "b": 2 }) -> |key, _value| { key == "a" }"#,
102                result: Ok(r#"{ "a": 1 }"#),
103            },
104            example! {
105                title: "Filter array",
106                source: "filter([1, 2]) -> |_index, value| { value < 2 }",
107                result: Ok("[1]"),
108            },
109        ]
110    }
111
112    fn compile(
113        &self,
114        _state: &state::TypeState,
115        _ctx: &mut FunctionCompileContext,
116        arguments: ArgumentList,
117    ) -> Compiled {
118        let value = arguments.required("value");
119        let closure = arguments.required_closure()?;
120
121        Ok(FilterFn { value, closure }.as_expr())
122    }
123
124    fn closure(&self) -> Option<closure::Definition> {
125        use closure::{Definition, Input, Output, Variable, VariableKind};
126
127        Some(Definition {
128            inputs: vec![Input {
129                parameter_keyword: "value",
130                kind: Kind::object(Collection::any()).or_array(Collection::any()),
131                variables: vec![
132                    Variable {
133                        kind: VariableKind::TargetInnerKey,
134                    },
135                    Variable {
136                        kind: VariableKind::TargetInnerValue,
137                    },
138                ],
139                output: Output::Kind(Kind::boolean()),
140                example: example! {
141                    title: "filter array",
142                    source: "filter([1, 2]) -> |index, _value| { index == 0 }",
143                    result: Ok("[1]"),
144                },
145            }],
146            is_iterator: true,
147        })
148    }
149}
150
151#[derive(Debug, Clone)]
152struct FilterFn {
153    value: Box<dyn Expression>,
154    closure: Closure,
155}
156
157impl FunctionExpression for FilterFn {
158    fn resolve(&self, ctx: &mut Context) -> ExpressionResult<Value> {
159        let value = self.value.resolve(ctx)?;
160        let Closure {
161            variables,
162            block,
163            block_type_def: _,
164        } = &self.closure;
165        let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));
166
167        filter(value, ctx, &runner)
168    }
169
170    fn type_def(&self, ctx: &state::TypeState) -> TypeDef {
171        let mut type_def = self.value.type_def(ctx);
172
173        // Erase any type information from the array or object, as we can't know
174        // which elements are removed at runtime.
175        if type_def.contains_array() {
176            type_def.kind_mut().add_array(Collection::any());
177        }
178
179        if type_def.contains_object() {
180            type_def.kind_mut().add_object(Collection::any());
181        }
182
183        type_def
184    }
185}