git.delta.rocks / jrsonnet / refs/commits / e88030decfb4

--- 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<T>(
-	values: &mut [T],
-	key_getter: impl Fn(&mut T) -> &mut Val,
-) -> Result<SortKeyType> {
+fn get_sort_type<T>(values: &[T], key_getter: impl Fn(&T) -> &Val) -> Result<SortKeyType> {
 	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<Val>) -> Result<Vec<Val>> {
+	// 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<Vec<Thunk<Val>>> {
+	// 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<Val>, key_getter: FuncVal) -> Result<Vec<Val>> {
+pub fn sort(values: ArrValue, key_getter: FuncVal) -> Result<ArrValue> {
 	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::<Result<Vec<Val>>>()?,
+		)?))
 	} 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<FuncVal>) -> Result<ArrValue> {
+pub fn builtin_sort(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {
 	if arr.len() <= 1 {
 		return Ok(arr);
 	}
-	Ok(ArrValue::eager(super::sort::sort(
-		ctx,
-		arr.iter().collect::<Result<Vec<_>>>()?,
+	Ok(super::sort::sort(
+		arr,
 		keyF.unwrap_or_else(FuncVal::identity),
-	)?))
+	)?)
+}
+
+fn uniq_identity(arr: Vec<Val>) -> Result<Vec<Val>> {
+	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<Vec<Thunk<Val>>> {
+	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<FuncVal>) -> Result<ArrValue> {
+	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::<Result<Vec<Val>>>()?,
+		)?))
+	} else {
+		Ok(ArrValue::lazy(Cc::new(uniq_keyf(arr, keyF)?)))
+	}
+}
+
+#[builtin]
+#[allow(non_snake_case)]
+pub fn builtin_set(arr: ArrValue, keyF: Option<FuncVal>) -> Result<ArrValue> {
+	let keyF = keyF.unwrap_or(FuncVal::identity());
+	if keyF.is_identity() {
+		let arr = arr.iter().collect::<Result<Vec<Val>>>()?;
+		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)))
+	}
 }

196196
197      aux(value),197      aux(value),
198198
199  uniq(arr, keyF=id)::
200    local f(a, b) =
201      if std.length(a) == 0 then
202        [b]
203      else if keyF(a[std.length(a) - 1]) == keyF(b) then
204        a
205      else
206        a + [b];
207    std.foldl(f, arr, []),
208
209  set(arr, keyF=id)::
210    std.uniq(std.sort(arr, keyF), keyF),
211
212  setUnion(a, b, keyF=id)::199  setUnion(a, b, keyF=id)::
213    // NOTE: order matters, values in `a` win200    // NOTE: order matters, values in `a` win
214    local aux(a, b, i, j, acc) =201    local aux(a, b, i, j, acc) =