git.delta.rocks / jrsonnet / refs/commits / 48e38361db94

difftreelog

source

crates/jrsonnet-stdlib/src/arrays.rs10.9 KiBsourcehistory
1#![allow(non_snake_case)]23use jrsonnet_evaluator::{4	Either, IStr, ObjValue, ObjValueBuilder, Result, ResultExt, Thunk, Val, bail,5	function::{FuncVal, NativeFn, builtin},6	runtime_error,7	typed::{BoundedI32, BoundedUsize, Either2, FromUntyped},8	val::{ArrValue, IndexableVal, equals},9};1011pub fn eval_on_empty(on_empty: Option<Thunk<Val>>) -> Result<Val> {12	if let Some(on_empty) = on_empty {13		on_empty.evaluate()14	} else {15		bail!("expected non-empty array")16	}17}1819#[builtin]20pub fn builtin_make_array(21	// Can't use usize because range_exclusive is over i3222	sz: BoundedI32<0, { i32::MAX }>,23	func: FuncVal,24) -> Result<ArrValue> {25	if *sz == 0 {26		return Ok(ArrValue::empty());27	}28	func.evaluate_trivial().map_or_else(29		// TODO: Different mapped array impl avoiding allocating unnecessary vals30		|| Ok(ArrValue::range_exclusive(0, *sz).map(FromUntyped::from_untyped(Val::Func(func))?)),31		|trivial| {32			#[expect(clippy::cast_sign_loss, reason = "sz is bounded to be larger than 0")]33			let mut out = Vec::with_capacity(*sz as usize);34			for _ in 0..*sz {35				out.push(trivial.clone());36			}37			Ok(ArrValue::new(out))38		},39	)40}4142#[builtin]43pub fn builtin_repeat(what: Either![IStr, ArrValue], count: usize) -> Result<Val> {44	Ok(match what {45		Either2::A(s) => Val::string(s.repeat(count)),46		Either2::B(arr) => Val::Arr(47			ArrValue::repeated(arr, count)48				.ok_or_else(|| runtime_error!("repeated length overflow"))?,49		),50	})51}5253#[builtin]54pub fn builtin_slice(55	indexable: IndexableVal,56	index: Option<Option<i32>>,57	end: Option<Option<i32>>,58	step: Option<Option<BoundedUsize<1, { i32::MAX as usize }>>>,59) -> Result<Val> {60	indexable61		.slice(index.flatten(), end.flatten(), step.flatten())62		.map(Val::from)63}6465#[builtin]66pub fn builtin_map(func: NativeFn!((Val) -> Val), arr: IndexableVal) -> ArrValue {67	let arr = arr.to_array();68	arr.map(func)69}7071#[builtin]72pub fn builtin_map_with_index(func: NativeFn!((u32, Val) -> Val), arr: IndexableVal) -> ArrValue {73	let arr = arr.to_array();74	arr.map_with_index(func)75}7677#[builtin]78pub fn builtin_map_with_key(79	func: NativeFn!((IStr, Val) -> Val),80	obj: ObjValue,81) -> Result<ObjValue> {82	let mut out = ObjValueBuilder::new();83	for (k, v) in obj.iter(84		// Makes sense mapped object should be ordered the same way, should not break anything when the output is not ordered (the default).85		// The thrown error might be different, but jsonnet86		// does not specify the evaluation order.87		#[cfg(feature = "exp-preserve-order")]88		true,89	) {90		let v = v?;91		out.field(k.clone()).value(func.call(k, v)?);92	}93	Ok(out.build())94}9596#[builtin]97pub fn builtin_flatmap(98	func: NativeFn!((Either![String, Val]) -> Val),99	arr: IndexableVal,100) -> Result<IndexableVal> {101	use std::fmt::Write;102	match arr {103		IndexableVal::Str(str) => {104			let mut out = String::new();105			for c in str.chars() {106				match func.call(Either2::A(c.to_string()))? {107					Val::Str(o) => write!(out, "{o}").unwrap(),108					Val::Null => {}109					_ => bail!("in std.join all items should be strings"),110				}111			}112			Ok(IndexableVal::Str(out.into()))113		}114		IndexableVal::Arr(a) => {115			let mut out = Vec::new();116			for el in a.iter() {117				let el = el?;118				match func.call(Either2::B(el))? {119					Val::Arr(o) => {120						for oe in o.iter() {121							out.push(oe?);122						}123					}124					Val::Null => {}125					_ => bail!("in std.join all items should be arrays"),126				}127			}128			Ok(IndexableVal::Arr(out.into()))129		}130	}131}132133type FilterFunc = NativeFn!((Thunk<Val>) -> bool);134135#[builtin]136pub fn builtin_filter(func: FilterFunc, arr: ArrValue) -> Result<ArrValue> {137	arr.filter(func)138}139140#[builtin]141pub fn builtin_filter_map(142	filter_func: FilterFunc,143	map_func: NativeFn!((Val) -> Val),144	arr: ArrValue,145) -> Result<ArrValue> {146	Ok(arr.filter(filter_func)?.map(map_func))147}148149#[builtin]150pub fn builtin_foldl(151	func: NativeFn!((Val, Either![Val, char]) -> Val),152	arr: Either![ArrValue, IStr],153	init: Val,154) -> Result<Val> {155	let mut acc = init;156	match arr {157		Either2::A(arr) => {158			for i in arr.iter() {159				acc = func.call(acc, Either2::A(i?))?;160			}161		}162		Either2::B(arr) => {163			for c in arr.chars() {164				acc = func.call(acc, Either2::B(c))?;165			}166		}167	}168	Ok(acc)169}170171#[builtin]172pub fn builtin_foldr(173	func: NativeFn!((Either![Val, char], Val) -> Val),174	arr: Either![ArrValue, IStr],175	init: Val,176) -> Result<Val> {177	let mut acc = init;178	match arr {179		Either2::A(arr) => {180			for i in arr.iter().rev() {181				acc = func.call(Either2::A(i?), acc)?;182			}183		}184		Either2::B(arr) => {185			for c in arr.chars().rev() {186				acc = func.call(Either2::B(c), acc)?;187			}188		}189	}190	Ok(acc)191}192193#[builtin]194pub fn builtin_range(from: i32, to: i32) -> Result<ArrValue> {195	if to < from {196		return Ok(ArrValue::empty());197	}198	Ok(ArrValue::range_inclusive(from, to))199}200201#[builtin]202pub fn builtin_join(sep: IndexableVal, arr: ArrValue) -> Result<IndexableVal> {203	use std::fmt::Write;204	Ok(match sep {205		IndexableVal::Arr(joiner_items) => {206			let mut out = Vec::new();207208			let mut first = true;209			for item in arr.iter() {210				let item = item?.clone();211				if let Val::Arr(items) = item {212					if !first {213						out.reserve(joiner_items.len());214						// TODO: extend215						for item in joiner_items.iter() {216							out.push(item?);217						}218					}219					first = false;220					out.reserve(items.len());221					for item in items.iter() {222						out.push(item?);223					}224				} else if matches!(item, Val::Null) {225				} else {226					bail!("in std.join all items should be arrays");227				}228			}229230			IndexableVal::Arr(out.into())231		}232		IndexableVal::Str(sep) => {233			let mut out = String::new();234235			let mut first = true;236			for item in arr.iter() {237				let item = item?.clone();238				if let Val::Str(item) = item {239					if !first {240						out += &sep;241					}242					first = false;243					write!(out, "{item}").unwrap();244				} else if matches!(item, Val::Null) {245				} else {246					bail!("in std.join all items should be strings");247				}248			}249250			IndexableVal::Str(out.into())251		}252	})253}254255#[builtin]256pub fn builtin_lines(arr: ArrValue) -> Result<IndexableVal> {257	builtin_join(258		IndexableVal::Str("\n".into()),259		ArrValue::extended(arr, ArrValue::new(vec![Val::string("")])),260	)261}262263#[builtin]264pub fn builtin_resolve_path(f: String, r: String) -> String {265	let Some(pos) = f.rfind('/') else {266		return r;267	};268	format!("{}{}", &f[..=pos], r)269}270271pub fn deep_join_inner(out: &mut String, arr: IndexableVal) -> Result<()> {272	use std::fmt::Write;273	match arr {274		IndexableVal::Str(s) => write!(out, "{s}").expect("no error"),275		IndexableVal::Arr(arr) => {276			for ele in arr.iter() {277				let indexable = IndexableVal::from_untyped(ele?)?;278				deep_join_inner(out, indexable)?;279			}280		}281	}282	Ok(())283}284285#[builtin]286pub fn builtin_deep_join(arr: IndexableVal) -> Result<String> {287	let mut out = String::new();288	deep_join_inner(&mut out, arr)?;289	Ok(out)290}291292#[builtin]293pub fn builtin_reverse(arr: ArrValue) -> ArrValue {294	arr.reversed()295}296297#[builtin]298pub fn builtin_any(arr: ArrValue) -> Result<bool> {299	for v in arr.iter() {300		let v = bool::from_untyped(v?)?;301		if v {302			return Ok(true);303		}304	}305	Ok(false)306}307308#[builtin]309pub fn builtin_all(arr: ArrValue) -> Result<bool> {310	for v in arr.iter() {311		let v = bool::from_untyped(v?)?;312		if !v {313			return Ok(false);314		}315	}316	Ok(true)317}318319#[builtin]320pub fn builtin_member(arr: IndexableVal, x: Val) -> Result<bool> {321	match arr {322		IndexableVal::Str(str) => {323			let x: IStr = IStr::from_untyped(x)?;324			Ok(!x.is_empty() && str.contains(&*x))325		}326		IndexableVal::Arr(a) => {327			for item in a.iter() {328				let item = item?;329				if equals(&item, &x)? {330					return Ok(true);331				}332			}333			Ok(false)334		}335	}336}337338#[builtin]339pub fn builtin_find(value: Val, arr: ArrValue) -> Result<Vec<usize>> {340	let mut out = Vec::new();341	for (i, ele) in arr.iter().enumerate() {342		let ele = ele?;343		if equals(&ele, &value)? {344			out.push(i);345		}346	}347	Ok(out)348}349350#[builtin]351pub fn builtin_contains(arr: IndexableVal, elem: Val) -> Result<bool> {352	builtin_member(arr, elem)353}354355#[builtin]356pub fn builtin_count(arr: ArrValue, x: Val) -> Result<usize> {357	let mut count = 0;358	for item in arr.iter() {359		if equals(&item?, &x)? {360			count += 1;361		}362	}363	Ok(count)364}365366#[builtin]367pub fn builtin_avg(arr: Vec<f64>, onEmpty: Option<Thunk<Val>>) -> Result<Val> {368	if arr.is_empty() {369		return eval_on_empty(onEmpty);370	}371	#[expect(372		clippy::cast_precision_loss,373		reason = "array sizes are bounded to i32 len"374	)]375	Ok(Val::try_num(arr.iter().sum::<f64>() / (arr.len() as f64))?)376}377378#[builtin]379pub fn builtin_remove_at(arr: ArrValue, at: i32) -> Result<ArrValue> {380	let newArrLeft = arr.clone().slice(None, Some(at), None);381	let newArrRight = arr.slice(Some(at + 1), None, None);382383	Ok(ArrValue::extended(newArrLeft, newArrRight))384}385386#[builtin]387pub fn builtin_remove(arr: ArrValue, elem: Val) -> Result<ArrValue> {388	for (index, item) in arr.iter().enumerate() {389		if equals(&item?, &elem)? {390			#[expect(391				clippy::cast_possible_truncation,392				clippy::cast_possible_wrap,393				reason = "array sizes are bounded to i32 len"394			)]395			return builtin_remove_at(arr.clone(), index as i32);396		}397	}398	Ok(arr)399}400401#[builtin]402pub fn builtin_flatten_arrays(arrs: Vec<ArrValue>) -> ArrValue {403	pub fn flatten_inner(values: &[ArrValue]) -> ArrValue {404		if values.len() == 1 {405			return values[0].clone();406		} else if values.len() == 2 {407			return ArrValue::extended(values[0].clone(), values[1].clone());408		}409		let (a, b) = values.split_at(values.len() / 2);410		ArrValue::extended(flatten_inner(a), flatten_inner(b))411	}412	if arrs.is_empty() {413		return ArrValue::empty();414	} else if arrs.len() == 1 {415		return arrs.into_iter().next().expect("single");416	}417	flatten_inner(&arrs)418}419420#[builtin]421pub fn builtin_flatten_deep_array(value: Val) -> Result<Vec<Val>> {422	fn process(value: Val, out: &mut Vec<Val>) -> Result<()> {423		match value {424			Val::Arr(arr) => {425				for ele in arr.iter() {426					process(ele?, out)?;427				}428			}429			_ => out.push(value),430		}431		Ok(())432	}433	let mut out = Vec::new();434	process(value, &mut out)?;435	Ok(out)436}437438#[builtin]439pub fn builtin_prune(440	a: Val,441442	#[default(false)]443	#[cfg(feature = "exp-preserve-order")]444	preserve_order: bool,445) -> Result<Val> {446	fn is_content(val: &Val) -> bool {447		match val {448			Val::Null => false,449			Val::Arr(a) => !a.is_empty(),450			Val::Obj(o) => !o.is_empty(),451			_ => true,452		}453	}454	Ok(match a {455		Val::Arr(a) => {456			let mut out = Vec::new();457			for (i, ele) in a.iter().enumerate() {458				let ele = ele459					.and_then(|v| {460						builtin_prune(461							v,462							#[cfg(feature = "exp-preserve-order")]463							preserve_order,464						)465					})466					.with_description(|| format!("elem <{i}> pruning"))?;467				if is_content(&ele) {468					out.push(ele);469				}470			}471			Val::arr(out)472		}473		Val::Obj(o) => {474			let mut out = ObjValueBuilder::new();475			for (name, value) in o.iter(476				#[cfg(feature = "exp-preserve-order")]477				preserve_order,478			) {479				let value = value480					.and_then(|v| {481						builtin_prune(482							v,483							#[cfg(feature = "exp-preserve-order")]484							preserve_order,485						)486					})487					.with_description(|| format!("field <{name}> pruning"))?;488				if !is_content(&value) {489					continue;490				}491				out.field(name).value(value);492			}493			Val::Obj(out.build())494		}495		_ => a,496	})497}