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

difftreelog

refactor turn Thunk structure around to reduce allocations

ztvvmqptYaroslav Bolyukin2026-02-12parent: #d5d04bb.patch.diff
in: master

9 files changed

modified.gitignorediffbeforeafterboth
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,10 @@
 .vscode
 .direnv
 
+# Nix artifacts
+/result
+/result-*
+
 cache
 
 jsonnet-cpp
modifiedcrates/jrsonnet-evaluator/src/arr/spec.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/arr/spec.rs
+++ b/crates/jrsonnet-evaluator/src/arr/spec.rs
@@ -7,7 +7,7 @@
 use super::ArrValue;
 use crate::{
 	error::ErrorKind::InfiniteRecursionDetected, evaluate, function::FuncVal, typed::Typed,
-	Context, Error, ObjValue, Result, Thunk, Val,
+	val::ThunkValue, Context, Error, ObjValue, Result, Thunk, Val,
 };
 
 pub trait ArrayLike: Any + Trace + Debug {
@@ -191,9 +191,25 @@
 			ArrayThunk::Waiting(_) | ArrayThunk::Pending => {}
 		};
 
-		let arr_thunk = self.clone();
-		Some(Thunk!(move || {
-			arr_thunk.get(index).transpose().expect("index checked")
+		#[derive(Trace)]
+		struct ExprArrThunk {
+			expr: ExprArray,
+			index: usize,
+		}
+		impl ThunkValue for ExprArrThunk {
+			type Output = Val;
+
+			fn get(&self) -> Result<Self::Output> {
+				self.expr
+					.get(self.index)
+					.transpose()
+					.expect("index checked")
+			}
+		}
+
+		Some(Thunk::new(ExprArrThunk {
+			expr: self.clone(),
+			index,
 		}))
 	}
 	fn get_cheap(&self, _index: usize) -> Option<Val> {
@@ -484,9 +500,22 @@
 			ArrayThunk::Waiting(()) | ArrayThunk::Pending => {}
 		};
 
-		let arr_thunk = self.clone();
-		Some(Thunk!(move || {
-			arr_thunk.get(index).transpose().expect("index checked")
+		#[derive(Trace)]
+		struct MappedArrayThunk<const WITH_INDEX: bool> {
+			arr: MappedArray<WITH_INDEX>,
+			index: usize,
+		}
+		impl<const WITH_INDEX: bool> ThunkValue for MappedArrayThunk<WITH_INDEX> {
+			type Output = Val;
+
+			fn get(&self) -> Result<Self::Output> {
+				self.arr.get(self.index).transpose().expect("index checked")
+			}
+		}
+
+		Some(Thunk::new(MappedArrayThunk {
+			arr: self.clone(),
+			index,
 		}))
 	}
 
modifiedcrates/jrsonnet-evaluator/src/dynamic.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/dynamic.rs
+++ b/crates/jrsonnet-evaluator/src/dynamic.rs
@@ -40,8 +40,9 @@
 impl<T: Trace + Clone> ThunkValue for Pending<T> {
 	type Output = T;
 
-	fn get(self: Box<Self>) -> Result<Self::Output> {
+	fn get(&self) -> Result<Self::Output> {
 		let Some(value) = self.0.get() else {
+			// TODO: Other error?
 			bail!(InfiniteRecursionDetected);
 		};
 		Ok(value.clone())
modifiedcrates/jrsonnet-evaluator/src/obj.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/obj.rs
+++ b/crates/jrsonnet-evaluator/src/obj.rs
@@ -23,7 +23,7 @@
 	gc::WithCapacityExt as _,
 	identity_hash, in_frame,
 	operator::evaluate_add_op,
-	val::ArrValue,
+	val::{ArrValue, ThunkValue},
 	CcUnbound, MaybeUnbound, Result, Thunk, Unbound, Val,
 };
 
@@ -753,13 +753,45 @@
 		if !self.has_field_ex(key.clone(), true) {
 			return None;
 		}
-		let obj = self.clone();
+		#[derive(Trace)]
+		struct ObjFieldThunk {
+			obj: ObjValue,
+			key: IStr,
+		}
+		impl ThunkValue for ObjFieldThunk {
+			type Output = Val;
 
-		Some(Thunk!(move || Ok(obj.get(key)?.expect("field exists"))))
+			fn get(&self) -> Result<Self::Output> {
+				self.obj
+					.get(self.key.clone())
+					.transpose()
+					.expect("field existence checked")
+			}
+		}
+
+		Some(Thunk::new(ObjFieldThunk {
+			obj: self.clone(),
+			key,
+		}))
 	}
 	pub fn get_lazy_or_bail(&self, key: IStr) -> Thunk<Val> {
-		let obj = self.clone();
-		Thunk!(move || obj.get_or_bail(key))
+		#[derive(Trace)]
+		struct ObjFieldThunk {
+			obj: ObjValue,
+			key: IStr,
+		}
+		impl ThunkValue for ObjFieldThunk {
+			type Output = Val;
+
+			fn get(&self) -> Result<Self::Output> {
+				self.obj.get_or_bail(self.key.clone())
+			}
+		}
+
+		Thunk::new(ObjFieldThunk {
+			obj: self.clone(),
+			key,
+		})
 	}
 	pub fn ptr_eq(a: &Self, b: &Self) -> bool {
 		Cc::ptr_eq(&a.0, &b.0)
modifiedcrates/jrsonnet-evaluator/src/val.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/val.rs
+++ b/crates/jrsonnet-evaluator/src/val.rs
@@ -2,13 +2,14 @@
 	cell::RefCell,
 	cmp::Ordering,
 	fmt::{self, Debug, Display},
+	marker::PhantomData,
 	mem::replace,
 	num::NonZeroU32,
 	ops::Deref,
 	rc::Rc,
 };
 
-use jrsonnet_gcmodule::{Acyclic, Cc, Trace, TraceBox};
+use jrsonnet_gcmodule::{cc_dyn, Acyclic, Cc, Trace};
 use jrsonnet_interner::IStr;
 pub use jrsonnet_macros::Thunk;
 use jrsonnet_types::ValType;
@@ -28,56 +29,106 @@
 
 pub trait ThunkValue: Trace {
 	type Output;
-	fn get(self: Box<Self>) -> Result<Self::Output>;
+	fn get(&self) -> Result<Self::Output>;
 }
 
 #[derive(Trace)]
-pub struct ThunkValueClosure<D: Trace, O: 'static> {
-	env: D,
-	// Carries no data, as it is not a real closure, all the
-	// captured environment is stored in `env` field.
-	#[trace(skip)]
-	closure: fn(D) -> Result<O>,
+enum MemoizedClusureThunkInner<D: Trace, T: Trace> {
+	Computed(T),
+	Errored(Error),
+	Waiting {
+		env: D,
+		// Carries no data, as it is not a real closure, all the
+		// captured environment is stored in `env` field.
+		#[trace(skip)]
+		closure: fn(D) -> Result<T>,
+	},
+	Pending,
 }
-impl<D: Trace, O: 'static> ThunkValueClosure<D, O> {
-	pub fn new(env: D, closure: fn(D) -> Result<O>) -> Self {
-		Self { env, closure }
+#[derive(Trace)]
+pub struct MemoizedClosureThunk<D: Trace, T: Trace>(RefCell<MemoizedClusureThunkInner<D, T>>);
+impl<D: Trace, T: Trace> MemoizedClosureThunk<D, T> {
+	pub fn new(env: D, closure: fn(D) -> Result<T>) -> Self {
+		Self(RefCell::new(MemoizedClusureThunkInner::Waiting {
+			env,
+			closure,
+		}))
 	}
 }
-impl<D: Trace, O: 'static> ThunkValue for ThunkValueClosure<D, O> {
-	type Output = O;
 
-	fn get(self: Box<Self>) -> Result<Self::Output> {
-		(self.closure)(self.env)
-	}
-}
+impl<D: Trace, T: Trace + Clone> ThunkValue for MemoizedClosureThunk<D, T> {
+	type Output = T;
 
-#[derive(Trace)]
-enum ThunkInner<T: Trace> {
-	Computed(T),
-	Errored(Error),
-	Waiting(TraceBox<dyn ThunkValue<Output = T>>),
-	Pending,
+	fn get(&self) -> Result<Self::Output> {
+		match &*self.0.borrow() {
+			MemoizedClusureThunkInner::Computed(v) => return Ok(v.clone()),
+			MemoizedClusureThunkInner::Errored(e) => return Err(e.clone()),
+			MemoizedClusureThunkInner::Pending => return Err(InfiniteRecursionDetected.into()),
+			MemoizedClusureThunkInner::Waiting { .. } => (),
+		};
+		let MemoizedClusureThunkInner::Waiting { env, closure } = replace(
+			&mut *self.0.borrow_mut(),
+			MemoizedClusureThunkInner::Pending,
+		) else {
+			unreachable!();
+		};
+		let new_value = match closure(env) {
+			Ok(v) => v,
+			Err(e) => {
+				*self.0.borrow_mut() = MemoizedClusureThunkInner::Errored(e.clone());
+				return Err(e);
+			}
+		};
+		*self.0.borrow_mut() = MemoizedClusureThunkInner::Computed(new_value.clone());
+		Ok(new_value)
+	}
 }
 
-/// Lazily evaluated value
-#[allow(clippy::module_name_repetitions)]
-#[derive(Clone, Trace)]
-pub struct Thunk<T: Trace>(Cc<RefCell<ThunkInner<T>>>);
+cc_dyn!(
+	/// Lazily evaluated value
+	#[derive(Clone)] Thunk<V: Trace>,
+	ThunkValue<Output = V>,
+	pub fn new() {...}
+);
 
 impl<T: Trace> Thunk<T> {
-	pub fn evaluated(val: T) -> Self {
-		Self(Cc::new(RefCell::new(ThunkInner::Computed(val))))
-	}
-	pub fn new(f: impl ThunkValue<Output = T> + 'static) -> Self {
-		Self(Cc::new(RefCell::new(ThunkInner::Waiting(TraceBox(
-			Box::new(f),
-		)))))
+	pub fn evaluated(val: T) -> Self
+	where
+		T: Clone,
+	{
+		#[derive(Trace)]
+		struct EvaluatedThunk<T: Trace>(T);
+		impl<T> ThunkValue for EvaluatedThunk<T>
+		where
+			T: Clone + Trace,
+		{
+			type Output = T;
+
+			fn get(&self) -> Result<Self::Output> {
+				Ok(self.0.clone())
+			}
+		}
+		Self::new(EvaluatedThunk(val))
 	}
 	pub fn errored(e: Error) -> Self {
-		Self(Cc::new(RefCell::new(ThunkInner::Errored(e))))
+		#[derive(Trace)]
+		struct ErroredThunk<T: Trace>(Error, PhantomData<T>);
+		impl<T> ThunkValue for ErroredThunk<T>
+		where
+			T: Trace,
+		{
+			type Output = T;
+
+			fn get(&self) -> Result<Self::Output> {
+				Err(self.0.clone())
+			}
+		}
+		Self::new(ErroredThunk(e, PhantomData))
 	}
-	pub fn result(res: Result<T, Error>) -> Self {
+	pub fn result(res: Result<T, Error>) -> Self
+	where
+		T: Clone,
+	{
 		match res {
 			Ok(o) => Self::evaluated(o),
 			Err(e) => Self::errored(e),
@@ -101,25 +152,7 @@
 	/// - Lazy value evaluation returned error
 	/// - This method was called during inner value evaluation
 	pub fn evaluate(&self) -> Result<T> {
-		match &*self.0.borrow() {
-			ThunkInner::Computed(v) => return Ok(v.clone()),
-			ThunkInner::Errored(e) => return Err(e.clone()),
-			ThunkInner::Pending => return Err(InfiniteRecursionDetected.into()),
-			ThunkInner::Waiting(..) => (),
-		};
-		let ThunkInner::Waiting(value) = replace(&mut *self.0.borrow_mut(), ThunkInner::Pending)
-		else {
-			unreachable!();
-		};
-		let new_value = match value.0.get() {
-			Ok(v) => v,
-			Err(e) => {
-				*self.0.borrow_mut() = ThunkInner::Errored(e.clone());
-				return Err(e);
-			}
-		};
-		*self.0.borrow_mut() = ThunkInner::Computed(new_value.clone());
-		Ok(new_value)
+		self.0.get()
 	}
 }
 
@@ -134,7 +167,7 @@
 	pub fn map<M>(self, mapper: M) -> Thunk<M::Output>
 	where
 		M: ThunkMapper<Input>,
-		M::Output: Trace,
+		M::Output: Trace + Clone,
 	{
 		let inner = self;
 		Thunk!(move || {
@@ -145,7 +178,7 @@
 	}
 }
 
-impl<T: Trace> From<Result<T>> for Thunk<T> {
+impl<T: Trace + Clone> From<Result<T>> for Thunk<T> {
 	fn from(value: Result<T>) -> Self {
 		match value {
 			Ok(o) => Self::evaluated(o),
@@ -162,7 +195,7 @@
 	}
 }
 
-impl<T: Trace + Default> Default for Thunk<T> {
+impl<T: Trace + Default + Clone> Default for Thunk<T> {
 	fn default() -> Self {
 		Self::evaluated(T::default())
 	}
modifiedcrates/jrsonnet-formatter/src/comments.rsdiffbeforeafterboth
--- a/crates/jrsonnet-formatter/src/comments.rs
+++ b/crates/jrsonnet-formatter/src/comments.rs
@@ -73,7 +73,10 @@
 						p!(out, str(" "));
 					}
 					p!(out, str("/* ") string(lines[0].trim().to_string()) str(" */"));
-					if matches!(loc, CommentLocation::AboveItem | CommentLocation::EndOfItems) {
+					if matches!(
+						loc,
+						CommentLocation::AboveItem | CommentLocation::EndOfItems
+					) {
 						p!(out, nl);
 					}
 				} else if !lines.is_empty() {
modifiedcrates/jrsonnet-macros/src/lib.rsdiffbeforeafterboth
--- a/crates/jrsonnet-macros/src/lib.rs
+++ b/crates/jrsonnet-macros/src/lib.rs
@@ -879,6 +879,6 @@
 	quote! {{
 		#move_check
 		#(#trace_check)*
-		::jrsonnet_evaluator::Thunk::new(::jrsonnet_evaluator::val::ThunkValueClosure::new(#env, #closure))
+		::jrsonnet_evaluator::Thunk::new(::jrsonnet_evaluator::val::MemoizedClosureThunk::new(#env, #closure))
 	}}.into()
 }
deletedresultdiffbeforeafterboth
before · result
1/nix/store/nd6v7jksg1dqhpx4x4vqgy5ry1nkb9lk-jrsonnet-current
modifiedxtask/src/sourcegen/mod.rsdiffbeforeafterboth
--- a/xtask/src/sourcegen/mod.rs
+++ b/xtask/src/sourcegen/mod.rs
@@ -203,49 +203,57 @@
 			});
 
 			let mut type_positions: HashMap<String, usize> = HashMap::new();
-			let field_positions: Vec<_> = node.fields.iter().map(|field| {
-				let ty_str = field.ty().to_string();
-				let pos = *type_positions.get(&ty_str).unwrap_or(&0);
-				type_positions.insert(ty_str, pos + 1);
-				pos
-			}).collect();
+			let field_positions: Vec<_> = node
+				.fields
+				.iter()
+				.map(|field| {
+					let ty_str = field.ty().to_string();
+					let pos = *type_positions.get(&ty_str).unwrap_or(&0);
+					type_positions.insert(ty_str, pos + 1);
+					pos
+				})
+				.collect();
 
-			let methods = node.fields.iter().zip(field_positions.iter()).map(|(field, &pos)| {
-				let method_name = field.method_name(kinds);
-				let ty = field.ty();
+			let methods = node
+				.fields
+				.iter()
+				.zip(field_positions.iter())
+				.map(|(field, &pos)| {
+					let method_name = field.method_name(kinds);
+					let ty = field.ty();
 
-				if field.is_many() {
-					quote! {
-						pub fn #method_name(&self) -> AstChildren<#ty> {
-							support::children(&self.syntax)
+					if field.is_many() {
+						quote! {
+							pub fn #method_name(&self) -> AstChildren<#ty> {
+								support::children(&self.syntax)
+							}
 						}
-					}
-				} else if let Some(token_kind) = field.token_kind(kinds) {
-					quote! {
-						pub fn #method_name(&self) -> Option<#ty> {
-							support::token(&self.syntax, #token_kind)
+					} else if let Some(token_kind) = field.token_kind(kinds) {
+						quote! {
+							pub fn #method_name(&self) -> Option<#ty> {
+								support::token(&self.syntax, #token_kind)
+							}
 						}
-					}
-				} else if field.is_token_enum(grammar) {
-					quote! {
-						pub fn #method_name(&self) -> Option<#ty> {
-							support::token_child(&self.syntax)
+					} else if field.is_token_enum(grammar) {
+						quote! {
+							pub fn #method_name(&self) -> Option<#ty> {
+								support::token_child(&self.syntax)
+							}
 						}
-					}
-				} else if pos == 0 {
-					quote! {
-						pub fn #method_name(&self) -> Option<#ty> {
-							support::children(&self.syntax).next()
+					} else if pos == 0 {
+						quote! {
+							pub fn #method_name(&self) -> Option<#ty> {
+								support::children(&self.syntax).next()
+							}
 						}
-					}
-				} else {
-					quote! {
-						pub fn #method_name(&self) -> Option<#ty> {
-							support::children(&self.syntax).nth(#pos)
+					} else {
+						quote! {
+							pub fn #method_name(&self) -> Option<#ty> {
+								support::children(&self.syntax).nth(#pos)
+							}
 						}
 					}
-				}
-			});
+				});
 			(
 				quote! {
 					#[pretty_doc_comment_placeholder_workaround]