Skip to main content

nu_protocol/
collection_columns.rs

1use crate::{CompareTypes, TypeRelation, TypeSet};
2use serde::{Deserialize, Serialize};
3use std::fmt::Display;
4
5#[allow(unused_imports)]
6use crate::SyntaxShape;
7
8/// Very basic ordered mapping, essentially a list of pairs.
9///
10/// Handles logic common to [`SyntaxShape::Record`], [`SyntaxShape::Table`], [`Type::Record`],
11/// [`Type::Table`], and possibly any other ordered mapping.
12///
13/// Implements [`Display`] for `T: Display`:
14/// ```rust
15/// # use nu_protocol::{CollectionColumns, Type};
16/// let cols = CollectionColumns::from(vec![
17///     ("a".to_string(), 1),
18///     ("b".to_string(), 2),
19/// ]);
20/// assert_eq!(cols.to_string(), "<a: 1, b: 2>");
21/// ```
22///
23/// Type widening (union) for [`Type`]:
24/// ```rust
25/// # use nu_protocol::{CollectionColumns, Type, TypeSet};
26/// let foo = CollectionColumns::from(vec![
27///     ("a".to_string(), Type::Int),
28///     ("b".to_string(), Type::String),
29/// ]);
30/// let bar = CollectionColumns::from(vec![
31///     ("a".to_string(), Type::Float),
32///     ("b".to_string(), Type::Int),
33///     ("c".to_string(), Type::Date),
34/// ]);
35/// assert_eq!(
36///     foo.union(bar),
37///     CollectionColumns::from(vec![
38///         ("a".to_string(), Type::Number),
39///         ("b".to_string(), Type::one_of([Type::String, Type::Int])),
40///     ])
41/// );
42/// ```
43#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash, Ord, PartialOrd)]
44#[serde(transparent)]
45pub struct CollectionColumns<T> {
46    fields: Box<[(String, T)]>,
47}
48
49impl<T> CollectionColumns<T> {
50    pub fn get<'s>(&'s self, key: &'_ str) -> Option<&'s T> {
51        self.iter()
52            .find(|(name, _)| name == key)
53            .map(|(_, val)| val)
54    }
55
56    pub fn map<U>(&self, f: impl Fn(&T) -> U) -> CollectionColumns<U> {
57        self.iter().map(|(k, v)| (k.clone(), f(v))).collect()
58    }
59
60    pub fn iter(&self) -> impl Iterator<Item = &(String, T)> {
61        self.into_iter()
62    }
63
64    pub fn is_empty(&self) -> bool {
65        self.fields.is_empty()
66    }
67
68    pub fn len(&self) -> usize {
69        self.fields.len()
70    }
71}
72
73impl<T> IntoIterator for CollectionColumns<T> {
74    type Item = (String, T);
75    type IntoIter = std::vec::IntoIter<Self::Item>;
76
77    fn into_iter(self) -> Self::IntoIter {
78        self.fields.into_iter()
79    }
80}
81
82impl<'a, T> IntoIterator for &'a CollectionColumns<T> {
83    type Item = &'a (String, T);
84    type IntoIter = std::slice::Iter<'a, (String, T)>;
85
86    fn into_iter(self) -> Self::IntoIter {
87        self.fields.iter()
88    }
89}
90
91impl<T> FromIterator<(String, T)> for CollectionColumns<T> {
92    fn from_iter<I: IntoIterator<Item = (String, T)>>(iter: I) -> Self {
93        Self {
94            fields: iter.into_iter().collect(),
95        }
96    }
97}
98
99impl<T> From<Vec<(String, T)>> for CollectionColumns<T> {
100    fn from(value: Vec<(String, T)>) -> Self {
101        Self {
102            fields: value.into_boxed_slice(),
103        }
104    }
105}
106
107impl<'a, T, const N: usize> From<[(&'a str, T); N]> for CollectionColumns<T> {
108    fn from(value: [(&'a str, T); N]) -> Self {
109        value
110            .into_iter()
111            .map(|(k, v)| (String::from(k), v))
112            .collect()
113    }
114}
115
116impl<T> CollectionColumns<T>
117where
118    T: TypeSet + Clone,
119{
120    fn widen_fields(lhs: Box<[(String, T)]>, rhs: Box<[(String, T)]>) -> Box<[(String, T)]> {
121        if lhs.is_empty() || rhs.is_empty() {
122            return [].into();
123        }
124
125        // iterate the shorter list to reduce quadratic behaviour
126        let (small, big) = if lhs.len() <= rhs.len() {
127            (lhs, rhs)
128        } else {
129            (rhs, lhs)
130        };
131
132        const MAP_THRESH: usize = 16;
133        if big.len() > MAP_THRESH {
134            use std::collections::HashMap;
135            let mut big_map: HashMap<String, T> = big.into_iter().collect();
136            small
137                .into_iter()
138                .filter_map(|(col, typ)| big_map.remove(&col).map(|b_typ| (col, typ.union(b_typ))))
139                .collect()
140        } else {
141            small
142                .into_iter()
143                .filter_map(|(col, typ)| {
144                    big.iter()
145                        .find_map(|(b_col, b_typ)| (&col == b_col).then(|| b_typ.clone()))
146                        .map(|b_typ| (col, typ.union(b_typ)))
147                })
148                .collect()
149        }
150    }
151}
152
153fn element_comparison_helper<T, F, O>(
154    lhs: &CollectionColumns<T>,
155    rhs: &CollectionColumns<T>,
156    f: F,
157) -> impl Iterator<Item = Option<O>>
158where
159    T: CompareTypes,
160    F: Fn(&T, &T) -> Option<O>,
161{
162    lhs.iter()
163        .map(move |(lhs_key, lhs_ty)| match rhs.get(lhs_key) {
164            Some(rhs_ty) => f(lhs_ty, rhs_ty),
165            // if `lhs` has a field `rhs` doesn't despite having at most the same number of
166            // columns as `rhs` (see NOTE[1]) then the sets of their keys are disjoint. they
167            // can't have a subtyping relation
168            None => None,
169        })
170}
171
172impl<T> CompareTypes for CollectionColumns<T>
173where
174    T: CompareTypes,
175{
176    fn compare_types(&self, other: &Self) -> Option<TypeRelation> {
177        // for structural subtyping, each field in a type is a "requirement". less
178        // fields in the type => less requirements => is supertype of more types
179        // e.g.: `{a: any, b: any}` is a supertype of `{a: any, b: any, c: any}`
180        //
181        // for `self` to be a subtype of `other`:
182        // - `self` must have all fields required by `other`. extra fields in `self` are irrelevant
183        // - for field `a` in `other`, `self.a` must be a subtype of `other.a`
184
185        match (self.is_empty(), other.is_empty()) {
186            (true, true) => return Some(TypeRelation::Equal),
187            (true, false) => return Some(TypeRelation::Supertype),
188            (false, true) => return Some(TypeRelation::Subtype),
189            (false, false) => (),
190        }
191
192        // NOTE[1]: with regards to number of columns `lhs` <= `rhs`
193        let (flipped, eq, (lhs, rhs)) = match self.fields.len().cmp(&other.fields.len()) {
194            std::cmp::Ordering::Less => (false, false, (self, other)),
195            std::cmp::Ordering::Equal => (false, true, (self, other)),
196            std::cmp::Ordering::Greater => (true, false, (other, self)),
197        };
198
199        let start = match eq {
200            true => TypeRelation::Equal,
201            false => TypeRelation::Supertype,
202        };
203
204        let out = element_comparison_helper(lhs, rhs, |lhs_ty, rhs_ty| {
205            if lhs_ty.is_any() || rhs_ty.is_any() {
206                // Not really" equal", just used to continue without affecting the outcome.
207                Some(TypeRelation::Equal)
208            } else {
209                lhs_ty.compare_types(rhs_ty)
210            }
211        })
212        .try_fold(start, |acc, e| acc.combine(e?))?;
213
214        Some(match flipped {
215            true => out.reverse(),
216            false => out,
217        })
218    }
219
220    /// Our type system uses the empty record as both the bottom and the top type of records
221    fn is_any(&self) -> bool {
222        self.fields.is_empty()
223    }
224
225    fn is_assignable_to(&self, dst: &Self) -> bool {
226        let src = self;
227
228        (src.is_any() || dst.is_any())
229            || element_comparison_helper(dst, src, |dst_ty, src_ty| {
230                Some(src_ty.is_assignable_to(dst_ty))
231            })
232            .try_fold(true, |acc, e| Some(acc && (e?)))
233            .unwrap_or(false)
234    }
235}
236
237impl<T> TypeSet for CollectionColumns<T>
238where
239    T: TypeSet + Clone,
240{
241    fn union(self, other: Self) -> Self {
242        let Self {
243            fields: self_fields,
244        } = self;
245        let Self {
246            fields: other_fields,
247        } = other;
248
249        Self {
250            fields: Self::widen_fields(self_fields, other_fields),
251        }
252    }
253}
254
255impl<T> Default for CollectionColumns<T> {
256    fn default() -> Self {
257        Self {
258            fields: Default::default(),
259        }
260    }
261}
262
263impl<T> Display for CollectionColumns<T>
264where
265    T: Display,
266{
267    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
268        match self.fields.as_ref() {
269            [] => Ok(()),
270            [(name, shape), tail @ ..] => {
271                write!(f, "<{name}: {shape}")?;
272                for (name, shape) in tail {
273                    write!(f, ", {name}: {shape}")?;
274                }
275
276                write!(f, ">")?;
277                Ok(())
278            }
279        }
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use pretty_assertions::assert_eq;
286    use rstest::rstest;
287
288    use super::*;
289    use crate::Type;
290
291    #[rstest]
292    #[case(Some(TypeRelation::Equal), [], [])]
293    #[case(Some(TypeRelation::Equal),
294        [("a", Type::Int)],
295        [("a", Type::Int)],
296    )]
297    #[case(None,
298        [("a", Type::Int)],
299        [("b", Type::Int)],
300    )]
301    #[case(Some(TypeRelation::Supertype),
302        [("a", Type::Int), ("b", Type::Int)],
303        [("a", Type::Int), ("b", Type::Int), ("c", Type::Int)],
304    )]
305    #[case(None,
306        [("name", Type::String), ("attrs", Type::list(Type::Any)), ("desc", Type::String)],
307        [("attrs", Type::list(Type::String)), ("desc", Type::String)],
308    )]
309    fn relations(
310        #[case] expected: Option<TypeRelation>,
311        #[case] lhs: impl IntoIterator<Item = (&'static str, Type)>,
312        #[case] rhs: impl IntoIterator<Item = (&'static str, Type)>,
313    ) {
314        let lhs = lhs
315            .into_iter()
316            .map(|(k, ty)| (k.to_owned(), ty))
317            .collect::<CollectionColumns<Type>>();
318        let rhs = rhs
319            .into_iter()
320            .map(|(k, ty)| (k.to_owned(), ty))
321            .collect::<CollectionColumns<Type>>();
322
323        assert_eq!(lhs.compare_types(&rhs), expected);
324        assert_eq!(rhs.compare_types(&lhs), expected.map(TypeRelation::reverse));
325    }
326
327    #[rstest]
328    #[case(true,
329        [("name", Type::String), ("attrs", Type::list(Type::Any)), ("desc", Type::String)],
330        [("attrs", Type::list(Type::String)), ("desc", Type::String)],
331    )]
332    fn is_assignable_to(
333        #[case] expected: bool,
334        #[case] src: impl IntoIterator<Item = (&'static str, Type)>,
335        #[case] dst: impl IntoIterator<Item = (&'static str, Type)>,
336    ) {
337        let src = src
338            .into_iter()
339            .map(|(k, ty)| (k.to_owned(), ty))
340            .collect::<CollectionColumns<Type>>();
341        let dst = dst
342            .into_iter()
343            .map(|(k, ty)| (k.to_owned(), ty))
344            .collect::<CollectionColumns<Type>>();
345
346        assert_eq!(src.is_assignable_to(&dst), expected)
347    }
348}