`aee64ce2020d`

source

crates/jrsonnet-stdlib/src/sort.rs6.4 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_gcmodule::Cc;14use jrsonnet_parser::BinaryOpType;1516use crate::eval_on_empty;1718#[derive(Copy, Clone)]19enum SortKeyType {20	Number,21	String,22	Unknown,23}2425#[derive(PartialEq)]26struct NonNaNf64(f64);27impl PartialOrd for NonNaNf64 {28	fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {29		self.0.partial_cmp(&other.0)30	}31}32impl Eq for NonNaNf64 {}33impl Ord for NonNaNf64 {34	fn cmp(&self, other: &Self) -> std::cmp::Ordering {35		self.partial_cmp(other).expect("non nan")36	}37}3839fn get_sort_type<T>(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result<SortKeyType> {40	let mut sort_type = SortKeyType::Unknown;41	for i in values.iter() {42		let i = key_getter(i);43		match (i, sort_type) {44			(Val::Str(_), SortKeyType::Unknown) => sort_type = SortKeyType::String,45			(Val::Num(_), SortKeyType::Unknown) => sort_type = SortKeyType::Number,46			(Val::Str(_), SortKeyType::String) | (Val::Num(_), SortKeyType::Number) => {}47			(Val::Str(_) | Val::Num(_), _) => {48				throw!("sort elements should have the same types")49			}50			_ => {}51		}52	}53	Ok(sort_type)54}5556fn sort_identity(mut values: Vec<Val>) -> Result<Vec<Val>> {57	// Fast path, identity key getter58	let sort_type = get_sort_type(&values, |k| k)?;59	match sort_type {60		SortKeyType::Number => values.sort_unstable_by_key(|v| match v {61			Val::Num(n) => NonNaNf64(*n),62			_ => unreachable!(),63		}),64		SortKeyType::String => values.sort_unstable_by_key(|v| match v {65			Val::Str(s) => s.clone(),66			_ => unreachable!(),67		}),68		SortKeyType::Unknown => {69			let mut err = None;70			// evaluate_compare_op will never return equal on types, which are different from71			// jsonnet perspective72			values.sort_unstable_by(|a, b| match evaluate_compare_op(a, b, BinaryOpType::Lt) {73				Ok(ord) => ord,74				Err(e) if err.is_none() => {75					let _ = err.insert(e);76					Ordering::Equal77				}78				Err(_) => Ordering::Equal,79			});80			if let Some(err) = err {81				return Err(err);82			}83		}84	};85	Ok(values)86}8788fn sort_keyf(values: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {89	// Slow path, user provided key getter90	let mut vk = Vec::with_capacity(values.len());91	for value in values.iter_lazy() {92		vk.push((93			value.clone(),94			keyf.evaluate_simple(&(value.clone(),), false)?,95		));96	}97	let sort_type = get_sort_type(&vk, |v| &v.1)?;98	match sort_type {99		SortKeyType::Number => vk.sort_by_key(|v| match v.1 {100			Val::Num(n) => NonNaNf64(n),101			_ => unreachable!(),102		}),103		SortKeyType::String => vk.sort_by_key(|v| match &v.1 {104			Val::Str(s) => s.clone(),105			_ => unreachable!(),106		}),107		SortKeyType::Unknown => {108			let mut err = None;109			// evaluate_compare_op will never return equal on types, which are different from110			// jsonnet perspective111			vk.sort_by(112				|(_a, ak), (_b, bk)| match evaluate_compare_op(ak, bk, BinaryOpType::Lt) {113					Ok(ord) => ord,114					Err(e) if err.is_none() => {115						let _ = err.insert(e);116						Ordering::Equal117					}118					Err(_) => Ordering::Equal,119				},120			);121			if let Some(err) = err {122				return Err(err);123			}124		}125	};126	Ok(vk.into_iter().map(|v| v.0).collect())127}128129/// * `key_getter` - None, if identity sort required130pub fn sort(values: ArrValue, key_getter: FuncVal) -> Result<ArrValue> {131	if values.len() <= 1 {132		return Ok(values);133	}134	if key_getter.is_identity() {135		Ok(ArrValue::eager(sort_identity(136			values.iter().collect::<Result<Vec<Val>>>()?,137		)?))138	} else {139		Ok(ArrValue::lazy(Cc::new(sort_keyf(values, key_getter)?)))140	}141}142143#[builtin]144pub fn builtin_sort(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {145	super::sort::sort(arr, keyF.unwrap_or_else(FuncVal::identity))146}147148fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {149	let mut out = Vec::new();150	let mut last = arr[0].clone();151	out.push(last.clone());152	for next in arr.into_iter().skip(1) {153		if !equals(&last, &next)? {154			out.push(next.clone());155		}156		last = next;157	}158	Ok(out)159}160161fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {162	let mut out = Vec::new();163	let last_value = arr.get_lazy(0).unwrap();164	let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?;165	out.push(last_value);166167	for next in arr.iter_lazy().skip(1) {168		let next_key = keyf.evaluate_simple(&(next.clone(),), false)?;169		if !equals(&last_key, &next_key)? {170			out.push(next.clone());171		}172		last_key = next_key;173	}174	Ok(out)175}176177#[builtin]178#[allow(non_snake_case)]179pub fn builtin_uniq(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {180	if arr.len() <= 1 {181		return Ok(arr);182	}183	let keyF = keyF.unwrap_or(FuncVal::identity());184	if keyF.is_identity() {185		Ok(ArrValue::eager(uniq_identity(186			arr.iter().collect::<Result<Vec<Val>>>()?,187		)?))188	} else {189		Ok(ArrValue::lazy(Cc::new(uniq_keyf(arr, keyF)?)))190	}191}192193#[builtin]194#[allow(non_snake_case)]195pub fn builtin_set(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {196	if arr.len() <= 1 {197		return Ok(arr);198	}199	let keyF = keyF.unwrap_or(FuncVal::identity());200	if keyF.is_identity() {201		let arr = arr.iter().collect::<Result<Vec<Val>>>()?;202		let arr = sort_identity(arr)?;203		let arr = uniq_identity(arr)?;204		Ok(ArrValue::eager(arr))205	} else {206		let arr = sort_keyf(arr, keyF.clone())?;207		let arr = uniq_keyf(ArrValue::lazy(Cc::new(arr)), keyF)?;208		Ok(ArrValue::lazy(Cc::new(arr)))209	}210}211212213fn eval_keyf(val: Val, key_f: &Option<FuncVal>) -> Result<Val> {214	if let Some(key_f) = key_f {215		key_f.evaluate_simple(&(val,), false)216	} else {217		Ok(val)218	}219}220221fn array_top1(arr: ArrValue, key_f: Option<FuncVal>, ordering: Ordering) -> Result<Val> {222	let mut iter = arr.iter();223	let mut min = iter.next().expect("not empty")?;224	let mut min_key = eval_keyf(min.clone(), &key_f)?;225	for item in iter {226		let cur = item?;227		let cur_key = eval_keyf(cur.clone(), &key_f)?;228		if evaluate_compare_op(&cur_key, &min_key, BinaryOpType::Lt)? == ordering {229			min = cur;230			min_key = cur_key;231		}232	}233	Ok(min)234}235236#[builtin]237pub fn builtin_min_array(238	arr: ArrValue,239	keyF: Option<FuncVal>,240	onEmpty: Option<Thunk<Val>>,241) -> Result<Val> {242	if arr.is_empty() {243		return eval_on_empty(onEmpty);244	}245	array_top1(arr, keyF, Ordering::Less)246}247#[builtin]248pub fn builtin_max_array(249	arr: ArrValue,250	keyF: Option<FuncVal>,251	onEmpty: Option<Thunk<Val>>,252) -> Result<Val> {253	if arr.is_empty() {254		return eval_on_empty(onEmpty);255	}256	array_top1(arr, keyF, Ordering::Greater)257}