--- a/crates/jrsonnet-stdlib/src/lib.rs +++ b/crates/jrsonnet-stdlib/src/lib.rs @@ -106,6 +106,8 @@ ("format", builtin_format::INST), // Sort ("sort", builtin_sort::INST), + ("uniq", builtin_uniq::INST), + ("set", builtin_set::INST), // Hash ("md5", builtin_md5::INST), #[cfg(feature = "exp-more-hashes")] --- a/crates/jrsonnet-stdlib/src/sort.rs +++ b/crates/jrsonnet-stdlib/src/sort.rs @@ -1,10 +1,11 @@ use jrsonnet_evaluator::{ error::Result, - function::{builtin, CallLocation, FuncVal}, + function::{builtin, FuncVal}, throw, - val::ArrValue, - Context, Val, + val::{equals, ArrValue}, + Thunk, Val, }; +use jrsonnet_gcmodule::Cc; #[derive(Copy, Clone)] enum SortKeyType { @@ -27,12 +28,9 @@ } } -fn get_sort_type( - values: &mut [T], - key_getter: impl Fn(&mut T) -> &mut Val, -) -> Result { +fn get_sort_type(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result { let mut sort_type = SortKeyType::Unknown; - for i in values.iter_mut() { + for i in values.iter() { let i = key_getter(i); match (i, sort_type) { (Val::Str(_), SortKeyType::Unknown) => sort_type = SortKeyType::String, @@ -47,65 +45,130 @@ Ok(sort_type) } +fn sort_identity(mut values: Vec) -> Result> { + // Fast path, identity key getter + let sort_type = get_sort_type(&values, |k| k)?; + match sort_type { + SortKeyType::Number => values.sort_unstable_by_key(|v| match v { + Val::Num(n) => NonNaNf64(*n), + _ => unreachable!(), + }), + SortKeyType::String => values.sort_unstable_by_key(|v| match v { + Val::Str(s) => s.clone(), + _ => unreachable!(), + }), + SortKeyType::Unknown => unreachable!(), + }; + Ok(values) +} + +fn sort_keyf(values: ArrValue, keyf: FuncVal) -> Result>> { + // Slow path, user provided key getter + let mut vk = Vec::with_capacity(values.len()); + for value in values.iter_lazy() { + vk.push(( + value.clone(), + keyf.evaluate_simple(&(value.clone(),), false)?, + )); + } + let sort_type = get_sort_type(&mut vk, |v| &v.1)?; + match sort_type { + SortKeyType::Number => vk.sort_by_key(|v| match v.1 { + Val::Num(n) => NonNaNf64(n), + _ => unreachable!(), + }), + SortKeyType::String => vk.sort_by_key(|v| match &v.1 { + Val::Str(s) => s.clone(), + _ => unreachable!(), + }), + SortKeyType::Unknown => unreachable!(), + }; + Ok(vk.into_iter().map(|v| v.0).collect()) +} + /// * `key_getter` - None, if identity sort required -pub fn sort(ctx: Context, mut values: Vec, key_getter: FuncVal) -> Result> { +pub fn sort(values: ArrValue, key_getter: FuncVal) -> Result { if values.len() <= 1 { return Ok(values); } if key_getter.is_identity() { - // Fast path, identity key getter - let sort_type = get_sort_type(&mut values, |k| k)?; - match sort_type { - SortKeyType::Number => values.sort_unstable_by_key(|v| match v { - Val::Num(n) => NonNaNf64(*n), - _ => unreachable!(), - }), - SortKeyType::String => values.sort_unstable_by_key(|v| match v { - Val::Str(s) => s.clone(), - _ => unreachable!(), - }), - SortKeyType::Unknown => unreachable!(), - }; - Ok(values) + Ok(ArrValue::eager(sort_identity( + values.iter().collect::>>()?, + )?)) } else { - // Slow path, user provided key getter - let mut vk = Vec::with_capacity(values.len()); - for value in values.iter() { - vk.push(( - value.clone(), - key_getter.evaluate( - ctx.clone(), - CallLocation::native(), - &(value.clone(),), - true, - )?, - )); - } - let sort_type = get_sort_type(&mut vk, |v| &mut v.1)?; - match sort_type { - SortKeyType::Number => vk.sort_by_key(|v| match v.1 { - Val::Num(n) => NonNaNf64(n), - _ => unreachable!(), - }), - SortKeyType::String => vk.sort_by_key(|v| match &v.1 { - Val::Str(s) => s.clone(), - _ => unreachable!(), - }), - SortKeyType::Unknown => unreachable!(), - }; - Ok(vk.into_iter().map(|v| v.0).collect()) + Ok(ArrValue::lazy(Cc::new(sort_keyf(values, key_getter)?))) } } #[builtin] #[allow(non_snake_case)] -pub fn builtin_sort(ctx: Context, arr: ArrValue, keyF: Option) -> Result { +pub fn builtin_sort(arr: ArrValue, keyF: Option) -> Result { if arr.len() <= 1 { return Ok(arr); } - Ok(ArrValue::eager(super::sort::sort( - ctx, - arr.iter().collect::>>()?, + Ok(super::sort::sort( + arr, keyF.unwrap_or_else(FuncVal::identity), - )?)) + )?) +} + +fn uniq_identity(arr: Vec) -> Result> { + let mut out = Vec::new(); + let mut last = arr[0].clone(); + out.push(last.clone()); + for next in arr.into_iter().skip(1) { + if !equals(&last, &next)? { + out.push(next.clone()); + } + last = next; + } + Ok(out) +} + +fn uniq_keyf(arr: ArrValue, keyf: FuncVal) -> Result>> { + let mut out = Vec::new(); + let last_value = arr.get_lazy(0).unwrap(); + let mut last_key = keyf.evaluate_simple(&(last_value.clone(),), false)?; + out.push(last_value.clone()); + + for next in arr.iter_lazy().skip(1) { + let next_key = keyf.evaluate_simple(&(next.clone(),), false)?; + if !equals(&last_key, &next_key)? { + out.push(next.clone()); + } + last_key = next_key; + } + Ok(out) +} + +#[builtin] +#[allow(non_snake_case)] +pub fn builtin_uniq(arr: ArrValue, keyF: Option) -> Result { + if arr.len() <= 1 { + return Ok(arr); + } + let keyF = keyF.unwrap_or(FuncVal::identity()); + if keyF.is_identity() { + Ok(ArrValue::eager(uniq_identity( + arr.iter().collect::>>()?, + )?)) + } else { + Ok(ArrValue::lazy(Cc::new(uniq_keyf(arr, keyF)?))) + } +} + +#[builtin] +#[allow(non_snake_case)] +pub fn builtin_set(arr: ArrValue, keyF: Option) -> Result { + let keyF = keyF.unwrap_or(FuncVal::identity()); + if keyF.is_identity() { + let arr = arr.iter().collect::>>()?; + let arr = sort_identity(arr)?; + let arr = uniq_identity(arr)?; + Ok(ArrValue::eager(arr)) + } else { + let arr = sort_keyf(arr, keyF.clone())?; + let arr = uniq_keyf(ArrValue::lazy(Cc::new(arr)), keyF)?; + Ok(ArrValue::lazy(Cc::new(arr))) + } } --- a/crates/jrsonnet-stdlib/src/std.jsonnet +++ b/crates/jrsonnet-stdlib/src/std.jsonnet @@ -196,19 +196,6 @@ aux(value), - uniq(arr, keyF=id):: - local f(a, b) = - if std.length(a) == 0 then - [b] - else if keyF(a[std.length(a) - 1]) == keyF(b) then - a - else - a + [b]; - std.foldl(f, arr, []), - - set(arr, keyF=id):: - std.uniq(std.sort(arr, keyF), keyF), - setUnion(a, b, keyF=id):: // NOTE: order matters, values in `a` win local aux(a, b, i, j, acc) =