`cfa49ab1b7b8`

source

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