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 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 _ => {}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 => {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: FuncVal) -> Result<Vec<Thunk<Val>>> {73 74 let mut vk = Vec::with_capacity(values.len());75 for value in values.iter_lazy() {76 vk.push((77 value.clone(),78 keyf.evaluate_simple(&(value.clone(),), false)?,79 ));80 }81 let sort_type = get_sort_type(&vk, |v| &v.1)?;82 match sort_type {83 SortKeyType::Number => vk.sort_by_key(|v| match v.1 {84 Val::Num(n) => n,85 _ => unreachable!(),86 }),87 SortKeyType::String => vk.sort_by_key(|v| match &v.1 {88 Val::Str(s) => s.clone(),89 _ => unreachable!(),90 }),91 SortKeyType::Unknown => {92 let mut err = None;93 94 95 vk.sort_by(96 |(_a, ak), (_b, bk)| match evaluate_compare_op(ak, bk, BinaryOpType::Lt) {97 Ok(ord) => ord,98 Err(e) if err.is_none() => {99 let _ = err.insert(e);100 Ordering::Equal101 }102 Err(_) => Ordering::Equal,103 },104 );105 if let Some(err) = err {106 return Err(err);107 }108 }109 };110 Ok(vk.into_iter().map(|v| v.0).collect())111}112113114pub fn sort(values: ArrValue, key_getter: FuncVal) -> Result<ArrValue> {115 if values.len() <= 1 {116 return Ok(values);117 }118 if key_getter.is_identity() {119 Ok(ArrValue::eager(sort_identity(120 values.iter().collect::<Result<Vec<Val>>>()?,121 )?))122 } else {123 Ok(ArrValue::lazy(sort_keyf(values, key_getter)?))124 }125}126127#[builtin]128pub fn builtin_sort(129 arr: ArrValue,130131 #[default(FuncVal::identity())] keyF: FuncVal,132) -> Result<ArrValue> {133 super::sort::sort(arr, keyF)134}135136fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {137 let mut out = Vec::new();138 let mut last = arr[0].clone();139 out.push(last.clone());140 for next in arr.into_iter().skip(1) {141 if !equals(&last, &next)? {142 out.push(next.clone());143 }144 last = next;145 }146 Ok(out)147}148149fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result<Vec<Thunk<Val>>> {150 let mut out = Vec::new();151 let last_value = arr.get_lazy(0).unwrap();152 let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?;153 out.push(last_value);154155 for next in arr.iter_lazy().skip(1) {156 let next_key = keyf.evaluate_simple(&(next.clone(),), false)?;157 if !equals(&last_key, &next_key)? {158 out.push(next.clone());159 }160 last_key = next_key;161 }162 Ok(out)163}164165#[builtin]166#[allow(non_snake_case)]167pub fn builtin_uniq(168 arr: ArrValue,169170 #[default(FuncVal::identity())] keyF: FuncVal,171) -> Result<ArrValue> {172 if arr.len() <= 1 {173 return Ok(arr);174 }175 if keyF.is_identity() {176 Ok(ArrValue::eager(uniq_identity(177 arr.iter().collect::<Result<Vec<Val>>>()?,178 )?))179 } else {180 Ok(ArrValue::lazy(uniq_keyf(arr, keyF)?))181 }182}183184#[builtin]185#[allow(non_snake_case)]186pub fn builtin_set(187 arr: ArrValue,188189 #[default(FuncVal::identity())] keyF: FuncVal,190) -> Result<ArrValue> {191 if arr.len() <= 1 {192 return Ok(arr);193 }194 if keyF.is_identity() {195 let arr = arr.iter().collect::<Result<Vec<Val>>>()?;196 let arr = sort_identity(arr)?;197 let arr = uniq_identity(arr)?;198 Ok(ArrValue::eager(arr))199 } else {200 let arr = sort_keyf(arr, keyF.clone())?;201 let arr = uniq_keyf(ArrValue::lazy(arr), keyF)?;202 Ok(ArrValue::lazy(arr))203 }204}205206fn eval_keyf(val: Val, key_f: &Option<FuncVal>) -> Result<Val> {207 if let Some(key_f) = key_f {208 key_f.evaluate_simple(&(val,), false)209 } else {210 Ok(val)211 }212}213214fn array_top1(arr: ArrValue, key_f: Option<FuncVal>, ordering: Ordering) -> Result<Val> {215 let mut iter = arr.iter();216 let mut min = iter.next().expect("not empty")?;217 let mut min_key = eval_keyf(min.clone(), &key_f)?;218 for item in iter {219 let cur = item?;220 let cur_key = eval_keyf(cur.clone(), &key_f)?;221 if evaluate_compare_op(&cur_key, &min_key, BinaryOpType::Lt)? == ordering {222 min = cur;223 min_key = cur_key;224 }225 }226 Ok(min)227}228229#[builtin]230pub fn builtin_min_array(231 arr: ArrValue,232 keyF: Option<FuncVal>,233 onEmpty: Option<Thunk<Val>>,234) -> Result<Val> {235 if arr.is_empty() {236 return eval_on_empty(onEmpty);237 }238 array_top1(arr, keyF, Ordering::Less)239}240#[builtin]241pub fn builtin_max_array(242 arr: ArrValue,243 keyF: Option<FuncVal>,244 onEmpty: Option<Thunk<Val>>,245) -> Result<Val> {246 if arr.is_empty() {247 return eval_on_empty(onEmpty);248 }249 array_top1(arr, keyF, Ordering::Greater)250}