git.delta.rocks / jrsonnet / refs/heads / master

difftreelog

source

crates/jrsonnet-stdlib/src/sets.rs3.9 KiBsourcehistory
1use std::cmp::Ordering;23use jrsonnet_evaluator::{Result, Thunk, Val, function::builtin, val::ArrValue};45use crate::keyf::KeyF;67#[builtin]8#[allow(non_snake_case)]9pub fn builtin_set_member(x: Thunk<Val>, arr: ArrValue, #[default] keyF: KeyF) -> Result<bool> {10	let mut low = 0;11	let mut high = arr.len32();1213	let x = keyF.eval(x)?;1415	while low < high {16		let middle = u32::midpoint(high, low);17		let comp = keyF.eval(arr.get_lazy32(middle).expect("in bounds"))?;18		match Val::try_cmp(&comp, &x)? {19			Ordering::Less => low = middle + 1,20			Ordering::Equal => return Ok(true),21			Ordering::Greater => high = middle,22		}23	}24	Ok(false)25}2627#[builtin]28#[allow(non_snake_case, clippy::redundant_closure)]29pub fn builtin_set_inter(30	a: ArrValue,31	b: ArrValue,32	#[default] keyF: KeyF,33) -> Result<Vec<Thunk<Val>>> {34	let mut a = a.iter_lazy();35	let mut b = b.iter_lazy();3637	let keyF = |v| keyF.eval(v);3839	let mut av = a.next();40	let mut bv = b.next();41	let mut ak = av.clone().map(keyF).transpose()?;42	let mut bk = bv.map(keyF).transpose()?;4344	let mut out = Vec::new();45	while let (Some(ac), Some(bc)) = (&ak, &bk) {46		match Val::try_cmp(ac, bc)? {47			Ordering::Less => {48				av = a.next();49				ak = av.clone().map(keyF).transpose()?;50			}51			Ordering::Greater => {52				bv = b.next();53				bk = bv.map(keyF).transpose()?;54			}55			Ordering::Equal => {56				out.push(av.clone().expect("ak != None => av != None"));57				av = a.next();58				ak = av.clone().map(keyF).transpose()?;59				bv = b.next();60				bk = bv.map(keyF).transpose()?;61			}62		}63	}64	Ok(out)65}6667#[builtin]68#[allow(non_snake_case, clippy::redundant_closure)]69pub fn builtin_set_diff(70	a: ArrValue,71	b: ArrValue,72	#[default] keyF: KeyF,73) -> Result<Vec<Thunk<Val>>> {74	let mut a = a.iter_lazy();75	let mut b = b.iter_lazy();7677	let keyF = |v| keyF.eval(v);7879	let mut av = a.next();80	let mut bv = b.next();81	let mut ak = av.clone().map(keyF).transpose()?;82	let mut bk = bv.map(keyF).transpose()?;8384	let mut out = Vec::new();85	while let (Some(ac), Some(bc)) = (&ak, &bk) {86		match Val::try_cmp(ac, bc)? {87			Ordering::Less => {88				// In a, but not in b89				out.push(av.clone().expect("ak != None"));90				av = a.next();91				ak = av.clone().map(keyF).transpose()?;92			}93			Ordering::Greater => {94				bv = b.next();95				bk = bv.map(keyF).transpose()?;96			}97			Ordering::Equal => {98				av = a.next();99				ak = av.clone().map(keyF).transpose()?;100				bv = b.next();101				bk = bv.map(keyF).transpose()?;102			}103		}104	}105	while let Some(_ac) = &ak {106		// In a, but not in b107		out.push(av.clone().expect("ak != None"));108		av = a.next();109		ak = av.clone().map(keyF).transpose()?;110	}111	Ok(out)112}113114#[builtin]115#[allow(non_snake_case, clippy::redundant_closure)]116pub fn builtin_set_union(117	a: ArrValue,118	b: ArrValue,119	#[default] keyF: KeyF,120) -> Result<Vec<Thunk<Val>>> {121	let mut a = a.iter_lazy();122	let mut b = b.iter_lazy();123124	let keyF = |v| keyF.eval(v);125126	let mut av = a.next();127	let mut bv = b.next();128	let mut ak = av.clone().map(keyF).transpose()?;129	let mut bk = bv.clone().map(keyF).transpose()?;130131	let mut out = Vec::new();132	while let (Some(ac), Some(bc)) = (&ak, &bk) {133		match Val::try_cmp(ac, bc)? {134			Ordering::Less => {135				out.push(av.clone().expect("ak != None"));136				av = a.next();137				ak = av.clone().map(keyF).transpose()?;138			}139			Ordering::Greater => {140				out.push(bv.clone().expect("bk != None"));141				bv = b.next();142				bk = bv.clone().map(keyF).transpose()?;143			}144			Ordering::Equal => {145				// NOTE: order matters, values in `a` win146				out.push(av.clone().expect("ak != None"));147				av = a.next();148				ak = av.clone().map(keyF).transpose()?;149				bv = b.next();150				bk = bv.clone().map(keyF).transpose()?;151			}152		}153	}154	// a.len() > b.len()155	while let Some(_ac) = &ak {156		out.push(av.clone().expect("ak != None"));157		av = a.next();158		ak = av.clone().map(keyF).transpose()?;159	}160	// b.len() > a.len()161	while let Some(_bc) = &bk {162		out.push(bv.clone().expect("ak != None"));163		bv = b.next();164		bk = bv.clone().map(keyF).transpose()?;165	}166	Ok(out)167}