`d349b9ef06a9`

source

crates/jrsonnet-stdlib/src/sort.rs6.3 KiBsource history
1#![allow(non_snake_case)]23use std::cmp::Ordering;45use jrsonnet_evaluator::{6	bail,7	function::{builtin, FuncVal},8	operator::evaluate_compare_op,9	val::{equals, ArrValue},10	Result, Thunk, Val,11};12use jrsonnet_parser::BinaryOpType;1314use crate::eval_on_empty;1516#[derive(Copy, Clone)]17enum SortKeyType {18	Number,19	String,20	Unknown,21}2223#[derive(PartialEq)]24struct NonNaNf64(f64);25impl PartialOrd for NonNaNf64 {26	fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {27		Some(self.cmp(other))28	}29}30impl Eq for NonNaNf64 {}31impl Ord for NonNaNf64 {32	fn cmp(&self, other: &Self) -> std::cmp::Ordering {33		self.0.partial_cmp(&other.0).expect("non nan")34	}35}3637fn get_sort_type<T>(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result<SortKeyType> {38	let mut sort_type = SortKeyType::Unknown;39	for i in values.iter() {40		let i = key_getter(i);41		match (i, sort_type) {42			(Val::Str(_), SortKeyType::Unknown) => sort_type = SortKeyType::String,43			(Val::Num(_), SortKeyType::Unknown) => sort_type = SortKeyType::Number,44			(Val::Str(_), SortKeyType::String) | (Val::Num(_), SortKeyType::Number) => {}45			(Val::Str(_) | Val::Num(_), _) => {46				bail!("sort elements should have the same types")47			}48			_ => {}49		}50	}51	Ok(sort_type)52}5354fn sort_identity(mut values: Vec<Val>) -> Result<Vec<Val>> {55	// Fast path, identity key getter56	let sort_type = get_sort_type(&values, |k| k)?;57	match sort_type {58		SortKeyType::Number => values.sort_unstable_by_key(|v| match v {59			Val::Num(n) => NonNaNf64(*n),60			_ => unreachable!(),61		}),62		SortKeyType::String => values.sort_unstable_by_key(|v| match v {63			Val::Str(s) => s.clone(),64			_ => unreachable!(),65		}),66		SortKeyType::Unknown => {67			let mut err = None;68			// evaluate_compare_op will never return equal on types, which are different from69			// jsonnet perspective70			values.sort_unstable_by(|a, b| match evaluate_compare_op(a, b, BinaryOpType::Lt) {71				Ok(ord) => ord,72				Err(e) if err.is_none() => {73					let _ = err.insert(e);74					Ordering::Equal75				}76				Err(_) => Ordering::Equal,77			});78			if let Some(err) = err {79				return Err(err);80			}81		}82	};83	Ok(values)84}8586fn sort_keyf(values: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {87	// Slow path, user provided key getter88	let mut vk = Vec::with_capacity(values.len());89	for value in values.iter_lazy() {90		vk.push((91			value.clone(),92			keyf.evaluate_simple(&(value.clone(),), false)?,93		));94	}95	let sort_type = get_sort_type(&vk, |v| &v.1)?;96	match sort_type {97		SortKeyType::Number => vk.sort_by_key(|v| match v.1 {98			Val::Num(n) => NonNaNf64(n),99			_ => unreachable!(),100		}),101		SortKeyType::String => vk.sort_by_key(|v| match &v.1 {102			Val::Str(s) => s.clone(),103			_ => unreachable!(),104		}),105		SortKeyType::Unknown => {106			let mut err = None;107			// evaluate_compare_op will never return equal on types, which are different from108			// jsonnet perspective109			vk.sort_by(110				|(_a, ak), (_b, bk)| match evaluate_compare_op(ak, bk, BinaryOpType::Lt) {111					Ok(ord) => ord,112					Err(e) if err.is_none() => {113						let _ = err.insert(e);114						Ordering::Equal115					}116					Err(_) => Ordering::Equal,117				},118			);119			if let Some(err) = err {120				return Err(err);121			}122		}123	};124	Ok(vk.into_iter().map(|v| v.0).collect())125}126127/// * `key_getter` - None, if identity sort required128pub fn sort(values: ArrValue, key_getter: FuncVal) -> Result<ArrValue> {129	if values.len() <= 1 {130		return Ok(values);131	}132	if key_getter.is_identity() {133		Ok(ArrValue::eager(sort_identity(134			values.iter().collect::<Result<Vec<Val>>>()?,135		)?))136	} else {137		Ok(ArrValue::lazy(sort_keyf(values, key_getter)?))138	}139}140141#[builtin]142pub fn builtin_sort(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {143	super::sort::sort(arr, keyF.unwrap_or_else(FuncVal::identity))144}145146fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {147	let mut out = Vec::new();148	let mut last = arr[0].clone();149	out.push(last.clone());150	for next in arr.into_iter().skip(1) {151		if !equals(&last, &next)? {152			out.push(next.clone());153		}154		last = next;155	}156	Ok(out)157}158159fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {160	let mut out = Vec::new();161	let last_value = arr.get_lazy(0).unwrap();162	let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?;163	out.push(last_value);164165	for next in arr.iter_lazy().skip(1) {166		let next_key = keyf.evaluate_simple(&(next.clone(),), false)?;167		if !equals(&last_key, &next_key)? {168			out.push(next.clone());169		}170		last_key = next_key;171	}172	Ok(out)173}174175#[builtin]176#[allow(non_snake_case)]177pub fn builtin_uniq(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {178	if arr.len() <= 1 {179		return Ok(arr);180	}181	let keyF = keyF.unwrap_or(FuncVal::identity());182	if keyF.is_identity() {183		Ok(ArrValue::eager(uniq_identity(184			arr.iter().collect::<Result<Vec<Val>>>()?,185		)?))186	} else {187		Ok(ArrValue::lazy(uniq_keyf(arr, keyF)?))188	}189}190191#[builtin]192#[allow(non_snake_case)]193pub fn builtin_set(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {194	if arr.len() <= 1 {195		return Ok(arr);196	}197	let keyF = keyF.unwrap_or(FuncVal::identity());198	if keyF.is_identity() {199		let arr = arr.iter().collect::<Result<Vec<Val>>>()?;200		let arr = sort_identity(arr)?;201		let arr = uniq_identity(arr)?;202		Ok(ArrValue::eager(arr))203	} else {204		let arr = sort_keyf(arr, keyF.clone())?;205		let arr = uniq_keyf(ArrValue::lazy(arr), keyF)?;206		Ok(ArrValue::lazy(arr))207	}208}209210fn eval_keyf(val: Val, key_f: &Option<FuncVal>) -> Result<Val> {211	if let Some(key_f) = key_f {212		key_f.evaluate_simple(&(val,), false)213	} else {214		Ok(val)215	}216}217218fn array_top1(arr: ArrValue, key_f: Option<FuncVal>, ordering: Ordering) -> Result<Val> {219	let mut iter = arr.iter();220	let mut min = iter.next().expect("not empty")?;221	let mut min_key = eval_keyf(min.clone(), &key_f)?;222	for item in iter {223		let cur = item?;224		let cur_key = eval_keyf(cur.clone(), &key_f)?;225		if evaluate_compare_op(&cur_key, &min_key, BinaryOpType::Lt)? == ordering {226			min = cur;227			min_key = cur_key;228		}229	}230	Ok(min)231}232233#[builtin]234pub fn builtin_min_array(235	arr: ArrValue,236	keyF: Option<FuncVal>,237	onEmpty: Option<Thunk<Val>>,238) -> Result<Val> {239	if arr.is_empty() {240		return eval_on_empty(onEmpty);241	}242	array_top1(arr, keyF, Ordering::Less)243}244#[builtin]245pub fn builtin_max_array(246	arr: ArrValue,247	keyF: Option<FuncVal>,248	onEmpty: Option<Thunk<Val>>,249) -> Result<Val> {250	if arr.is_empty() {251		return eval_on_empty(onEmpty);252	}253	array_top1(arr, keyF, Ordering::Greater)254}