`a3646b331c9c`

source

crates/jrsonnet-stdlib/src/sort.rs5.5 KiBsource history
1#![allow(non_snake_case)]23use std::cmp::Ordering;45use jrsonnet_evaluator::{6	Result, Thunk, Val, bail,7	function::builtin,8	val::{ArrValue, equals},9};1011use crate::{eval_on_empty, keyf::KeyF};1213#[derive(Copy, Clone)]14enum SortKeyType {15	Number,16	String,17	Unspecialized,18	Unknown,19}2021fn get_sort_type<T>(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result<SortKeyType> {22	let mut sort_type = SortKeyType::Unknown;23	for i in values {24		let i = key_getter(i);25		match (i, sort_type) {26			(Val::Str(_), SortKeyType::Unknown) => sort_type = SortKeyType::String,27			(Val::Num(_), SortKeyType::Unknown) => sort_type = SortKeyType::Number,28			(Val::Str(_), SortKeyType::String) | (Val::Num(_), SortKeyType::Number) => {}29			(Val::Str(_) | Val::Num(_), _) => {30				bail!("sort elements should have the same types")31			}32			(_, _) => return Ok(SortKeyType::Unspecialized),33		}34	}35	Ok(sort_type)36}3738fn sort_identity(mut values: Vec<Val>) -> Result<Vec<Val>> {39	// Fast path, identity key getter40	let sort_type = get_sort_type(&values, |k| k)?;41	match sort_type {42		SortKeyType::Number => values.sort_unstable_by_key(|v| match v {43			Val::Num(n) => *n,44			_ => unreachable!(),45		}),46		SortKeyType::String => values.sort_unstable_by_key(|v| match v {47			Val::Str(s) => s.clone(),48			_ => unreachable!(),49		}),50		SortKeyType::Unknown | SortKeyType::Unspecialized => {51			let mut err = None;52			// evaluate_compare_op will never return equal on types, which are different from53			// jsonnet perspective54			values.sort_unstable_by(|a, b| match Val::try_cmp(a, b) {55				Ok(ord) => ord,56				Err(e) if err.is_none() => {57					let _ = err.insert(e);58					Ordering::Equal59				}60				Err(_) => Ordering::Equal,61			});62			if let Some(err) = err {63				return Err(err);64			}65		}66	}67	Ok(values)68}6970fn sort_keyf(values: ArrValue, keyf: KeyF) -> Result<Vec<Thunk<Val>>> {71	// Slow path, user provided key getter72	let mut vk = Vec::with_capacity(values.len() as usize);73	for value in values.iter_lazy() {74		vk.push((value.clone(), keyf.eval(value)?));75	}76	let sort_type = get_sort_type(&vk, |v| &v.1)?;77	match sort_type {78		SortKeyType::Number => vk.sort_by_key(|v| match v.1 {79			Val::Num(n) => n,80			_ => unreachable!(),81		}),82		SortKeyType::String => vk.sort_by_key(|v| match &v.1 {83			Val::Str(s) => s.clone(),84			_ => unreachable!(),85		}),86		SortKeyType::Unknown | SortKeyType::Unspecialized => {87			let mut err = None;88			// evaluate_compare_op will never return equal on types, which are different from89			// jsonnet perspective90			vk.sort_by(|(_a, ak), (_b, bk)| match Val::try_cmp(ak, bk) {91				Ok(ord) => ord,92				Err(e) if err.is_none() => {93					let _ = err.insert(e);94					Ordering::Equal95				}96				Err(_) => Ordering::Equal,97			});98			if let Some(err) = err {99				return Err(err);100			}101		}102	}103	Ok(vk.into_iter().map(|v| v.0).collect())104}105106/// * `key_getter` - None, if identity sort required107pub fn sort(values: ArrValue, key_getter: KeyF) -> Result<ArrValue> {108	if values.len() <= 1 {109		return Ok(values);110	}111	if key_getter.is_identity() {112		Ok(ArrValue::new(sort_identity(113			values.iter().collect::<Result<Vec<Val>>>()?,114		)?))115	} else {116		Ok(ArrValue::new(sort_keyf(values, key_getter)?))117	}118}119120#[builtin]121pub fn builtin_sort(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {122	super::sort::sort(arr, keyF)123}124125fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {126	let mut out = Vec::new();127	let mut last = arr[0].clone();128	out.push(last.clone());129	for next in arr.into_iter().skip(1) {130		if !equals(&last, &next)? {131			out.push(next.clone());132		}133		last = next;134	}135	Ok(out)136}137138fn uniq_keyf(arr: ArrValue, keyf: KeyF) -> Result<Vec<Thunk<Val>>> {139	let mut out = Vec::new();140	let last_value = arr.get_lazy(0).unwrap();141	let mut last_key = keyf.eval(last_value.clone())?;142	out.push(last_value);143144	for next in arr.iter_lazy().skip(1) {145		let next_key = keyf.eval(next.clone())?;146		if !equals(&last_key, &next_key)? {147			out.push(next.clone());148		}149		last_key = next_key;150	}151	Ok(out)152}153154#[builtin]155#[allow(non_snake_case)]156pub fn builtin_uniq(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {157	if arr.len() <= 1 {158		return Ok(arr);159	}160	if keyF.is_identity() {161		Ok(ArrValue::new(uniq_identity(162			arr.iter().collect::<Result<Vec<Val>>>()?,163		)?))164	} else {165		Ok(ArrValue::new(uniq_keyf(arr, keyF)?))166	}167}168169#[builtin]170#[allow(non_snake_case)]171pub fn builtin_set(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {172	if arr.len() <= 1 {173		return Ok(arr);174	}175	if keyF.is_identity() {176		let arr = arr.iter().collect::<Result<Vec<Val>>>()?;177		let arr = sort_identity(arr)?;178		let arr = uniq_identity(arr)?;179		Ok(ArrValue::new(arr))180	} else {181		let arr = sort_keyf(arr, keyF.clone())?;182		let arr = uniq_keyf(ArrValue::new(arr), keyF)?;183		Ok(ArrValue::new(arr))184	}185}186187fn array_top1(arr: ArrValue, keyf: KeyF, ordering: Ordering) -> Result<Val> {188	let mut iter = arr.iter();189	let mut min = iter.next().expect("not empty")?;190	let mut min_key = keyf.eval(Thunk::evaluated(min.clone()))?;191	for item in iter {192		let cur = item?;193		let cur_key = keyf.eval(Thunk::evaluated(cur.clone()))?;194		if Val::try_cmp(&cur_key, &min_key)? == ordering {195			min = cur;196			min_key = cur_key;197		}198	}199	Ok(min)200}201202#[builtin]203pub fn builtin_min_array(204	arr: ArrValue,205	#[default] keyF: KeyF,206	onEmpty: Option<Thunk<Val>>,207) -> Result<Val> {208	if arr.is_empty() {209		return eval_on_empty(onEmpty);210	}211	array_top1(arr, keyF, Ordering::Less)212}213#[builtin]214pub fn builtin_max_array(215	arr: ArrValue,216	#[default] keyF: KeyF,217	onEmpty: Option<Thunk<Val>>,218) -> Result<Val> {219	if arr.is_empty() {220		return eval_on_empty(onEmpty);221	}222	array_top1(arr, keyF, Ordering::Greater)223}