1#![allow(non_snake_case)]23use std::cmp::Ordering;45use jrsonnet_evaluator::{6 bail,7 function::builtin,8 operator::evaluate_compare_op,9 val::{equals, ArrValue},10 Result, Thunk, Val,11};12use jrsonnet_parser::BinaryOpType;1314use crate::{eval_on_empty, keyf::KeyF};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: KeyF) -> Result<Vec<Thunk<Val>>> {74 75 let mut vk = Vec::with_capacity(values.len());76 for value in values.iter_lazy() {77 vk.push((value.clone(), keyf.eval(value)?));78 }79 let sort_type = get_sort_type(&vk, |v| &v.1)?;80 match sort_type {81 SortKeyType::Number => vk.sort_by_key(|v| match v.1 {82 Val::Num(n) => n,83 _ => unreachable!(),84 }),85 SortKeyType::String => vk.sort_by_key(|v| match &v.1 {86 Val::Str(s) => s.clone(),87 _ => unreachable!(),88 }),89 SortKeyType::Unknown | SortKeyType::Unspecialized => {90 let mut err = None;91 92 93 vk.sort_by(94 |(_a, ak), (_b, bk)| match evaluate_compare_op(ak, bk, BinaryOpType::Lt) {95 Ok(ord) => ord,96 Err(e) if err.is_none() => {97 let _ = err.insert(e);98 Ordering::Equal99 }100 Err(_) => Ordering::Equal,101 },102 );103 if let Some(err) = err {104 return Err(err);105 }106 }107 }108 Ok(vk.into_iter().map(|v| v.0).collect())109}110111112pub fn sort(values: ArrValue, key_getter: KeyF) -> Result<ArrValue> {113 if values.len() <= 1 {114 return Ok(values);115 }116 if key_getter.is_identity() {117 Ok(ArrValue::eager(sort_identity(118 values.iter().collect::<Result<Vec<Val>>>()?,119 )?))120 } else {121 Ok(ArrValue::lazy(sort_keyf(values, key_getter)?))122 }123}124125#[builtin]126pub fn builtin_sort(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {127 super::sort::sort(arr, keyF)128}129130fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {131 let mut out = Vec::new();132 let mut last = arr[0].clone();133 out.push(last.clone());134 for next in arr.into_iter().skip(1) {135 if !equals(&last, &next)? {136 out.push(next.clone());137 }138 last = next;139 }140 Ok(out)141}142143fn uniq_keyf(arr: ArrValue, keyf: KeyF) -> Result<Vec<Thunk<Val>>> {144 let mut out = Vec::new();145 let last_value = arr.get_lazy(0).unwrap();146 let mut last_key = keyf.eval(last_value.clone())?;147 out.push(last_value);148149 for next in arr.iter_lazy().skip(1) {150 let next_key = keyf.eval(next.clone())?;151 if !equals(&last_key, &next_key)? {152 out.push(next.clone());153 }154 last_key = next_key;155 }156 Ok(out)157}158159#[builtin]160#[allow(non_snake_case)]161pub fn builtin_uniq(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {162 if arr.len() <= 1 {163 return Ok(arr);164 }165 if keyF.is_identity() {166 Ok(ArrValue::eager(uniq_identity(167 arr.iter().collect::<Result<Vec<Val>>>()?,168 )?))169 } else {170 Ok(ArrValue::lazy(uniq_keyf(arr, keyF)?))171 }172}173174#[builtin]175#[allow(non_snake_case)]176pub fn builtin_set(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {177 if arr.len() <= 1 {178 return Ok(arr);179 }180 if keyF.is_identity() {181 let arr = arr.iter().collect::<Result<Vec<Val>>>()?;182 let arr = sort_identity(arr)?;183 let arr = uniq_identity(arr)?;184 Ok(ArrValue::eager(arr))185 } else {186 let arr = sort_keyf(arr, keyF.clone())?;187 let arr = uniq_keyf(ArrValue::lazy(arr), keyF)?;188 Ok(ArrValue::lazy(arr))189 }190}191192fn array_top1(arr: ArrValue, keyf: KeyF, ordering: Ordering) -> Result<Val> {193 let mut iter = arr.iter();194 let mut min = iter.next().expect("not empty")?;195 let mut min_key = keyf.eval(Thunk::evaluated(min.clone()))?;196 for item in iter {197 let cur = item?;198 let cur_key = keyf.eval(Thunk::evaluated(cur.clone()))?;199 if evaluate_compare_op(&cur_key, &min_key, BinaryOpType::Lt)? == ordering {200 min = cur;201 min_key = cur_key;202 }203 }204 Ok(min)205}206207#[builtin]208pub fn builtin_min_array(209 arr: ArrValue,210 #[default] keyF: KeyF,211 onEmpty: Option<Thunk<Val>>,212) -> Result<Val> {213 if arr.is_empty() {214 return eval_on_empty(onEmpty);215 }216 array_top1(arr, keyF, Ordering::Less)217}218#[builtin]219pub fn builtin_max_array(220 arr: ArrValue,221 #[default] keyF: KeyF,222 onEmpty: Option<Thunk<Val>>,223) -> Result<Val> {224 if arr.is_empty() {225 return eval_on_empty(onEmpty);226 }227 array_top1(arr, keyF, Ordering::Greater)228}