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

difftreelog

source

crates/jrsonnet-stdlib/src/arrays.rs9.7 KiBsourcehistory
1#![allow(non_snake_case)]23use jrsonnet_evaluator::{4	bail,5	function::{builtin, FuncVal},6	runtime_error,7	typed::{BoundedI32, BoundedUsize, Either2, NativeFn, Typed},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: BoundedI32<0, { i32::MAX }>, func: FuncVal) -> Result<ArrValue> {22	if *sz == 0 {23		return Ok(ArrValue::empty());24	}25	func.evaluate_trivial().map_or_else(26		|| Ok(ArrValue::range_exclusive(0, *sz).map(func)),27		|trivial| {28			let mut out = Vec::with_capacity(*sz as usize);29			for _ in 0..*sz {30				out.push(trivial.clone());31			}32			Ok(ArrValue::eager(out))33		},34	)35}3637#[builtin]38pub fn builtin_repeat(what: Either![IStr, ArrValue], count: usize) -> Result<Val> {39	Ok(match what {40		Either2::A(s) => Val::string(s.repeat(count)),41		Either2::B(arr) => Val::Arr(42			ArrValue::repeated(arr, count)43				.ok_or_else(|| runtime_error!("repeated length overflow"))?,44		),45	})46}4748#[builtin]49pub fn builtin_slice(50	indexable: IndexableVal,51	index: Option<i32>,52	end: Option<i32>,53	step: Option<BoundedUsize<1, { i32::MAX as usize }>>,54) -> Result<Val> {55	indexable.slice(index, end, step).map(Val::from)56}5758#[builtin]59pub fn builtin_map(func: FuncVal, arr: IndexableVal) -> ArrValue {60	let arr = arr.to_array();61	arr.map(func)62}6364#[builtin]65pub fn builtin_map_with_index(func: FuncVal, arr: IndexableVal) -> ArrValue {66	let arr = arr.to_array();67	arr.map_with_index(func)68}6970#[builtin]71pub fn builtin_map_with_key(func: FuncVal, obj: ObjValue) -> Result<ObjValue> {72	let mut out = ObjValueBuilder::new();73	for (k, v) in obj.iter() {74		let v = v?;75		out.field(k).value(func.evaluate_simple(&(v,), false)?);76	}77	Ok(out.build())78}7980#[builtin]81pub fn builtin_flatmap(82	func: NativeFn<((Either![String, Val],), Val)>,83	arr: IndexableVal,84) -> Result<IndexableVal> {85	use std::fmt::Write;86	match arr {87		IndexableVal::Str(str) => {88			let mut out = String::new();89			for c in str.chars() {90				match func(Either2::A(c.to_string()))? {91					Val::Str(o) => write!(out, "{o}").unwrap(),92					Val::Null => continue,93					_ => bail!("in std.join all items should be strings"),94				};95			}96			Ok(IndexableVal::Str(out.into()))97		}98		IndexableVal::Arr(a) => {99			let mut out = Vec::new();100			for el in a.iter() {101				let el = el?;102				match func(Either2::B(el))? {103					Val::Arr(o) => {104						for oe in o.iter() {105							out.push(oe?);106						}107					}108					Val::Null => continue,109					_ => bail!("in std.join all items should be arrays"),110				};111			}112			Ok(IndexableVal::Arr(out.into()))113		}114	}115}116117#[builtin]118pub fn builtin_filter(func: FuncVal, arr: ArrValue) -> Result<ArrValue> {119	arr.filter(|val| bool::from_untyped(func.evaluate_simple(&(val.clone(),), false)?))120}121122#[builtin]123pub fn builtin_filter_map(124	filter_func: FuncVal,125	map_func: FuncVal,126	arr: ArrValue,127) -> Result<ArrValue> {128	Ok(builtin_filter(filter_func, arr)?.map(map_func))129}130131#[builtin]132pub fn builtin_foldl(func: FuncVal, arr: ArrValue, init: Val) -> Result<Val> {133	let mut acc = init;134	for i in arr.iter() {135		acc = func.evaluate_simple(&(acc, i?), false)?;136	}137	Ok(acc)138}139140#[builtin]141pub fn builtin_foldr(func: FuncVal, arr: ArrValue, init: Val) -> Result<Val> {142	let mut acc = init;143	for i in arr.iter().rev() {144		acc = func.evaluate_simple(&(i?, acc), false)?;145	}146	Ok(acc)147}148149#[builtin]150pub fn builtin_range(from: i32, to: i32) -> Result<ArrValue> {151	if to < from {152		return Ok(ArrValue::empty());153	}154	Ok(ArrValue::range_inclusive(from, to))155}156157#[builtin]158pub fn builtin_join(sep: IndexableVal, arr: ArrValue) -> Result<IndexableVal> {159	use std::fmt::Write;160	Ok(match sep {161		IndexableVal::Arr(joiner_items) => {162			let mut out = Vec::new();163164			let mut first = true;165			for item in arr.iter() {166				let item = item?.clone();167				if let Val::Arr(items) = item {168					if !first {169						out.reserve(joiner_items.len());170						// TODO: extend171						for item in joiner_items.iter() {172							out.push(item?);173						}174					}175					first = false;176					out.reserve(items.len());177					for item in items.iter() {178						out.push(item?);179					}180				} else if matches!(item, Val::Null) {181					continue;182				} else {183					bail!("in std.join all items should be arrays");184				}185			}186187			IndexableVal::Arr(out.into())188		}189		IndexableVal::Str(sep) => {190			let mut out = String::new();191192			let mut first = true;193			for item in arr.iter() {194				let item = item?.clone();195				if let Val::Str(item) = item {196					if !first {197						out += &sep;198					}199					first = false;200					write!(out, "{item}").unwrap();201				} else if matches!(item, Val::Null) {202					continue;203				} else {204					bail!("in std.join all items should be strings");205				}206			}207208			IndexableVal::Str(out.into())209		}210	})211}212213#[builtin]214pub fn builtin_lines(arr: ArrValue) -> Result<IndexableVal> {215	builtin_join(216		IndexableVal::Str("\n".into()),217		ArrValue::extended(arr, ArrValue::eager(vec![Val::string("")])),218	)219}220221#[builtin]222pub fn builtin_resolve_path(f: String, r: String) -> String {223	let Some(pos) = f.rfind('/') else {224		return r;225	};226	format!("{}{}", &f[..=pos], r)227}228229pub fn deep_join_inner(out: &mut String, arr: IndexableVal) -> Result<()> {230	use std::fmt::Write;231	match arr {232		IndexableVal::Str(s) => write!(out, "{s}").expect("no error"),233		IndexableVal::Arr(arr) => {234			for ele in arr.iter() {235				let indexable = IndexableVal::from_untyped(ele?)?;236				deep_join_inner(out, indexable)?;237			}238		}239	}240	Ok(())241}242243#[builtin]244pub fn builtin_deep_join(arr: IndexableVal) -> Result<String> {245	let mut out = String::new();246	deep_join_inner(&mut out, arr)?;247	Ok(out)248}249250#[builtin]251pub fn builtin_reverse(arr: ArrValue) -> ArrValue {252	arr.reversed()253}254255#[builtin]256pub fn builtin_any(arr: ArrValue) -> Result<bool> {257	for v in arr.iter() {258		let v = bool::from_untyped(v?)?;259		if v {260			return Ok(true);261		}262	}263	Ok(false)264}265266#[builtin]267pub fn builtin_all(arr: ArrValue) -> Result<bool> {268	for v in arr.iter() {269		let v = bool::from_untyped(v?)?;270		if !v {271			return Ok(false);272		}273	}274	Ok(true)275}276277#[builtin]278pub fn builtin_member(arr: IndexableVal, x: Val) -> Result<bool> {279	match arr {280		IndexableVal::Str(str) => {281			let x: IStr = IStr::from_untyped(x)?;282			Ok(!x.is_empty() && str.contains(&*x))283		}284		IndexableVal::Arr(a) => {285			for item in a.iter() {286				let item = item?;287				if equals(&item, &x)? {288					return Ok(true);289				}290			}291			Ok(false)292		}293	}294}295296#[builtin]297pub fn builtin_find(value: Val, arr: ArrValue) -> Result<Vec<usize>> {298	let mut out = Vec::new();299	for (i, ele) in arr.iter().enumerate() {300		let ele = ele?;301		if equals(&ele, &value)? {302			out.push(i);303		}304	}305	Ok(out)306}307308#[builtin]309pub fn builtin_contains(arr: IndexableVal, elem: Val) -> Result<bool> {310	builtin_member(arr, elem)311}312313#[builtin]314pub fn builtin_count(arr: ArrValue, x: Val) -> Result<usize> {315	let mut count = 0;316	for item in arr.iter() {317		if equals(&item?, &x)? {318			count += 1;319		}320	}321	Ok(count)322}323324#[builtin]325pub fn builtin_avg(arr: Vec<f64>, onEmpty: Option<Thunk<Val>>) -> Result<Val> {326	if arr.is_empty() {327		return eval_on_empty(onEmpty);328	}329	Ok(Val::try_num(arr.iter().sum::<f64>() / (arr.len() as f64))?)330}331332#[builtin]333pub fn builtin_remove_at(arr: ArrValue, at: i32) -> Result<ArrValue> {334	let newArrLeft = arr.clone().slice(None, Some(at), None);335	let newArrRight = arr.slice(Some(at + 1), None, None);336337	Ok(ArrValue::extended(newArrLeft, newArrRight))338}339340#[builtin]341pub fn builtin_remove(arr: ArrValue, elem: Val) -> Result<ArrValue> {342	for (index, item) in arr.iter().enumerate() {343		if equals(&item?, &elem)? {344			return builtin_remove_at(arr.clone(), index as i32);345		}346	}347	Ok(arr)348}349350#[builtin]351pub fn builtin_flatten_arrays(arrs: Vec<ArrValue>) -> ArrValue {352	pub fn flatten_inner(values: &[ArrValue]) -> ArrValue {353		if values.len() == 1 {354			return values[0].clone();355		} else if values.len() == 2 {356			return ArrValue::extended(values[0].clone(), values[1].clone());357		}358		let (a, b) = values.split_at(values.len() / 2);359		ArrValue::extended(flatten_inner(a), flatten_inner(b))360	}361	if arrs.is_empty() {362		return ArrValue::empty();363	} else if arrs.len() == 1 {364		return arrs.into_iter().next().expect("single");365	}366	flatten_inner(&arrs)367}368369#[builtin]370pub fn builtin_flatten_deep_array(value: Val) -> Result<Vec<Val>> {371	fn process(value: Val, out: &mut Vec<Val>) -> Result<()> {372		match value {373			Val::Arr(arr) => {374				for ele in arr.iter() {375					process(ele?, out)?;376				}377			}378			_ => out.push(value),379		}380		Ok(())381	}382	let mut out = Vec::new();383	process(value, &mut out)?;384	Ok(out)385}386387#[builtin]388pub fn builtin_prune(389	a: Val,390391	#[default(false)]392	#[cfg(feature = "exp-preserve-order")]393	preserve_order: bool,394) -> Result<Val> {395	fn is_content(val: &Val) -> bool {396		match val {397			Val::Null => false,398			Val::Arr(a) => !a.is_empty(),399			Val::Obj(o) => !o.is_empty(),400			_ => true,401		}402	}403	Ok(match a {404		Val::Arr(a) => {405			let mut out = Vec::new();406			for (i, ele) in a.iter().enumerate() {407				let ele = ele408					.and_then(|v| {409						builtin_prune(410							v,411							#[cfg(feature = "exp-preserve-order")]412							preserve_order,413						)414					})415					.with_description(|| format!("elem <{i}> pruning"))?;416				if is_content(&ele) {417					out.push(ele);418				}419			}420			Val::Arr(ArrValue::eager(out))421		}422		Val::Obj(o) => {423			let mut out = ObjValueBuilder::new();424			for (name, value) in o.iter(425				#[cfg(feature = "exp-preserve-order")]426				preserve_order,427			) {428				let value = value429					.and_then(|v| {430						builtin_prune(431							v,432							#[cfg(feature = "exp-preserve-order")]433							preserve_order,434						)435					})436					.with_description(|| format!("field <{name}> pruning"))?;437				if !is_content(&value) {438					continue;439				}440				out.field(name).value(value);441			}442			Val::Obj(out.build())443		}444		_ => a,445	})446}