git.delta.rocks / jrsonnet / refs/commits / 112adb2810f2

difftreelog

source

crates/jrsonnet-stdlib/src/arrays.rs10.9 KiBsourcehistory
1#![allow(non_snake_case)]23use jrsonnet_evaluator::{4	bail, error,5	function::{builtin, NativeFn},6	runtime_error,7	typed::{BoundedUsize, Either2, FromUntyped},8	val::{equals, ArrValue, IndexableVal},9	Either, IStr, ObjValue, ObjValueBuilder, Result, ResultExt, Thunk, Val,10};1112pub fn eval_on_empty(on_empty: Option<Thunk<Val>>) -> Result<Val> {13	if let Some(on_empty) = on_empty {14		on_empty.evaluate()15	} else {16		bail!("expected non-empty array")17	}18}1920#[builtin]21pub fn builtin_make_array(sz: u32, func: NativeFn!((u32,) -> Val)) -> Result<ArrValue> {22	if sz == 0 {23		return Ok(ArrValue::empty());24	}25	// Try eager evaluation: call func(i) immediately for each element.26	'eager: {27		let mut out = Vec::with_capacity(sz as usize);28		for i in 0..sz {29			match func.call(i) {30				Ok(v) => out.push(v),31				Err(_) => break 'eager,32			}33		}34		return Ok(ArrValue::new(out));35	}36	Ok(ArrValue::make(sz, func))37}3839#[builtin]40pub fn builtin_repeat(what: Either![IStr, ArrValue], count: u32) -> Result<Val> {41	Ok(match what {42		Either2::A(s) => Val::string(s.repeat(count as usize)),43		Either2::B(arr) => Val::Arr(44			ArrValue::repeated(arr, count)45				.ok_or_else(|| runtime_error!("repeated length overflow"))?,46		),47	})48}4950#[builtin]51pub fn builtin_slice(52	indexable: IndexableVal,53	index: Option<Option<i32>>,54	end: Option<Option<i32>>,55	step: Option<Option<BoundedUsize<1, { i32::MAX as usize }>>>,56) -> Result<Val> {57	indexable58		.slice(index.flatten(), end.flatten(), step.flatten())59		.map(Val::from)60}6162#[builtin]63pub fn builtin_map(func: NativeFn!((Val) -> Val), arr: IndexableVal) -> ArrValue {64	let arr = arr.to_array();65	arr.map(func)66}6768#[builtin]69pub fn builtin_map_with_index(func: NativeFn!((u32, Val) -> Val), arr: IndexableVal) -> ArrValue {70	let arr = arr.to_array();71	arr.map_with_index(func)72}7374#[builtin]75pub fn builtin_map_with_key(76	func: NativeFn!((IStr, Val) -> Val),77	obj: ObjValue,78) -> Result<ObjValue> {79	let mut out = ObjValueBuilder::new();80	for (k, v) in obj.iter(81		// Makes sense mapped object should be ordered the same way, should not break anything when the output is not ordered (the default).82		// The thrown error might be different, but jsonnet83		// does not specify the evaluation order.84		#[cfg(feature = "exp-preserve-order")]85		true,86	) {87		let v = v?;88		out.field(k.clone()).value(func.call(k, v)?);89	}90	Ok(out.build())91}9293#[builtin]94pub fn builtin_flatmap(95	func: NativeFn!((Either![String, Val]) -> Val),96	arr: IndexableVal,97) -> Result<IndexableVal> {98	use std::fmt::Write;99	match arr {100		IndexableVal::Str(str) => {101			let mut out = String::new();102			for c in str.chars() {103				match func.call(Either2::A(c.to_string()))? {104					Val::Str(o) => write!(out, "{o}").unwrap(),105					Val::Null => {}106					_ => bail!("in std.join all items should be strings"),107				}108			}109			Ok(IndexableVal::Str(out.into()))110		}111		IndexableVal::Arr(a) => {112			let mut out = Vec::new();113			for el in a.iter() {114				let el = el?;115				match func.call(Either2::B(el))? {116					Val::Arr(o) => {117						for oe in o.iter() {118							out.push(oe?);119						}120					}121					Val::Null => {}122					_ => bail!("in std.join all items should be arrays"),123				}124			}125			Ok(IndexableVal::Arr(out.into()))126		}127	}128}129130type FilterFunc = NativeFn!((Thunk<Val>) -> bool);131132#[builtin]133pub fn builtin_filter(func: FilterFunc, arr: ArrValue) -> Result<ArrValue> {134	arr.filter(func)135}136137#[builtin]138pub fn builtin_filter_map(139	filter_func: FilterFunc,140	map_func: NativeFn!((Val) -> Val),141	arr: ArrValue,142) -> Result<ArrValue> {143	Ok(arr.filter(filter_func)?.map(map_func))144}145146#[builtin]147pub fn builtin_foldl(148	func: NativeFn!((Val, Either![Val, char]) -> Val),149	arr: Either![ArrValue, IStr],150	init: Val,151) -> Result<Val> {152	let mut acc = init;153	match arr {154		Either2::A(arr) => {155			for i in arr.iter() {156				acc = func.call(acc, Either2::A(i?))?;157			}158		}159		Either2::B(arr) => {160			for c in arr.chars() {161				acc = func.call(acc, Either2::B(c))?;162			}163		}164	}165	Ok(acc)166}167168#[builtin]169pub fn builtin_foldr(170	func: NativeFn!((Either![Val, char], Val) -> Val),171	arr: Either![ArrValue, IStr],172	init: Val,173) -> Result<Val> {174	let mut acc = init;175	match arr {176		Either2::A(arr) => {177			for i in arr.iter().rev() {178				acc = func.call(Either2::A(i?), acc)?;179			}180		}181		Either2::B(arr) => {182			for c in arr.chars().rev() {183				acc = func.call(Either2::B(c), acc)?;184			}185		}186	}187	Ok(acc)188}189190#[builtin]191pub fn builtin_range(from: i32, to: i32) -> Result<ArrValue> {192	if to < from {193		return Ok(ArrValue::empty());194	}195	Ok(ArrValue::range_inclusive(from, to))196}197198#[builtin]199pub fn builtin_join(sep: IndexableVal, arr: ArrValue) -> Result<IndexableVal> {200	use std::fmt::Write;201	Ok(match sep {202		IndexableVal::Arr(joiner_items) => {203			let mut out = Vec::new();204205			let mut first = true;206			for item in arr.iter() {207				let item = item?.clone();208				if let Val::Arr(items) = item {209					if !first {210						out.reserve(joiner_items.len() as usize);211						// TODO: extend212						for item in joiner_items.iter() {213							out.push(item?);214						}215					}216					first = false;217					out.reserve(items.len() as usize);218					for item in items.iter() {219						out.push(item?);220					}221				} else if matches!(item, Val::Null) {222				} else {223					bail!("in std.join all items should be arrays");224				}225			}226227			IndexableVal::Arr(out.into())228		}229		IndexableVal::Str(sep) => {230			let mut out = String::new();231232			let mut first = true;233			for item in arr.iter() {234				let item = item?.clone();235				if let Val::Str(item) = item {236					if !first {237						out += &sep;238					}239					first = false;240					write!(out, "{item}").unwrap();241				} else if matches!(item, Val::Null) {242				} else {243					bail!("in std.join all items should be strings");244				}245			}246247			IndexableVal::Str(out.into())248		}249	})250}251252#[builtin]253pub fn builtin_lines(arr: ArrValue) -> Result<IndexableVal> {254	builtin_join(255		IndexableVal::Str("\n".into()),256		ArrValue::extended(arr, ArrValue::new(vec![Val::string("")]))257			.ok_or_else(|| error!("array is too large"))?,258	)259}260261#[builtin]262pub fn builtin_resolve_path(f: String, r: String) -> String {263	let Some(pos) = f.rfind('/') else {264		return r;265	};266	format!("{}{}", &f[..=pos], r)267}268269pub fn deep_join_inner(out: &mut String, arr: IndexableVal) -> Result<()> {270	use std::fmt::Write;271	match arr {272		IndexableVal::Str(s) => write!(out, "{s}").expect("no error"),273		IndexableVal::Arr(arr) => {274			for ele in arr.iter() {275				let indexable = IndexableVal::from_untyped(ele?)?;276				deep_join_inner(out, indexable)?;277			}278		}279	}280	Ok(())281}282283#[builtin]284pub fn builtin_deep_join(arr: IndexableVal) -> Result<String> {285	let mut out = String::new();286	deep_join_inner(&mut out, arr)?;287	Ok(out)288}289290#[builtin]291pub fn builtin_reverse(arr: ArrValue) -> ArrValue {292	arr.reversed()293}294295#[builtin]296pub fn builtin_any(arr: ArrValue) -> Result<bool> {297	for v in arr.iter() {298		let v = bool::from_untyped(v?)?;299		if v {300			return Ok(true);301		}302	}303	Ok(false)304}305306#[builtin]307pub fn builtin_all(arr: ArrValue) -> Result<bool> {308	for v in arr.iter() {309		let v = bool::from_untyped(v?)?;310		if !v {311			return Ok(false);312		}313	}314	Ok(true)315}316317#[builtin]318pub fn builtin_member(arr: IndexableVal, x: Val) -> Result<bool> {319	match arr {320		IndexableVal::Str(str) => {321			let x: IStr = IStr::from_untyped(x)?;322			Ok(!x.is_empty() && str.contains(&*x))323		}324		IndexableVal::Arr(a) => {325			for item in a.iter() {326				let item = item?;327				if equals(&item, &x)? {328					return Ok(true);329				}330			}331			Ok(false)332		}333	}334}335336#[builtin]337pub fn builtin_find(value: Val, arr: ArrValue) -> Result<Vec<usize>> {338	let mut out = Vec::new();339	for (i, ele) in arr.iter().enumerate() {340		let ele = ele?;341		if equals(&ele, &value)? {342			out.push(i);343		}344	}345	Ok(out)346}347348#[builtin]349pub fn builtin_contains(arr: IndexableVal, elem: Val) -> Result<bool> {350	builtin_member(arr, elem)351}352353#[builtin]354pub fn builtin_count(arr: ArrValue, x: Val) -> Result<usize> {355	let mut count = 0;356	for item in arr.iter() {357		if equals(&item?, &x)? {358			count += 1;359		}360	}361	Ok(count)362}363364#[builtin]365pub fn builtin_avg(arr: Vec<f64>, onEmpty: Option<Thunk<Val>>) -> Result<Val> {366	if arr.is_empty() {367		return eval_on_empty(onEmpty);368	}369	#[expect(370		clippy::cast_precision_loss,371		reason = "array sizes are bounded to i32 len"372	)]373	Ok(Val::try_num(arr.iter().sum::<f64>() / (arr.len() as f64))?)374}375376#[builtin]377pub fn builtin_remove_at(arr: ArrValue, at: i32) -> Result<ArrValue> {378	let newArrLeft = arr.clone().slice(None, Some(at), None);379	let newArrRight = arr.slice(Some(at + 1), None, None);380381	Ok(ArrValue::extended(newArrLeft, newArrRight).ok_or_else(|| error!("array is too large"))?)382}383384#[builtin]385pub fn builtin_remove(arr: ArrValue, elem: Val) -> Result<ArrValue> {386	for (index, item) in arr.iter().enumerate() {387		if equals(&item?, &elem)? {388			#[expect(389				clippy::cast_possible_truncation,390				clippy::cast_possible_wrap,391				reason = "array sizes are bounded to i32 len"392			)]393			return builtin_remove_at(arr.clone(), index as i32);394		}395	}396	Ok(arr)397}398399#[builtin]400pub fn builtin_flatten_arrays(arrs: Vec<ArrValue>) -> Result<ArrValue> {401	pub fn flatten_inner(values: &[ArrValue]) -> Result<ArrValue> {402		if values.len() == 1 {403			return Ok(values[0].clone());404		} else if values.len() == 2 {405			return ArrValue::extended(values[0].clone(), values[1].clone())406				.ok_or_else(|| error!("array is too large"));407		}408		let (a, b) = values.split_at(values.len() / 2);409		ArrValue::extended(flatten_inner(a)?, flatten_inner(b)?)410			.ok_or_else(|| error!("array is too large"))411	}412	if arrs.is_empty() {413		return Ok(ArrValue::empty());414	} else if arrs.len() == 1 {415		return Ok(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}