git.delta.rocks / jrsonnet / refs/commits / e88030decfb4

1use jrsonnet_evaluator::{2	error::Result,3	function::{builtin, FuncVal},4	throw,5	val::{equals, ArrValue},6	Thunk, Val,7};8use jrsonnet_gcmodule::Cc;910#[derive(Copy, Clone)]11enum SortKeyType {12	Number,13	String,14	Unknown,15}1617#[derive(PartialEq)]18struct NonNaNf64(f64);19impl PartialOrd for NonNaNf64 {20	fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {21		self.0.partial_cmp(&other.0)22	}23}24impl Eq for NonNaNf64 {}25impl Ord for NonNaNf64 {26	fn cmp(&self, other: &Self) -> std::cmp::Ordering {27		self.partial_cmp(other).expect("non nan")28	}29}3031fn get_sort_type<T>(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result<SortKeyType> {32	let mut sort_type = SortKeyType::Unknown;33	for i in values.iter() {34		let i = key_getter(i);35		match (i, sort_type) {36			(Val::Str(_), SortKeyType::Unknown) => sort_type = SortKeyType::String,37			(Val::Num(_), SortKeyType::Unknown) => sort_type = SortKeyType::Number,38			(Val::Str(_), SortKeyType::String) | (Val::Num(_), SortKeyType::Number) => {}39			(Val::Str(_) | Val::Num(_), _) => {40				throw!("sort elements should have the same types")41			}42			_ => throw!("sort key should either be a string or a number"),43		}44	}45	Ok(sort_type)46}4748fn sort_identity(mut values: Vec<Val>) -> Result<Vec<Val>> {49	// Fast path, identity key getter50	let sort_type = get_sort_type(&values, |k| k)?;51	match sort_type {52		SortKeyType::Number => values.sort_unstable_by_key(|v| match v {53			Val::Num(n) => NonNaNf64(*n),54			_ => unreachable!(),55		}),56		SortKeyType::String => values.sort_unstable_by_key(|v| match v {57			Val::Str(s) => s.clone(),58			_ => unreachable!(),59		}),60		SortKeyType::Unknown => unreachable!(),61	};62	Ok(values)63}6465fn sort_keyf(values: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {66	// Slow path, user provided key getter67	let mut vk = Vec::with_capacity(values.len());68	for value in values.iter_lazy() {69		vk.push((70			value.clone(),71			keyf.evaluate_simple(&(value.clone(),), false)?,72		));73	}74	let sort_type = get_sort_type(&mut vk, |v| &v.1)?;75	match sort_type {76		SortKeyType::Number => vk.sort_by_key(|v| match v.1 {77			Val::Num(n) => NonNaNf64(n),78			_ => unreachable!(),79		}),80		SortKeyType::String => vk.sort_by_key(|v| match &v.1 {81			Val::Str(s) => s.clone(),82			_ => unreachable!(),83		}),84		SortKeyType::Unknown => unreachable!(),85	};86	Ok(vk.into_iter().map(|v| v.0).collect())87}8889/// * `key_getter` - None, if identity sort required90pub fn sort(values: ArrValue, key_getter: FuncVal) -> Result<ArrValue> {91	if values.len() <= 1 {92		return Ok(values);93	}94	if key_getter.is_identity() {95		Ok(ArrValue::eager(sort_identity(96			values.iter().collect::<Result<Vec<Val>>>()?,97		)?))98	} else {99		Ok(ArrValue::lazy(Cc::new(sort_keyf(values, key_getter)?)))100	}101}102103#[builtin]104#[allow(non_snake_case)]105pub fn builtin_sort(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {106	if arr.len() <= 1 {107		return Ok(arr);108	}109	Ok(super::sort::sort(110		arr,111		keyF.unwrap_or_else(FuncVal::identity),112	)?)113}114115fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {116	let mut out = Vec::new();117	let mut last = arr[0].clone();118	out.push(last.clone());119	for next in arr.into_iter().skip(1) {120		if !equals(&last, &next)? {121			out.push(next.clone());122		}123		last = next;124	}125	Ok(out)126}127128fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {129	let mut out = Vec::new();130	let last_value = arr.get_lazy(0).unwrap();131	let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?;132	out.push(last_value.clone());133134	for next in arr.iter_lazy().skip(1) {135		let next_key = keyf.evaluate_simple(&(next.clone(),), false)?;136		if !equals(&last_key, &next_key)? {137			out.push(next.clone());138		}139		last_key = next_key;140	}141	Ok(out)142}143144#[builtin]145#[allow(non_snake_case)]146pub fn builtin_uniq(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {147	if arr.len() <= 1 {148		return Ok(arr);149	}150	let keyF = keyF.unwrap_or(FuncVal::identity());151	if keyF.is_identity() {152		Ok(ArrValue::eager(uniq_identity(153			arr.iter().collect::<Result<Vec<Val>>>()?,154		)?))155	} else {156		Ok(ArrValue::lazy(Cc::new(uniq_keyf(arr, keyF)?)))157	}158}159160#[builtin]161#[allow(non_snake_case)]162pub fn builtin_set(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {163	let keyF = keyF.unwrap_or(FuncVal::identity());164	if keyF.is_identity() {165		let arr = arr.iter().collect::<Result<Vec<Val>>>()?;166		let arr = sort_identity(arr)?;167		let arr = uniq_identity(arr)?;168		Ok(ArrValue::eager(arr))169	} else {170		let arr = sort_keyf(arr, keyF.clone())?;171		let arr = uniq_keyf(ArrValue::lazy(Cc::new(arr)), keyF)?;172		Ok(ArrValue::lazy(Cc::new(arr)))173	}174}