diff --git a/lib/vrl/compiler/src/type_def.rs b/lib/vrl/compiler/src/type_def.rs index 38b09aab41b0b..215cca6a836d7 100644 --- a/lib/vrl/compiler/src/type_def.rs +++ b/lib/vrl/compiler/src/type_def.rs @@ -436,6 +436,36 @@ impl TypeKind { Object(_) => Kind::Object, } } + + /// Collects the kinds into a single kind. + /// Array and objects may have different kinds for each key/index, this collects those + /// into a single kind. + pub fn collect_kinds(self) -> TypeKind { + match self { + TypeKind::Array(kinds) => { + let mut newkinds = BTreeMap::new(); + newkinds.insert( + Index::Any, + kinds + .into_iter() + .fold(KindInfo::Unknown, |acc, (_, k)| acc.merge(k, false, true)), + ); + TypeKind::Array(newkinds) + } + + TypeKind::Object(kinds) => { + let mut newkinds = BTreeMap::new(); + newkinds.insert( + Field::Any, + kinds + .into_iter() + .fold(KindInfo::Unknown, |acc, (_, k)| acc.merge(k, false, true)), + ); + TypeKind::Object(newkinds) + } + _ => self, + } + } } #[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Ord)] @@ -784,6 +814,20 @@ impl TypeDef { self } + /// Collects any subtypes that can contain multiple indexed types (array, object) and collects them into + /// a single type for all indexes. + /// Used for functions that cant determine which indexes of a collection have been used in the result. + pub fn collect_subtypes(mut self) -> Self { + self.kind = match self.kind { + KindInfo::Known(set) => { + KindInfo::Known(set.into_iter().map(|k| k.collect_kinds()).collect()) + } + v => v, + }; + + self + } + #[inline] pub fn is_unknown(&self) -> bool { matches!(self.kind, KindInfo::Unknown) @@ -934,6 +978,35 @@ mod tests { use path::{self, Segment}; use std::str::FromStr; + #[test] + fn collect_subtypes() { + let kind = TypeKind::Array({ + let mut set1 = BTreeSet::new(); + set1.insert(TypeKind::Integer); + let mut set2 = BTreeSet::new(); + set2.insert(TypeKind::Bytes); + + let mut map = BTreeMap::new(); + map.insert(Index::Index(1), KindInfo::Known(set1)); + map.insert(Index::Index(2), KindInfo::Known(set2)); + map + }); + + let kind = kind.collect_kinds(); + + let expected = TypeKind::Array({ + let mut set = BTreeSet::new(); + set.insert(TypeKind::Integer); + set.insert(TypeKind::Bytes); + + let mut map = BTreeMap::new(); + map.insert(Index::Any, KindInfo::Known(set)); + map + }); + + assert_eq!(kind, expected); + } + mod kind_info { use super::*; diff --git a/lib/vrl/stdlib/src/slice.rs b/lib/vrl/stdlib/src/slice.rs index 7c048bb4625b8..df02372674a59 100644 --- a/lib/vrl/stdlib/src/slice.rs +++ b/lib/vrl/stdlib/src/slice.rs @@ -114,7 +114,8 @@ impl Expression for SliceFn { let td = TypeDef::new().fallible(); match self.value.type_def(state) { - v if v.is_bytes() || v.is_array() => td.merge(v), + v if v.is_bytes() => td.merge(v), + v if v.is_array() => td.merge(v).collect_subtypes(), _ => td.bytes().add_array_mapped::<(), Kind>(map! { (): Kind::all(), }), @@ -218,10 +219,7 @@ mod tests { start: 0 ], want: Ok(vec![0, 1, 2]), - tdef: TypeDef::new().fallible().array_mapped::(map! { 0: Kind::Integer, - 1: Kind::Integer, - 2: Kind::Integer - }), + tdef: TypeDef::new().fallible().array_mapped::<(), Kind>(map! { (): Kind::Integer }), } array_1 { @@ -229,11 +227,7 @@ mod tests { start: 1 ], want: Ok(vec![1, 2]), - // TODO: This is wrong! See https://github.com/timberio/vector/issues/6676 - tdef: TypeDef::new().fallible().array_mapped::(map! { 0: Kind::Integer, - 1: Kind::Integer, - 2: Kind::Integer - }), + tdef: TypeDef::new().fallible().array_mapped::<(), Kind>(map! { (): Kind::Integer }), } array_minus_2 { @@ -241,11 +235,17 @@ mod tests { start: -2 ], want: Ok(vec![1, 2]), - // TODO: This is wrong! See https://github.com/timberio/vector/issues/6676 - tdef: TypeDef::new().fallible().array_mapped::(map! { 0: Kind::Integer, - 1: Kind::Integer, - 2: Kind::Integer - }), + tdef: TypeDef::new().fallible().array_mapped::<(), Kind>(map! { (): Kind::Integer }), + } + + array_mixed_types { + args: func_args![value: value!([0, "ook", true]), + start: 1 + ], + want: Ok(value!(["ook", true])), + tdef: TypeDef::new().fallible().array_mapped::<(), Kind>( + map! { (): Kind::Integer | Kind::Bytes | Kind::Boolean } + ), } error_after_end {