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 50 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 67 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}888990pub 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}