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

difftreelog

perf move mapWithIndex to native

Yaroslav Bolyukin2024-06-18parent: #d5d1e03.patch.diff
in: master

5 files changed

modifiedcrates/jrsonnet-evaluator/src/arr/mod.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/arr/mod.rs
+++ b/crates/jrsonnet-evaluator/src/arr/mod.rs
@@ -54,7 +54,12 @@
 
 	#[must_use]
 	pub fn map(self, mapper: FuncVal) -> Self {
-		Self::new(MappedArray::new(self, mapper))
+		Self::new(<MappedArray<false>>::new(self, mapper))
+	}
+
+	#[must_use]
+	pub fn map_with_index(self, mapper: FuncVal) -> Self {
+		Self::new(<MappedArray<true>>::new(self, mapper))
 	}
 
 	pub fn filter(self, filter: impl Fn(&Val) -> Result<bool>) -> Result<Self> {
modifiedcrates/jrsonnet-evaluator/src/arr/spec.rsdiffbeforeafterboth
after · crates/jrsonnet-evaluator/src/arr/spec.rs
1use std::{any::Any, cell::RefCell, fmt::Debug, iter, mem::replace};23use jrsonnet_gcmodule::{Cc, Trace};4use jrsonnet_interner::{IBytes, IStr};5use jrsonnet_parser::LocExpr;67use super::ArrValue;8use crate::{9	error::ErrorKind::InfiniteRecursionDetected, evaluate, function::FuncVal, typed::Typed,10	val::ThunkValue, Context, Error, ObjValue, Result, Thunk, Val,11};1213pub trait ArrayLike: Any + Trace + Debug {14	fn len(&self) -> usize;15	fn is_empty(&self) -> bool {16		self.len() == 017	}18	fn get(&self, index: usize) -> Result<Option<Val>>;19	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>>;20	fn get_cheap(&self, index: usize) -> Option<Val>;2122	fn is_cheap(&self) -> bool;23}2425#[derive(Debug, Trace)]26pub struct SliceArray {27	pub(crate) inner: ArrValue,28	pub(crate) from: u32,29	pub(crate) to: u32,30	pub(crate) step: u32,31}3233impl SliceArray {34	fn iter(&self) -> impl Iterator<Item = Result<Val>> + '_ {35		self.inner36			.iter()37			.skip(self.from as usize)38			.take((self.to - self.from) as usize)39			.step_by(self.step as usize)40	}4142	fn iter_lazy(&self) -> impl Iterator<Item = Thunk<Val>> + '_ {43		self.inner44			.iter_lazy()45			.skip(self.from as usize)46			.take((self.to - self.from) as usize)47			.step_by(self.step as usize)48	}4950	fn iter_cheap(&self) -> Option<impl crate::arr::ArrayLikeIter<Val> + '_> {51		Some(52			self.inner53				.iter_cheap()?54				.skip(self.from as usize)55				.take((self.to - self.from) as usize)56				.step_by(self.step as usize),57		)58	}59}60impl ArrayLike for SliceArray {61	fn len(&self) -> usize {62		iter::repeat(())63			.take((self.to - self.from) as usize)64			.step_by(self.step as usize)65			.count()66	}6768	fn get(&self, index: usize) -> Result<Option<Val>> {69		self.iter().nth(index).transpose()70	}7172	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {73		self.iter_lazy().nth(index)74	}7576	fn get_cheap(&self, index: usize) -> Option<Val> {77		self.iter_cheap()?.nth(index)78	}79	fn is_cheap(&self) -> bool {80		self.inner.is_cheap()81	}82}8384#[derive(Trace, Debug)]85pub struct CharArray(pub Vec<char>);86impl ArrayLike for CharArray {87	fn len(&self) -> usize {88		self.0.len()89	}9091	fn get(&self, index: usize) -> Result<Option<Val>> {92		Ok(self.get_cheap(index))93	}9495	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {96		self.get_cheap(index).map(Thunk::evaluated)97	}9899	fn get_cheap(&self, index: usize) -> Option<Val> {100		self.0.get(index).map(|v| Val::string(*v))101	}102	fn is_cheap(&self) -> bool {103		true104	}105}106107#[derive(Trace, Debug)]108pub struct BytesArray(pub IBytes);109impl ArrayLike for BytesArray {110	fn len(&self) -> usize {111		self.0.len()112	}113114	fn get(&self, index: usize) -> Result<Option<Val>> {115		Ok(self.get_cheap(index))116	}117118	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {119		self.get_cheap(index).map(Thunk::evaluated)120	}121122	fn get_cheap(&self, index: usize) -> Option<Val> {123		self.0.get(index).map(|v| Val::Num(f64::from(*v)))124	}125	fn is_cheap(&self) -> bool {126		true127	}128}129130#[derive(Debug, Trace, Clone)]131enum ArrayThunk<T: 'static + Trace> {132	Computed(Val),133	Errored(Error),134	Waiting(T),135	Pending,136}137138#[derive(Debug, Trace, Clone)]139pub struct ExprArray {140	ctx: Context,141	cached: Cc<RefCell<Vec<ArrayThunk<LocExpr>>>>,142}143impl ExprArray {144	pub fn new(ctx: Context, items: impl IntoIterator<Item = LocExpr>) -> Self {145		Self {146			ctx,147			cached: Cc::new(RefCell::new(148				items.into_iter().map(ArrayThunk::Waiting).collect(),149			)),150		}151	}152}153impl ArrayLike for ExprArray {154	fn len(&self) -> usize {155		self.cached.borrow().len()156	}157	fn get(&self, index: usize) -> Result<Option<Val>> {158		if index >= self.len() {159			return Ok(None);160		}161		match &self.cached.borrow()[index] {162			ArrayThunk::Computed(c) => return Ok(Some(c.clone())),163			ArrayThunk::Errored(e) => return Err(e.clone()),164			ArrayThunk::Pending => return Err(InfiniteRecursionDetected.into()),165			ArrayThunk::Waiting(..) => {}166		};167168		let ArrayThunk::Waiting(expr) =169			replace(&mut self.cached.borrow_mut()[index], ArrayThunk::Pending)170		else {171			unreachable!()172		};173174		let new_value = match evaluate(self.ctx.clone(), &expr) {175			Ok(v) => v,176			Err(e) => {177				self.cached.borrow_mut()[index] = ArrayThunk::Errored(e.clone());178				return Err(e);179			}180		};181		self.cached.borrow_mut()[index] = ArrayThunk::Computed(new_value.clone());182		Ok(Some(new_value))183	}184	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {185		#[derive(Trace)]186		struct ArrayElement {187			arr_thunk: ExprArray,188			index: usize,189		}190191		impl ThunkValue for ArrayElement {192			type Output = Val;193194			fn get(self: Box<Self>) -> Result<Self::Output> {195				self.arr_thunk196					.get(self.index)197					.transpose()198					.expect("index checked")199			}200		}201202		if index >= self.len() {203			return None;204		}205		match &self.cached.borrow()[index] {206			ArrayThunk::Computed(c) => return Some(Thunk::evaluated(c.clone())),207			ArrayThunk::Errored(e) => return Some(Thunk::errored(e.clone())),208			ArrayThunk::Waiting(_) | ArrayThunk::Pending => {}209		};210211		Some(Thunk::new(ArrayElement {212			arr_thunk: self.clone(),213			index,214		}))215	}216	fn get_cheap(&self, _index: usize) -> Option<Val> {217		None218	}219	fn is_cheap(&self) -> bool {220		false221	}222}223224#[derive(Trace, Debug)]225pub struct ExtendedArray {226	pub a: ArrValue,227	pub b: ArrValue,228	split: usize,229	len: usize,230}231impl ExtendedArray {232	pub fn new(a: ArrValue, b: ArrValue) -> Self {233		let a_len = a.len();234		let b_len = b.len();235		Self {236			a,237			b,238			split: a_len,239			len: a_len.checked_add(b_len).expect("too large array value"),240		}241	}242}243244struct WithExactSize<I>(I, usize);245impl<I, T> Iterator for WithExactSize<I>246where247	I: Iterator<Item = T>,248{249	type Item = T;250251	fn next(&mut self) -> Option<Self::Item> {252		self.0.next()253	}254	fn nth(&mut self, n: usize) -> Option<Self::Item> {255		self.0.nth(n)256	}257	fn size_hint(&self) -> (usize, Option<usize>) {258		(self.1, Some(self.1))259	}260}261impl<I> DoubleEndedIterator for WithExactSize<I>262where263	I: DoubleEndedIterator,264{265	fn next_back(&mut self) -> Option<Self::Item> {266		self.0.next_back()267	}268	fn nth_back(&mut self, n: usize) -> Option<Self::Item> {269		self.0.nth_back(n)270	}271}272impl<I> ExactSizeIterator for WithExactSize<I>273where274	I: Iterator,275{276	fn len(&self) -> usize {277		self.1278	}279}280impl ArrayLike for ExtendedArray {281	fn get(&self, index: usize) -> Result<Option<Val>> {282		if self.split > index {283			self.a.get(index)284		} else {285			self.b.get(index - self.split)286		}287	}288	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {289		if self.split > index {290			self.a.get_lazy(index)291		} else {292			self.b.get_lazy(index - self.split)293		}294	}295296	fn len(&self) -> usize {297		self.len298	}299300	fn get_cheap(&self, index: usize) -> Option<Val> {301		if self.split > index {302			self.a.get_cheap(index)303		} else {304			self.b.get_cheap(index - self.split)305		}306	}307	fn is_cheap(&self) -> bool {308		self.a.is_cheap() && self.b.is_cheap()309	}310}311312#[derive(Trace, Debug)]313pub struct LazyArray(pub Vec<Thunk<Val>>);314impl ArrayLike for LazyArray {315	fn len(&self) -> usize {316		self.0.len()317	}318	fn get(&self, index: usize) -> Result<Option<Val>> {319		let Some(v) = self.0.get(index) else {320			return Ok(None);321		};322		v.evaluate().map(Some)323	}324	fn get_cheap(&self, _index: usize) -> Option<Val> {325		None326	}327	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {328		self.0.get(index).cloned()329	}330	fn is_cheap(&self) -> bool {331		false332	}333}334335#[derive(Trace, Debug)]336pub struct EagerArray(pub Vec<Val>);337impl ArrayLike for EagerArray {338	fn len(&self) -> usize {339		self.0.len()340	}341342	fn get(&self, index: usize) -> Result<Option<Val>> {343		Ok(self.0.get(index).cloned())344	}345346	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {347		self.0.get(index).cloned().map(Thunk::evaluated)348	}349350	fn get_cheap(&self, index: usize) -> Option<Val> {351		self.0.get(index).cloned()352	}353	fn is_cheap(&self) -> bool {354		true355	}356}357358/// Inclusive range type359#[derive(Debug, Trace, PartialEq, Eq)]360pub struct RangeArray {361	start: i32,362	end: i32,363}364impl RangeArray {365	pub fn empty() -> Self {366		Self::new_exclusive(0, 0)367	}368	pub fn new_exclusive(start: i32, end: i32) -> Self {369		end.checked_sub(1)370			.map_or_else(Self::empty, |end| Self { start, end })371	}372	pub fn new_inclusive(start: i32, end: i32) -> Self {373		Self { start, end }374	}375	fn range(&self) -> impl ExactSizeIterator<Item = i32> + DoubleEndedIterator {376		WithExactSize(377			self.start..=self.end,378			(self.end as usize)379				.wrapping_sub(self.start as usize)380				.wrapping_add(1),381		)382	}383}384385impl ArrayLike for RangeArray {386	fn len(&self) -> usize {387		self.range().len()388	}389	fn is_empty(&self) -> bool {390		self.range().len() == 0391	}392393	fn get(&self, index: usize) -> Result<Option<Val>> {394		Ok(self.get_cheap(index))395	}396397	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {398		self.get_cheap(index).map(Thunk::evaluated)399	}400401	fn get_cheap(&self, index: usize) -> Option<Val> {402		self.range().nth(index).map(|i| Val::Num(f64::from(i)))403	}404	fn is_cheap(&self) -> bool {405		true406	}407}408409#[derive(Debug, Trace)]410pub struct ReverseArray(pub ArrValue);411impl ArrayLike for ReverseArray {412	fn len(&self) -> usize {413		self.0.len()414	}415416	fn get(&self, index: usize) -> Result<Option<Val>> {417		self.0.get(self.0.len() - index - 1)418	}419420	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {421		self.0.get_lazy(self.0.len() - index - 1)422	}423424	fn get_cheap(&self, index: usize) -> Option<Val> {425		self.0.get_cheap(self.0.len() - index - 1)426	}427	fn is_cheap(&self) -> bool {428		self.0.is_cheap()429	}430}431432#[derive(Trace, Debug, Clone)]433pub struct MappedArray<const WithIndex: bool> {434	inner: ArrValue,435	cached: Cc<RefCell<Vec<ArrayThunk<()>>>>,436	mapper: FuncVal,437}438impl<const WithIndex: bool> MappedArray<WithIndex> {439	pub fn new(inner: ArrValue, mapper: FuncVal) -> Self {440		let len = inner.len();441		Self {442			inner,443			cached: Cc::new(RefCell::new(vec![ArrayThunk::Waiting(()); len])),444			mapper,445		}446	}447	fn evaluate(&self, index: usize, value: Val) -> Result<Val> {448		if WithIndex {449			self.mapper.evaluate_simple(&(index, value), false)450		} else {451			self.mapper.evaluate_simple(&(value,), false)452		}453	}454}455impl<const WithIndex: bool> ArrayLike for MappedArray<WithIndex> {456	fn len(&self) -> usize {457		self.cached.borrow().len()458	}459460	fn get(&self, index: usize) -> Result<Option<Val>> {461		if index >= self.len() {462			return Ok(None);463		}464		match &self.cached.borrow()[index] {465			ArrayThunk::Computed(c) => return Ok(Some(c.clone())),466			ArrayThunk::Errored(e) => return Err(e.clone()),467			ArrayThunk::Pending => return Err(InfiniteRecursionDetected.into()),468			ArrayThunk::Waiting(..) => {}469		};470471		let ArrayThunk::Waiting(()) =472			replace(&mut self.cached.borrow_mut()[index], ArrayThunk::Pending)473		else {474			unreachable!()475		};476477		let val = self478			.inner479			.get(index)480			.transpose()481			.expect("index checked")482			.and_then(|r| self.evaluate(index, r));483484		let new_value = match val {485			Ok(v) => v,486			Err(e) => {487				self.cached.borrow_mut()[index] = ArrayThunk::Errored(e.clone());488				return Err(e);489			}490		};491		self.cached.borrow_mut()[index] = ArrayThunk::Computed(new_value.clone());492		Ok(Some(new_value))493	}494	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {495		#[derive(Trace)]496		struct ArrayElement<const WithIndex: bool> {497			arr_thunk: MappedArray<WithIndex>,498			index: usize,499		}500501		impl<const WithIndex: bool> ThunkValue for ArrayElement<WithIndex> {502			type Output = Val;503504			fn get(self: Box<Self>) -> Result<Self::Output> {505				self.arr_thunk506					.get(self.index)507					.transpose()508					.expect("index checked")509			}510		}511512		if index >= self.len() {513			return None;514		}515		match &self.cached.borrow()[index] {516			ArrayThunk::Computed(c) => return Some(Thunk::evaluated(c.clone())),517			ArrayThunk::Errored(e) => return Some(Thunk::errored(e.clone())),518			ArrayThunk::Waiting(()) | ArrayThunk::Pending => {}519		};520521		Some(Thunk::new(ArrayElement {522			arr_thunk: self.clone(),523			index,524		}))525	}526527	fn get_cheap(&self, _index: usize) -> Option<Val> {528		None529	}530	fn is_cheap(&self) -> bool {531		false532	}533}534535#[derive(Trace, Debug)]536pub struct RepeatedArray {537	data: ArrValue,538	repeats: usize,539	total_len: usize,540}541impl RepeatedArray {542	pub fn new(data: ArrValue, repeats: usize) -> Option<Self> {543		let total_len = data.len().checked_mul(repeats)?;544		Some(Self {545			data,546			repeats,547			total_len,548		})549	}550}551552impl ArrayLike for RepeatedArray {553	fn len(&self) -> usize {554		self.total_len555	}556557	fn get(&self, index: usize) -> Result<Option<Val>> {558		if index > self.total_len {559			return Ok(None);560		}561		self.data.get(index % self.data.len())562	}563564	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {565		if index > self.total_len {566			return None;567		}568		self.data.get_lazy(index % self.data.len())569	}570571	fn get_cheap(&self, index: usize) -> Option<Val> {572		if index > self.total_len {573			return None;574		}575		self.data.get_cheap(index % self.data.len())576	}577	fn is_cheap(&self) -> bool {578		self.data.is_cheap()579	}580}581582#[derive(Trace, Debug)]583pub struct PickObjectValues {584	obj: ObjValue,585	keys: Vec<IStr>,586}587588impl PickObjectValues {589	pub fn new(obj: ObjValue, keys: Vec<IStr>) -> Self {590		Self { obj, keys }591	}592}593594impl ArrayLike for PickObjectValues {595	fn len(&self) -> usize {596		self.keys.len()597	}598599	fn get(&self, index: usize) -> Result<Option<Val>> {600		let Some(key) = self.keys.get(index) else {601			return Ok(None);602		};603		Ok(Some(self.obj.get_or_bail(key.clone())?))604	}605606	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {607		let key = self.keys.get(index)?;608		Some(self.obj.get_lazy_or_bail(key.clone()))609	}610611	fn get_cheap(&self, _index: usize) -> Option<Val> {612		None613	}614615	fn is_cheap(&self) -> bool {616		false617	}618}619620#[derive(Trace, Debug)]621pub struct PickObjectKeyValues {622	obj: ObjValue,623	keys: Vec<IStr>,624}625626impl PickObjectKeyValues {627	pub fn new(obj: ObjValue, keys: Vec<IStr>) -> Self {628		Self { obj, keys }629	}630}631632#[derive(Typed)]633pub struct KeyValue {634	key: IStr,635	value: Thunk<Val>,636}637638impl ArrayLike for PickObjectKeyValues {639	fn len(&self) -> usize {640		self.keys.len()641	}642643	fn get(&self, index: usize) -> Result<Option<Val>> {644		let Some(key) = self.keys.get(index) else {645			return Ok(None);646		};647		Ok(Some(648			KeyValue::into_untyped(KeyValue {649				key: key.clone(),650				value: Thunk::evaluated(self.obj.get_or_bail(key.clone())?),651			})652			.expect("convertible"),653		))654	}655656	fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {657		let key = self.keys.get(index)?;658		// Nothing can fail in the key part, yet value is still659		// lazy-evaluated660		Some(Thunk::evaluated(661			KeyValue::into_untyped(KeyValue {662				key: key.clone(),663				value: self.obj.get_lazy_or_bail(key.clone()),664			})665			.expect("convertible"),666		))667	}668669	fn get_cheap(&self, _index: usize) -> Option<Val> {670		None671	}672673	fn is_cheap(&self) -> bool {674		false675	}676}
modifiedcrates/jrsonnet-stdlib/src/arrays.rsdiffbeforeafterboth
--- a/crates/jrsonnet-stdlib/src/arrays.rs
+++ b/crates/jrsonnet-stdlib/src/arrays.rs
@@ -62,6 +62,12 @@
 }
 
 #[builtin]
