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 Unspecialized,21 Unknown,22}2324fn get_sort_type<T>(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result<SortKeyType> {25 let mut sort_type = SortKeyType::Unknown;26 for i in values {27 let i = key_getter(i);28 match (i, sort_type) {29 (Val::Str(_), SortKeyType::Unknown) => sort_type = SortKeyType::String,30 (Val::Num(_), SortKeyType::Unknown) => sort_type = SortKeyType::Number,31 (Val::Str(_), SortKeyType::String) | (Val::Num(_), SortKeyType::Number) => {}32 (Val::Str(_) | Val::Num(_), _) => {33 bail!("sort elements should have the same types")34 }35 (_, _) => return Ok(SortKeyType::Unspecialized),36 }37 }38 Ok(sort_type)39}4041fn sort_identity(mut values: Vec<Val>) -> Result<Vec<Val>> {42 43 let sort_type = get_sort_type(&values, |k| k)?;44 match sort_type {45 SortKeyType::Number => values.sort_unstable_by_key(|v| match v {46 Val::Num(n) => *n,47 _ => unreachable!(),48 }),49 SortKeyType::String => values.sort_unstable_by_key(|v| match v {50 Val::Str(s) => s.clone(),51 _ => unreachable!(),52 }),53 SortKeyType::Unknown | SortKeyType::Unspecialized => {54 let mut err = None;55 56 57 values.sort_unstable_by(|a, b| match evaluate_compare_op(a, b, BinaryOpType::Lt) {58 Ok(ord) => ord,59 Err(e) if err.is_none() => {60 let _ = err.insert(e);61 Ordering::Equal62 }63 Err(_) => Ordering::Equal,64 });65 if let Some(err) = err {66 return Err(err);67 }68 }69 }70 Ok(values)71}7273fn sort_keyf(values: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {74 75 let mut vk = Vec::with_capacity(values.len());76 for value in values.iter_lazy() {77 vk.push((78 value.clone(),79 keyf.evaluate_simple(&(value.clone(),), false)?,80 ));81 }82 let sort_type = get_sort_type(&vk, |v| &v.1)?;83 match sort_type {84 SortKeyType::Number => vk.sort_by_key(|v| match v.1 {85 Val::Num(n) => n,86 _ => unreachable!(),87 }),88 SortKeyType::String => vk.sort_by_key(|v| match &v.1 {89 Val::Str(s) => s.clone(),90 _ => unreachable!(),91 }),92 SortKeyType::Unknown | SortKeyType::Unspecialized => {93 let mut err = None;94 95 96 vk.sort_by(97 |(_a, ak), (_b, bk)| match evaluate_compare_op(ak, bk, BinaryOpType::Lt) {98 Ok(ord) => ord,99 Err(e) if err.is_none() => {100 let _ = err.insert(e);101 Ordering::Equal102 }103 Err(_) => Ordering::Equal,104 },105 );106 if let Some(err) = err {107 return Err(err);108 }109 }110 }111 Ok(vk.into_iter().map(|v| v.0).collect())112}113114115pub fn sort(values: ArrValue, key_getter: FuncVal) -> Result<ArrValue> {116 if values.len() <= 1 {117 return Ok(values);118 }119 if key_getter.is_identity() {120 Ok(ArrValue::eager(sort_identity(121 values.iter().collect::<Result<Vec<Val>>>()?,122 )?))123 } else {124 Ok(ArrValue::lazy(sort_keyf(values, key_getter)?))125 }126}127128#[builtin]129pub fn builtin_sort(130 arr: ArrValue,131132 #[default(FuncVal::identity())] keyF: FuncVal,133) -> Result<ArrValue> {134 super::sort::sort(arr, keyF)135}136137fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {138 let mut out = Vec::new();139 let mut last = arr[0].clone();140 out.push(last.clone());141 for next in arr.into_iter().skip(1) {142 if !equals(&last, &next)? {143 out.push(next.clone());144 }145 last = next;146 }147 Ok(out)148}149150fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {151 let mut out = Vec::new();152 let last_value = arr.get_lazy(0).unwrap();153 let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?;154 out.push(last_value);155156 for next in arr.iter_lazy().skip(1) {157 let next_key = keyf.evaluate_simple(&(next.clone(),), false)?;158 if !equals(&last_key, &next_key)? {159 out.push(next.clone());160 }161 last_key = next_key;162 }163 Ok(out)164}165166#[builtin]167#[allow(non_snake_case)]168pub fn builtin_uniq(169 arr: ArrValue,170171 #[default(FuncVal::identity())] keyF: FuncVal,172) -> Result<ArrValue> {173 if arr.len() <= 1 {174 return Ok(arr);175 }176 if keyF.is_identity() {177 Ok(ArrValue::eager(uniq_identity(178 arr.iter().collect::<Result<Vec<Val>>>()?,179 )?))180 } else {181 Ok(ArrValue::lazy(uniq_keyf(arr, keyF)?))182 }183}184185#[builtin]186#[allow(non_snake_case)]187pub fn builtin_set(188 arr: ArrValue,189190 #[default(FuncVal::identity())] keyF: FuncVal,191) -> Result<ArrValue> {192 if arr.len() <= 1 {193 return Ok(arr);194 }195 if keyF.is_identity() {196 let arr = arr.iter().collect::<Result<Vec<Val>>>()?;197 let arr = sort_identity(arr)?;198 let arr = uniq_identity(arr)?;199 Ok(ArrValue::eager(arr))200 } else {201 let arr = sort_keyf(arr, keyF.clone())?;202 let arr = uniq_keyf(ArrValue::lazy(arr), keyF)?;203 Ok(ArrValue::lazy(arr))204 }205}206207fn eval_keyf(val: Val, key_f: Option<&FuncVal>) -> Result<Val> {208 if let Some(key_f) = key_f {209 key_f.evaluate_simple(&(val,), false)210 } else {211 Ok(val)212 }213}214215fn array_top1(arr: ArrValue, key_f: Option<&FuncVal>, ordering: Ordering) -> Result<Val> {216 let mut iter = arr.iter();217 let mut min = iter.next().expect("not empty")?;218 let mut min_key = eval_keyf(min.clone(), key_f)?;219 for item in iter {220 let cur = item?;221 let cur_key = eval_keyf(cur.clone(), key_f)?;222 if evaluate_compare_op(&cur_key, &min_key, BinaryOpType::Lt)? == ordering {223 min = cur;224 min_key = cur_key;225 }226 }227 Ok(min)228}229230#[builtin]231pub fn builtin_min_array(232 arr: ArrValue,233 keyF: Option<FuncVal>,234 onEmpty: Option<Thunk<Val>>,235) -> Result<Val> {236 if arr.is_empty() {237 return eval_on_empty(onEmpty);238 }239 array_top1(arr, keyF.as_ref(), Ordering::Less)240}241#[builtin]242pub fn builtin_max_array(243 arr: ArrValue,244 keyF: Option<FuncVal>,245 onEmpty: Option<Thunk<Val>>,246) -> Result<Val> {247 if arr.is_empty() {248 return eval_on_empty(onEmpty);249 }250 array_top1(arr, keyF.as_ref(), Ordering::Greater)251}