1#![allow(non_snake_case)]23use std::cmp::Ordering;45use jrsonnet_evaluator::{6 Result, Thunk, Val, bail,7 function::builtin,8 operator::evaluate_compare_op,9 val::{ArrValue, equals},10};11use jrsonnet_ir::BinaryOpType;1213use crate::{eval_on_empty, keyf::KeyF};1415#[derive(Copy, Clone)]16enum SortKeyType {17 Number,18 String,19 Unspecialized,20 Unknown,21}2223fn get_sort_type<T>(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result<SortKeyType> {24 let mut sort_type = SortKeyType::Unknown;25 for i in values {26 let i = key_getter(i);27 match (i, sort_type) {28 (Val::Str(_), SortKeyType::Unknown) => sort_type = SortKeyType::String,29 (Val::Num(_), SortKeyType::Unknown) => sort_type = SortKeyType::Number,30 (Val::Str(_), SortKeyType::String) | (Val::Num(_), SortKeyType::Number) => {}31 (Val::Str(_) | Val::Num(_), _) => {32 bail!("sort elements should have the same types")33 }34 (_, _) => return Ok(SortKeyType::Unspecialized),35 }36 }37 Ok(sort_type)38}3940fn sort_identity(mut values: Vec<Val>) -> Result<Vec<Val>> {41 42 let sort_type = get_sort_type(&values, |k| k)?;43 match sort_type {44 SortKeyType::Number => values.sort_unstable_by_key(|v| match v {45 Val::Num(n) => *n,46 _ => unreachable!(),47 }),48 SortKeyType::String => values.sort_unstable_by_key(|v| match v {49 Val::Str(s) => s.clone(),50 _ => unreachable!(),51 }),52 SortKeyType::Unknown | SortKeyType::Unspecialized => {53 let mut err = None;54 55 56 values.sort_unstable_by(|a, b| match evaluate_compare_op(a, b, BinaryOpType::Lt) {57 Ok(ord) => ord,58 Err(e) if err.is_none() => {59 let _ = err.insert(e);60 Ordering::Equal61 }62 Err(_) => Ordering::Equal,63 });64 if let Some(err) = err {65 return Err(err);66 }67 }68 }69 Ok(values)70}7172fn sort_keyf(values: ArrValue, keyf: KeyF) -> Result<Vec<Thunk<Val>>> {73 74 let mut vk = Vec::with_capacity(values.len());75 for value in values.iter_lazy() {76 vk.push((value.clone(), keyf.eval(value)?));77 }78 let sort_type = get_sort_type(&vk, |v| &v.1)?;79 match sort_type {80 SortKeyType::Number => vk.sort_by_key(|v| match v.1 {81 Val::Num(n) => n,82 _ => unreachable!(),83 }),84 SortKeyType::String => vk.sort_by_key(|v| match &v.1 {85 Val::Str(s) => s.clone(),86 _ => unreachable!(),87 }),88 SortKeyType::Unknown | SortKeyType::Unspecialized => {89 let mut err = None;90 91 92 vk.sort_by(93 |(_a, ak), (_b, bk)| match evaluate_compare_op(ak, bk, BinaryOpType::Lt) {94 Ok(ord) => ord,95 Err(e) if err.is_none() => {96 let _ = err.insert(e);97 Ordering::Equal98 }99 Err(_) => Ordering::Equal,100 },101 );102 if let Some(err) = err {103 return Err(err);104 }105 }106 }107 Ok(vk.into_iter().map(|v| v.0).collect())108}109110111pub fn sort(values: ArrValue, key_getter: KeyF) -> Result<ArrValue> {112 if values.len() <= 1 {113 return Ok(values);114 }115 if key_getter.is_identity() {116 Ok(ArrValue::new(sort_identity(117 values.iter().collect::<Result<Vec<Val>>>()?,118 )?))119 } else {120 Ok(ArrValue::new(sort_keyf(values, key_getter)?))121 }122}123124#[builtin]125pub fn builtin_sort(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {126 super::sort::sort(arr, keyF)127}128129fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {130 let mut out = Vec::new();131 let mut last = arr[0].clone();132 out.push(last.clone());133 for next in arr.into_iter().skip(1) {134 if !equals(&last, &next)? {135 out.push(next.clone());136 }137 last = next;138 }139 Ok(out)140}141142fn uniq_keyf(arr: ArrValue, keyf: KeyF) -> Result<Vec<Thunk<Val>>> {143 let mut out = Vec::new();144 let last_value = arr.get_lazy(0).unwrap();145 let mut last_key = keyf.eval(last_value.clone())?;146 out.push(last_value);147148 for next in arr.iter_lazy().skip(1) {149 let next_key = keyf.eval(next.clone())?;150 if !equals(&last_key, &next_key)? {151 out.push(next.clone());152 }153 last_key = next_key;154 }155 Ok(out)156}157158#[builtin]159#[allow(non_snake_case)]160pub fn builtin_uniq(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {161 if arr.len() <= 1 {162 return Ok(arr);163 }164 if keyF.is_identity() {165 Ok(ArrValue::new(uniq_identity(166 arr.iter().collect::<Result<Vec<Val>>>()?,167 )?))168 } else {169 Ok(ArrValue::new(uniq_keyf(arr, keyF)?))170 }171}172173#[builtin]174#[allow(non_snake_case)]175pub fn builtin_set(arr: ArrValue, #[default] keyF: KeyF) -> Result<ArrValue> {176 if arr.len() <= 1 {177 return Ok(arr);178 }179 if keyF.is_identity() {180 let arr = arr.iter().collect::<Result<Vec<Val>>>()?;181 let arr = sort_identity(arr)?;182 let arr = uniq_identity(arr)?;183 Ok(ArrValue::new(arr))184 } else {185 let arr = sort_keyf(arr, keyF.clone())?;186 let arr = uniq_keyf(ArrValue::new(arr), keyF)?;187 Ok(ArrValue::new(arr))188 }189}190191fn array_top1(arr: ArrValue, keyf: KeyF, ordering: Ordering) -> Result<Val> {192 let mut iter = arr.iter();193 let mut min = iter.next().expect("not empty")?;194 let mut min_key = keyf.eval(Thunk::evaluated(min.clone()))?;195 for item in iter {196 let cur = item?;197 let cur_key = keyf.eval(Thunk::evaluated(cur.clone()))?;198 if evaluate_compare_op(&cur_key, &min_key, BinaryOpType::Lt)? == ordering {199 min = cur;200 min_key = cur_key;201 }202 }203 Ok(min)204}205206#[builtin]207pub fn builtin_min_array(208 arr: ArrValue,209 #[default] keyF: KeyF,210 onEmpty: Option<Thunk<Val>>,211) -> Result<Val> {212 if arr.is_empty() {213 return eval_on_empty(onEmpty);214 }215 array_top1(arr, keyF, Ordering::Less)216}217#[builtin]218pub fn builtin_max_array(219 arr: ArrValue,220 #[default] keyF: KeyF,221 onEmpty: Option<Thunk<Val>>,222) -> Result<Val> {223 if arr.is_empty() {224 return eval_on_empty(onEmpty);225 }226 array_top1(arr, keyF, Ordering::Greater)227}