`b5d51b90d7ba`

source

crates/jrsonnet-stdlib/src/sort.rs6.2 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 {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(143	arr: ArrValue,144145	#[default(FuncVal::identity())] keyF: FuncVal,146) -> Result<ArrValue> {147	super::sort::sort(arr, keyF)148}149150fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {151	let mut out = Vec::new();152	let mut last = arr[0].clone();153	out.push(last.clone());154	for next in arr.into_iter().skip(1) {155		if !equals(&last, &next)? {156			out.push(next.clone());157		}158		last = next;159	}160	Ok(out)161}162163fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {164	let mut out = Vec::new();165	let last_value = arr.get_lazy(0).unwrap();166	let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?;167	out.push(last_value);168169	for next in arr.iter_lazy().skip(1) {170		let next_key = keyf.evaluate_simple(&(next.clone(),), false)?;171		if !equals(&last_key, &next_key)? {172			out.push(next.clone());173		}174		last_key = next_key;175	}176	Ok(out)177}178179#[builtin]180#[allow(non_snake_case)]181pub fn builtin_uniq(182	arr: ArrValue,183184	#[default(FuncVal::identity())] keyF: FuncVal,185) -> Result<ArrValue> {186	if arr.len() <= 1 {187		return Ok(arr);188	}189	if keyF.is_identity() {190		Ok(ArrValue::eager(uniq_identity(191			arr.iter().collect::<Result<Vec<Val>>>()?,192		)?))193	} else {194		Ok(ArrValue::lazy(uniq_keyf(arr, keyF)?))195	}196}197198#[builtin]199#[allow(non_snake_case)]200pub fn builtin_set(201	arr: ArrValue,202203	#[default(FuncVal::identity())] keyF: FuncVal,204) -> Result<ArrValue> {205	if arr.len() <= 1 {206		return Ok(arr);207	}208	if keyF.is_identity() {209		let arr = arr.iter().collect::<Result<Vec<Val>>>()?;210		let arr = sort_identity(arr)?;211		let arr = uniq_identity(arr)?;212		Ok(ArrValue::eager(arr))213	} else {214		let arr = sort_keyf(arr, keyF.clone())?;215		let arr = uniq_keyf(ArrValue::lazy(arr), keyF)?;216		Ok(ArrValue::lazy(arr))217	}218}219220fn eval_keyf(val: Val, key_f: &Option<FuncVal>) -> Result<Val> {221	if let Some(key_f) = key_f {222		key_f.evaluate_simple(&(val,), false)223	} else {224		Ok(val)225	}226}227228fn array_top1(arr: ArrValue, key_f: Option<FuncVal>, ordering: Ordering) -> Result<Val> {229	let mut iter = arr.iter();230	let mut min = iter.next().expect("not empty")?;231	let mut min_key = eval_keyf(min.clone(), &key_f)?;232	for item in iter {233		let cur = item?;234		let cur_key = eval_keyf(cur.clone(), &key_f)?;235		if evaluate_compare_op(&cur_key, &min_key, BinaryOpType::Lt)? == ordering {236			min = cur;237			min_key = cur_key;238		}239	}240	Ok(min)241}242243#[builtin]244pub fn builtin_min_array(245	arr: ArrValue,246	keyF: Option<FuncVal>,247	onEmpty: Option<Thunk<Val>>,248) -> Result<Val> {249	if arr.is_empty() {250		return eval_on_empty(onEmpty);251	}252	array_top1(arr, keyF, Ordering::Less)253}254#[builtin]255pub fn builtin_max_array(256	arr: ArrValue,257	keyF: Option<FuncVal>,258	onEmpty: Option<Thunk<Val>>,259) -> Result<Val> {260	if arr.is_empty() {261		return eval_on_empty(onEmpty);262	}263	array_top1(arr, keyF, Ordering::Greater)264}