difftreelog
feat apperently, std.sort now supports arrays
in: master
1 file changed
crates/jrsonnet-stdlib/src/sort.rsdiffbeforeafterboth1use 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!("list is not empty, as checked in sort"),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(&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!("list is not empty, as checked in sort"),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 super::sort::sort(arr, keyF.unwrap_or_else(FuncVal::identity))107}108109fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {110 let mut out = Vec::new();111 let mut last = arr[0].clone();112 out.push(last.clone());113 for next in arr.into_iter().skip(1) {114 if !equals(&last, &next)? {115 out.push(next.clone());116 }117 last = next;118 }119 Ok(out)120}121122fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {123 let mut out = Vec::new();124 let last_value = arr.get_lazy(0).unwrap();125 let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?;126 out.push(last_value);127128 for next in arr.iter_lazy().skip(1) {129 let next_key = keyf.evaluate_simple(&(next.clone(),), false)?;130 if !equals(&last_key, &next_key)? {131 out.push(next.clone());132 }133 last_key = next_key;134 }135 Ok(out)136}137138#[builtin]139#[allow(non_snake_case)]140pub fn builtin_uniq(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {141 if arr.len() <= 1 {142 return Ok(arr);143 }144 let keyF = keyF.unwrap_or(FuncVal::identity());145 if keyF.is_identity() {146 Ok(ArrValue::eager(uniq_identity(147 arr.iter().collect::<Result<Vec<Val>>>()?,148 )?))149 } else {150 Ok(ArrValue::lazy(Cc::new(uniq_keyf(arr, keyF)?)))151 }152}153154#[builtin]155#[allow(non_snake_case)]156pub fn builtin_set(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {157 if arr.len() <= 1 {158 return Ok(arr);159 }160 let keyF = keyF.unwrap_or(FuncVal::identity());161 if keyF.is_identity() {162 let arr = arr.iter().collect::<Result<Vec<Val>>>()?;163 let arr = sort_identity(arr)?;164 let arr = uniq_identity(arr)?;165 Ok(ArrValue::eager(arr))166 } else {167 let arr = sort_keyf(arr, keyF.clone())?;168 let arr = uniq_keyf(ArrValue::lazy(Cc::new(arr)), keyF)?;169 Ok(ArrValue::lazy(Cc::new(arr)))170 }171}