`f3c8ed70f69f`

source

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