git.delta.rocks / jrsonnet / refs/commits / 6fbe7c61b4ca

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, error,5	function::{NativeFn, builtin},6	typed::{BoundedUsize, Either2, FromUntyped},7	val::{ArrValue, IndexableVal, equals},8};910pub fn eval_on_empty(on_empty: Option<Thunk<Val>>) -> Result<Val> {11	if let Some(on_empty) = on_empty {12		on_empty.evaluate()13	} else {14		bail!("expected non-empty array")15	}16}1718#[builtin]19pub fn builtin_make_array(sz: u32, func: NativeFn!((u32,) -> Val)) -> Result<ArrValue> {20	if sz == 0 {21		return Ok(ArrValue::empty());22	}23	// Try eager evaluation: call func(i) immediately for each element.24	'eager: {25		let mut out = Vec::with_capacity(sz as usize);26		for i in 0..sz {27			match func.call(i) {28				Ok(v) => out.push(v),29				Err(_) => break 'eager,30			}31		}32		return Ok(ArrValue::new(out));33	}34	Ok(ArrValue::make(sz, func))35}3637#[builtin]38pub fn builtin_repeat(what: Either![IStr, ArrValue], count: u32) -> Result<Val> {39	Ok(match what {40		Either2::A(s) => Val::string(s.repeat(count as usize)),41		Either2::B(arr) => Val::Arr(42			ArrValue::repeated(arr, count).ok_or_else(|| error!("repeated length overflow"))?,43		),44	})45}4647#[builtin]48pub fn builtin_slice(49	indexable: IndexableVal,50	index: Option<Option<i32>>,51	end: Option<Option<i32>>,52	step: Option<Option<BoundedUsize<1, { i32::MAX as usize }>>>,53) -> Result<Val> {54	indexable55		.slice(index.flatten(), end.flatten(), step.flatten())56		.map(Val::from)57}5859#[builtin]60pub fn builtin_map(func: NativeFn!((Val) -> Val), arr: IndexableVal) -> ArrValue {61	let arr = arr.to_array();62	arr.map(func)63}6465#[builtin]66pub fn builtin_map_with_index(func: NativeFn!((u32, Val) -> Val), arr: IndexableVal) -> ArrValue {67	let arr = arr.to_array();68	arr.map_with_index(func)69}7071#[builtin]72pub fn builtin_map_with_key(73	func: NativeFn!((IStr, Val) -> Val),74	obj: ObjValue,75) -> Result<ObjValue> {76	let mut out = ObjValueBuilder::new();77	for (k, v) in obj.iter(78		// Makes sense mapped object should be ordered the same way, should not break anything when the output is not ordered (the default).79		// The thrown error might be different, but jsonnet80		// does not specify the evaluation order.81		#[cfg(feature = "exp-preserve-order")]82		true,83	) {84		let v = v?;85		out.field(k.clone()).value(func.call(k, v)?);86	}87	Ok(out.build())88}8990#[builtin]91pub fn builtin_flatmap(92	func: NativeFn!((Either![String, Val]) -> Val),93	arr: IndexableVal,94) -> Result<IndexableVal> {95	use std::fmt::Write;96	match arr {97		IndexableVal::Str(str) => {98			let mut out = String::new();99			for c in str.chars() {100				match func.call(Either2::A(c.to_string()))? {101					Val::Str(o) => write!(out, "{o}").unwrap(),102					Val::Null => {}103					_ => bail!("in std.join all items should be strings"),104				}105			}106			Ok(IndexableVal::Str(out.into()))107		}108		IndexableVal::Arr(a) => {109			let mut out = Vec::new();110			for el in a.iter() {111				let el = el?;112				match func.call(Either2::B(el))? {113					Val::Arr(o) => {114						for oe in o.iter() {115							out.push(oe?);116						}117					}118					Val::Null => {}119					_ => bail!("in std.join all items should be arrays"),120				}121			}122			Ok(IndexableVal::Arr(out.into()))123		}124	}125}126127type FilterFunc = NativeFn!((Thunk<Val>) -> bool);128129#[builtin]130pub fn builtin_filter(func: FilterFunc, arr: ArrValue) -> Result<ArrValue> {131	arr.filter(func)132}133134#[builtin]135pub fn builtin_filter_map(136	filter_func: FilterFunc,137	map_func: NativeFn!((Val) -> Val),138	arr: ArrValue,139) -> Result<ArrValue> {140	Ok(arr.filter(filter_func)?.map(map_func))141}142143#[builtin]144pub fn builtin_foldl(145	func: NativeFn!((Val, Either![Val, char]) -> Val),146	arr: Either![ArrValue, IStr],147	init: Val,148) -> Result<Val> {149	let mut acc = init;150	match arr {151		Either2::A(arr) => {152			for i in arr.iter() {153				acc = func.call(acc, Either2::A(i?))?;154			}155		}156		Either2::B(arr) => {157			for c in arr.chars() {158				acc = func.call(acc, Either2::B(c))?;159			}160		}161	}162	Ok(acc)163}164165#[builtin]166pub fn builtin_foldr(167	func: NativeFn!((Either![Val, char], Val) -> Val),168	arr: Either![ArrValue, IStr],169	init: Val,170) -> Result<Val> {171	let mut acc = init;172	match arr {173		Either2::A(arr) => {174			for i in arr.iter().rev() {175				acc = func.call(Either2::A(i?), acc)?;176			}177		}178		Either2::B(arr) => {179			for c in arr.chars().rev() {180				acc = func.call(Either2::B(c), acc)?;181			}182		}183	}184	Ok(acc)185}186187#[builtin]188pub fn builtin_range(from: i32, to: i32) -> Result<ArrValue> {189	if to < from {190		return Ok(ArrValue::empty());191	}192	Ok(ArrValue::range_inclusive(from, to))193}194195#[builtin]196pub fn builtin_join(sep: IndexableVal, arr: ArrValue) -> Result<IndexableVal> {197	use std::fmt::Write;198	Ok(match sep {199		IndexableVal::Arr(joiner_items) => {200			let mut out = Vec::new();201202			let mut first = true;203			for item in arr.iter() {204				let item = item?.clone();205				if let Val::Arr(items) = item {206					if !first {207						out.reserve(joiner_items.len() as usize);208						// TODO: extend209						for item in joiner_items.iter() {210							out.push(item?);211						}212					}213					first = false;214					out.reserve(items.len() as usize);215					for item in items.iter() {216						out.push(item?);217					}218				} else if matches!(item, Val::Null) {219				} else {220					bail!("in std.join all items should be arrays");221				}222			}223224			IndexableVal::Arr(out.into())225		}226		IndexableVal::Str(sep) => {227			let mut out = String::new();228229			let mut first = true;230			for item in arr.iter() {231				let item = item?.clone();232				if let Val::Str(item) = item {233					if !first {234						out += &sep;235					}236					first = false;237					write!(out, "{item}").unwrap();238				} else if matches!(item, Val::Null) {239				} else {240					bail!("in std.join all items should be strings");241				}242			}243244			IndexableVal::Str(out.into())245		}246	})247}248249#[builtin]250pub fn builtin_lines(arr: ArrValue) -> Result<IndexableVal> {251	builtin_join(252		IndexableVal::Str("\n".into()),253		ArrValue::extended(arr, ArrValue::new(vec![Val::string("")]))254			.ok_or_else(|| error!("array is too large"))?,255	)256}257258#[builtin]259pub fn builtin_resolve_path(f: String, r: String) -> String {260	let Some(pos) = f.rfind('/') else {261		return r;262	};263	format!("{}{}", &f[..=pos], r)264}265266pub fn deep_join_inner(out: &mut String, arr: IndexableVal) -> Result<()> {267	use std::fmt::Write;268	match arr {269		IndexableVal::Str(s) => write!(out, "{s}").expect("no error"),270		IndexableVal::Arr(arr) => {271			for ele in arr.iter() {272				let indexable = IndexableVal::from_untyped(ele?)?;273				deep_join_inner(out, indexable)?;274			}275		}276	}277	Ok(())278}279280#[builtin]281pub fn builtin_deep_join(arr: IndexableVal) -> Result<String> {282	let mut out = String::new();283	deep_join_inner(&mut out, arr)?;284	Ok(out)285}286287#[builtin]288pub fn builtin_reverse(arr: ArrValue) -> ArrValue {289	arr.reversed()290}291292#[builtin]293pub fn builtin_any(arr: ArrValue) -> Result<bool> {294	for v in arr.iter() {295		let v = bool::from_untyped(v?)?;296		if v {297			return Ok(true);298		}299	}300	Ok(false)301}302303#[builtin]304pub fn builtin_all(arr: ArrValue) -> Result<bool> {305	for v in arr.iter() {306		let v = bool::from_untyped(v?)?;307		if !v {308			return Ok(false);309		}310	}311	Ok(true)312}313314#[builtin]315pub fn builtin_member(arr: IndexableVal, x: Val) -> Result<bool> {316	match arr {317		IndexableVal::Str(str) => {318			let x: IStr = IStr::from_untyped(x)?;319			Ok(!x.is_empty() && str.contains(&*x))320		}321		IndexableVal::Arr(a) => {322			for item in a.iter() {323				let item = item?;324				if equals(&item, &x)? {325					return Ok(true);326				}327			}328			Ok(false)329		}330	}331}332333#[builtin]334pub fn builtin_find(value: Val, arr: ArrValue) -> Result<Vec<usize>> {335	let mut out = Vec::new();336	for (i, ele) in arr.iter().enumerate() {337		let ele = ele?;338		if equals(&ele, &value)? {339			out.push(i);340		}341	}342	Ok(out)343}344345#[builtin]346pub fn builtin_contains(arr: IndexableVal, elem: Val) -> Result<bool> {347	builtin_member(arr, elem)348}349350#[builtin]351pub fn builtin_count(arr: ArrValue, x: Val) -> Result<usize> {352	let mut count = 0;353	for item in arr.iter() {354		if equals(&item?, &x)? {355			count += 1;356		}357	}358	Ok(count)359}360361#[builtin]362pub fn builtin_avg(arr: Vec<f64>, onEmpty: Option<Thunk<Val>>) -> Result<Val> {363	if arr.is_empty() {364		return eval_on_empty(onEmpty);365	}366	#[expect(367		clippy::cast_precision_loss,368		reason = "array sizes are bounded to i32 len"369	)]370	Ok(Val::try_num(arr.iter().sum::<f64>() / (arr.len() as f64))?)371}372373#[builtin]374pub fn builtin_remove_at(arr: ArrValue, at: i32) -> Result<ArrValue> {375	let newArrLeft = arr.clone().slice(None, Some(at), None);376	let newArrRight = arr.slice(Some(at + 1), None, None);377378	Ok(ArrValue::extended(newArrLeft, newArrRight).ok_or_else(|| error!("array is too large"))?)379}380381#[builtin]382pub fn builtin_remove(arr: ArrValue, elem: Val) -> Result<ArrValue> {383	for (index, item) in arr.iter().enumerate() {384		if equals(&item?, &elem)? {385			#[expect(386				clippy::cast_possible_truncation,387				clippy::cast_possible_wrap,388				reason = "array sizes are bounded to i32 len"389			)]390			return builtin_remove_at(arr.clone(), index as i32);391		}392	}393	Ok(arr)394}395396#[builtin]397pub fn builtin_flatten_arrays(arrs: Vec<ArrValue>) -> Result<ArrValue> {398	pub fn flatten_inner(values: &[ArrValue]) -> Result<ArrValue> {399		if values.len() == 1 {400			return Ok(values[0].clone());401		} else if values.len() == 2 {402			return ArrValue::extended(values[0].clone(), values[1].clone())403				.ok_or_else(|| error!("array is too large"));404		}405		let (a, b) = values.split_at(values.len() / 2);406		ArrValue::extended(flatten_inner(a)?, flatten_inner(b)?)407			.ok_or_else(|| error!("array is too large"))408	}409	if arrs.is_empty() {410		return Ok(ArrValue::empty());411	} else if arrs.len() == 1 {412		return Ok(arrs.into_iter().next().expect("single"));413	}414	flatten_inner(&arrs)415}416417#[builtin]418pub fn builtin_flatten_deep_array(value: Val) -> Result<Vec<Val>> {419	fn process(value: Val, out: &mut Vec<Val>) -> Result<()> {420		match value {421			Val::Arr(arr) => {422				for ele in arr.iter() {423					process(ele?, out)?;424				}425			}426			_ => out.push(value),427		}428		Ok(())429	}430	let mut out = Vec::new();431	process(value, &mut out)?;432	Ok(out)433}434435#[builtin]436pub fn builtin_prune(437	a: Val,438439	#[default(false)]440	#[cfg(feature = "exp-preserve-order")]441	preserve_order: bool,442) -> Result<Val> {443	fn is_content(val: &Val) -> bool {444		match val {445			Val::Null => false,446			Val::Arr(a) => !a.is_empty(),447			Val::Obj(o) => !o.is_empty(),448			_ => true,449		}450	}451	Ok(match a {452		Val::Arr(a) => {453			let mut out = Vec::new();454			for (i, ele) in a.iter().enumerate() {455				let ele = ele456					.and_then(|v| {457						builtin_prune(458							v,459							#[cfg(feature = "exp-preserve-order")]460							preserve_order,461						)462					})463					.with_description(|| format!("elem <{i}> pruning"))?;464				if is_content(&ele) {465					out.push(ele);466				}467			}468			Val::arr(out)469		}470		Val::Obj(o) => {471			let mut out = ObjValueBuilder::new();472			for (name, value) in o.iter(473				#[cfg(feature = "exp-preserve-order")]474				preserve_order,475			) {476				let value = value477					.and_then(|v| {478						builtin_prune(479							v,480							#[cfg(feature = "exp-preserve-order")]481							preserve_order,482						)483					})484					.with_description(|| format!("field <{name}> pruning"))?;485				if !is_content(&value) {486					continue;487				}488				out.field(name).value(value);489			}490			Val::Obj(out.build())491		}492		_ => a,493	})494}