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

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	runtime_error,7	typed::{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(sz: u32, func: NativeFn!((u32,) -> Val)) -> Result<ArrValue> {21	if sz == 0 {22		return Ok(ArrValue::empty());23	}24	// Try eager evaluation: call func(i) immediately for each element.25	'eager: {26		let mut out = Vec::with_capacity(sz as usize);27		for i in 0..sz {28			match func.call(i) {29				Ok(v) => out.push(v),30				Err(_) => break 'eager,31			}32		}33		return Ok(ArrValue::new(out));34	}35	Ok(ArrValue::make(sz, func))36}3738#[builtin]39pub fn builtin_repeat(what: Either![IStr, ArrValue], count: u32) -> Result<Val> {40	Ok(match what {41		Either2::A(s) => Val::string(s.repeat(count as usize)),42		Either2::B(arr) => Val::Arr(43			ArrValue::repeated(arr, count)44				.ok_or_else(|| runtime_error!("repeated length overflow"))?,45		),46	})47}4849#[builtin]50pub fn builtin_slice(51	indexable: IndexableVal,52	index: Option<Option<i32>>,53	end: Option<Option<i32>>,54	step: Option<Option<BoundedUsize<1, { i32::MAX as usize }>>>,55) -> Result<Val> {56	indexable57		.slice(index.flatten(), end.flatten(), step.flatten())58		.map(Val::from)59}6061#[builtin]62pub fn builtin_map(func: NativeFn!((Val) -> Val), arr: IndexableVal) -> ArrValue {63	let arr = arr.to_array();64	arr.map(func)65}6667#[builtin]68pub fn builtin_map_with_index(func: NativeFn!((u32, Val) -> Val), arr: IndexableVal) -> ArrValue {69	let arr = arr.to_array();70	arr.map_with_index(func)71}7273#[builtin]74pub fn builtin_map_with_key(75	func: NativeFn!((IStr, Val) -> Val),76	obj: ObjValue,77) -> Result<ObjValue> {78	let mut out = ObjValueBuilder::new();79	for (k, v) in obj.iter(80		// Makes sense mapped object should be ordered the same way, should not break anything when the output is not ordered (the default).81		// The thrown error might be different, but jsonnet82		// does not specify the evaluation order.83		#[cfg(feature = "exp-preserve-order")]84		true,85	) {86		let v = v?;87		out.field(k.clone()).value(func.call(k, v)?);88	}89	Ok(out.build())90}9192#[builtin]93pub fn builtin_flatmap(94	func: NativeFn!((Either![String, Val]) -> Val),95	arr: IndexableVal,96) -> Result<IndexableVal> {97	use std::fmt::Write;98	match arr {99		IndexableVal::Str(str) => {100			let mut out = String::new();101			for c in str.chars() {102				match func.call(Either2::A(c.to_string()))? {103					Val::Str(o) => write!(out, "{o}").unwrap(),104					Val::Null => {}105					_ => bail!("in std.join all items should be strings"),106				}107			}108			Ok(IndexableVal::Str(out.into()))109		}110		IndexableVal::Arr(a) => {111			let mut out = Vec::new();112			for el in a.iter() {113				let el = el?;114				match func.call(Either2::B(el))? {115					Val::Arr(o) => {116						for oe in o.iter() {117							out.push(oe?);118						}119					}120					Val::Null => {}121					_ => bail!("in std.join all items should be arrays"),122				}123			}124			Ok(IndexableVal::Arr(out.into()))125		}126	}127}128129type FilterFunc = NativeFn!((Thunk<Val>) -> bool);130131#[builtin]132pub fn builtin_filter(func: FilterFunc, arr: ArrValue) -> Result<ArrValue> {133	arr.filter(func)134}135136#[builtin]137pub fn builtin_filter_map(138	filter_func: FilterFunc,139	map_func: NativeFn!((Val) -> Val),140	arr: ArrValue,141) -> Result<ArrValue> {142	Ok(arr.filter(filter_func)?.map(map_func))143}144145#[builtin]146pub fn builtin_foldl(147	func: NativeFn!((Val, Either![Val, char]) -> Val),148	arr: Either![ArrValue, IStr],149	init: Val,150) -> Result<Val> {151	let mut acc = init;152	match arr {153		Either2::A(arr) => {154			for i in arr.iter() {155				acc = func.call(acc, Either2::A(i?))?;156			}157		}158		Either2::B(arr) => {159			for c in arr.chars() {160				acc = func.call(acc, Either2::B(c))?;161			}162		}163	}164	Ok(acc)165}166167#[builtin]168pub fn builtin_foldr(169	func: NativeFn!((Either![Val, char], Val) -> Val),170	arr: Either![ArrValue, IStr],171	init: Val,172) -> Result<Val> {173	let mut acc = init;174	match arr {175		Either2::A(arr) => {176			for i in arr.iter().rev() {177				acc = func.call(Either2::A(i?), acc)?;178			}179		}180		Either2::B(arr) => {181			for c in arr.chars().rev() {182				acc = func.call(Either2::B(c), acc)?;183			}184		}185	}186	Ok(acc)187}188189#[builtin]190pub fn builtin_range(from: i32, to: i32) -> Result<ArrValue> {191	if to < from {192		return Ok(ArrValue::empty());193	}194	Ok(ArrValue::range_inclusive(from, to))195}196197#[builtin]198pub fn builtin_join(sep: IndexableVal, arr: ArrValue) -> Result<IndexableVal> {199	use std::fmt::Write;200	Ok(match sep {201		IndexableVal::Arr(joiner_items) => {202			let mut out = Vec::new();203204			let mut first = true;205			for item in arr.iter() {206				let item = item?.clone();207				if let Val::Arr(items) = item {208					if !first {209						out.reserve(joiner_items.len() as usize);210						// TODO: extend211						for item in joiner_items.iter() {212							out.push(item?);213						}214					}215					first = false;216					out.reserve(items.len() as usize);217					for item in items.iter() {218						out.push(item?);219					}220				} else if matches!(item, Val::Null) {221				} else {222					bail!("in std.join all items should be arrays");223				}224			}225226			IndexableVal::Arr(out.into())227		}228		IndexableVal::Str(sep) => {229			let mut out = String::new();230231			let mut first = true;232			for item in arr.iter() {233				let item = item?.clone();234				if let Val::Str(item) = item {235					if !first {236						out += &sep;237					}238					first = false;239					write!(out, "{item}").unwrap();240				} else if matches!(item, Val::Null) {241				} else {242					bail!("in std.join all items should be strings");243				}244			}245246			IndexableVal::Str(out.into())247		}248	})249}250251#[builtin]252pub fn builtin_lines(arr: ArrValue) -> Result<IndexableVal> {253	builtin_join(254		IndexableVal::Str("\n".into()),255		ArrValue::extended(arr, ArrValue::new(vec![Val::string("")]))256			.ok_or_else(|| error!("array is too large"))?,257	)258}259260#[builtin]261pub fn builtin_resolve_path(f: String, r: String) -> String {262	let Some(pos) = f.rfind('/') else {263		return r;264	};265	format!("{}{}", &f[..=pos], r)266}267268pub fn deep_join_inner(out: &mut String, arr: IndexableVal) -> Result<()> {269	use std::fmt::Write;270	match arr {271		IndexableVal::Str(s) => write!(out, "{s}").expect("no error"),272		IndexableVal::Arr(arr) => {273			for ele in arr.iter() {274				let indexable = IndexableVal::from_untyped(ele?)?;275				deep_join_inner(out, indexable)?;276			}277		}278	}279	Ok(())280}281282#[builtin]283pub fn builtin_deep_join(arr: IndexableVal) -> Result<String> {284	let mut out = String::new();285	deep_join_inner(&mut out, arr)?;286	Ok(out)287}288289#[builtin]290pub fn builtin_reverse(arr: ArrValue) -> ArrValue {291	arr.reversed()292}293294#[builtin]295pub fn builtin_any(arr: ArrValue) -> Result<bool> {296	for v in arr.iter() {297		let v = bool::from_untyped(v?)?;298		if v {299			return Ok(true);300		}301	}302	Ok(false)303}304305#[builtin]306pub fn builtin_all(arr: ArrValue) -> Result<bool> {307	for v in arr.iter() {308		let v = bool::from_untyped(v?)?;309		if !v {310			return Ok(false);311		}312	}313	Ok(true)314}315316#[builtin]317pub fn builtin_member(arr: IndexableVal, x: Val) -> Result<bool> {318	match arr {319		IndexableVal::Str(str) => {320			let x: IStr = IStr::from_untyped(x)?;321			Ok(!x.is_empty() && str.contains(&*x))322		}323		IndexableVal::Arr(a) => {324			for item in a.iter() {325				let item = item?;326				if equals(&item, &x)? {327					return Ok(true);328				}329			}330			Ok(false)331		}332	}333}334335#[builtin]336pub fn builtin_find(value: Val, arr: ArrValue) -> Result<Vec<usize>> {337	let mut out = Vec::new();338	for (i, ele) in arr.iter().enumerate() {339		let ele = ele?;340		if equals(&ele, &value)? {341			out.push(i);342		}343	}344	Ok(out)345}346347#[builtin]348pub fn builtin_contains(arr: IndexableVal, elem: Val) -> Result<bool> {349	builtin_member(arr, elem)350}351352#[builtin]353pub fn builtin_count(arr: ArrValue, x: Val) -> Result<usize> {354	let mut count = 0;355	for item in arr.iter() {356		if equals(&item?, &x)? {357			count += 1;358		}359	}360	Ok(count)361}362363#[builtin]364pub fn builtin_avg(arr: Vec<f64>, onEmpty: Option<Thunk<Val>>) -> Result<Val> {365	if arr.is_empty() {366		return eval_on_empty(onEmpty);367	}368	#[expect(369		clippy::cast_precision_loss,370		reason = "array sizes are bounded to i32 len"371	)]372	Ok(Val::try_num(arr.iter().sum::<f64>() / (arr.len() as f64))?)373}374375#[builtin]376pub fn builtin_remove_at(arr: ArrValue, at: i32) -> Result<ArrValue> {377	let newArrLeft = arr.clone().slice(None, Some(at), None);378	let newArrRight = arr.slice(Some(at + 1), None, None);379380	Ok(ArrValue::extended(newArrLeft, newArrRight).ok_or_else(|| error!("array is too large"))?)381}382383#[builtin]384pub fn builtin_remove(arr: ArrValue, elem: Val) -> Result<ArrValue> {385	for (index, item) in arr.iter().enumerate() {386		if equals(&item?, &elem)? {387			#[expect(388				clippy::cast_possible_truncation,389				clippy::cast_possible_wrap,390				reason = "array sizes are bounded to i32 len"391			)]392			return builtin_remove_at(arr.clone(), index as i32);393		}394	}395	Ok(arr)396}397398#[builtin]399pub fn builtin_flatten_arrays(arrs: Vec<ArrValue>) -> Result<ArrValue> {400	pub fn flatten_inner(values: &[ArrValue]) -> Result<ArrValue> {401		if values.len() == 1 {402			return Ok(values[0].clone());403		} else if values.len() == 2 {404			return ArrValue::extended(values[0].clone(), values[1].clone())405				.ok_or_else(|| error!("array is too large"));406		}407		let (a, b) = values.split_at(values.len() / 2);408		ArrValue::extended(flatten_inner(a)?, flatten_inner(b)?)409			.ok_or_else(|| error!("array is too large"))410	}411	if arrs.is_empty() {412		return Ok(ArrValue::empty());413	} else if arrs.len() == 1 {414		return Ok(arrs.into_iter().next().expect("single"));415	}416	flatten_inner(&arrs)417}418419#[builtin]420pub fn builtin_flatten_deep_array(value: Val) -> Result<Vec<Val>> {421	fn process(value: Val, out: &mut Vec<Val>) -> Result<()> {422		match value {423			Val::Arr(arr) => {424				for ele in arr.iter() {425					process(ele?, out)?;426				}427			}428			_ => out.push(value),429		}430		Ok(())431	}432	let mut out = Vec::new();433	process(value, &mut out)?;434	Ok(out)435}436437#[builtin]438pub fn builtin_prune(439	a: Val,440441	#[default(false)]442	#[cfg(feature = "exp-preserve-order")]443	preserve_order: bool,444) -> Result<Val> {445	fn is_content(val: &Val) -> bool {446		match val {447			Val::Null => false,448			Val::Arr(a) => !a.is_empty(),449			Val::Obj(o) => !o.is_empty(),450			_ => true,451		}452	}453	Ok(match a {454		Val::Arr(a) => {455			let mut out = Vec::new();456			for (i, ele) in a.iter().enumerate() {457				let ele = ele458					.and_then(|v| {459						builtin_prune(460							v,461							#[cfg(feature = "exp-preserve-order")]462							preserve_order,463						)464					})465					.with_description(|| format!("elem <{i}> pruning"))?;466				if is_content(&ele) {467					out.push(ele);468				}469			}470			Val::arr(out)471		}472		Val::Obj(o) => {473			let mut out = ObjValueBuilder::new();474			for (name, value) in o.iter(475				#[cfg(feature = "exp-preserve-order")]476				preserve_order,477			) {478				let value = value479					.and_then(|v| {480						builtin_prune(481							v,482							#[cfg(feature = "exp-preserve-order")]483							preserve_order,484						)485					})486					.with_description(|| format!("field <{name}> pruning"))?;487				if !is_content(&value) {488					continue;489				}490				out.field(name).value(value);491			}492			Val::Obj(out.build())493		}494		_ => a,495	})496}