+pub fn builtin_map_with_index(func: FuncVal, arr: IndexableVal) -> ArrValue {
+	let arr = arr.to_array();
+	arr.map_with_index(func)
+}
+
+#[builtin]
 pub fn builtin_flatmap(
 	func: NativeFn<((Either![String, Val],), Val)>,
 	arr: IndexableVal,
modifiedcrates/jrsonnet-stdlib/src/lib.rsdiffbeforeafterboth
--- a/crates/jrsonnet-stdlib/src/lib.rs
+++ b/crates/jrsonnet-stdlib/src/lib.rs
@@ -78,6 +78,7 @@
 		("repeat", builtin_repeat::INST),
 		("slice", builtin_slice::INST),
 		("map", builtin_map::INST),
+		("mapWithIndex", builtin_map_with_index::INST),
 		("flatMap", builtin_flatmap::INST),
 		("filter", builtin_filter::INST),
 		("foldl", builtin_foldl::INST),
modifiedcrates/jrsonnet-stdlib/src/std.jsonnetdiffbeforeafterboth
--- a/crates/jrsonnet-stdlib/src/std.jsonnet
+++ b/crates/jrsonnet-stdlib/src/std.jsonnet
@@ -3,14 +3,6 @@
 
   thisFile:: error 'std.thisFile is deprecated, to enable its support in jrsonnet - recompile it with "legacy-this-file" support.\nThis will slow down stdlib caching a bit, though',
 
-  mapWithIndex(func, arr)::
-    if !std.isFunction(func) then
-      error ('std.mapWithIndex first param must be function, got ' + std.type(func))
-    else if !std.isArray(arr) && !std.isString(arr) then
-      error ('std.mapWithIndex second param must be array, got ' + std.type(arr))
-    else
-      std.makeArray(std.length(arr), function(i) func(i, arr[i])),
-
   mapWithKey(func, obj)::
     if !std.isFunction(func) then
       error ('std.mapWithKey first param must be function, got ' + std.type(func))