From f66194917e5d44a9677062a5cc2309bc2376f475 Mon Sep 17 00:00:00 2001
From: Yaroslav Bolyukin <iam@lach.pw>
Date: Thu, 27 Jul 2023 14:22:35 +0000
Subject: [PATCH] feat: unify Arg and Typed handling for Thunk


---
--- a/crates/jrsonnet-evaluator/src/ctx.rs
+++ b/crates/jrsonnet-evaluator/src/ctx.rs
@@ -81,7 +81,7 @@
 	#[must_use]
 	pub fn into_future(self, ctx: Pending<Self>) -> Self {
 		{
-			ctx.0.borrow_mut().replace(self);
+			ctx.clone().fill(self);
 		}
 		ctx.unwrap()
 	}
--- a/crates/jrsonnet-evaluator/src/dynamic.rs
+++ b/crates/jrsonnet-evaluator/src/dynamic.rs
@@ -1,29 +1,49 @@
-use std::cell::RefCell;
+use std::cell::OnceCell;
 
 use jrsonnet_gcmodule::{Cc, Trace};
 
+use crate::{error::ErrorKind::InfiniteRecursionDetected, throw, val::ThunkValue, Result, Thunk};
+
 // TODO: Replace with OnceCell once in std
 #[derive(Clone, Trace)]
-pub struct Pending<V: Trace + 'static>(pub Cc<RefCell<Option<V>>>);
+pub struct Pending<V: Trace + 'static>(pub Cc<OnceCell<V>>);
 impl<T: Trace + 'static> Pending<T> {
 	pub fn new() -> Self {
-		Self(Cc::new(RefCell::new(None)))
+		Self(Cc::new(OnceCell::new()))
 	}
 	pub fn new_filled(v: T) -> Self {
-		Self(Cc::new(RefCell::new(Some(v))))
+		let cell = OnceCell::new();
+		let _ = cell.set(v);
+		Self(Cc::new(cell))
 	}
 	/// # Panics
 	/// If wrapper is filled already
 	pub fn fill(self, value: T) {
-		assert!(self.0.borrow().is_none(), "wrapper is filled already");
-		self.0.borrow_mut().replace(value);
+		self.0
+			.set(value)
+			.map_err(|_| ())
+			.expect("wrapper is filled already")
 	}
 }
 impl<T: Clone + Trace + 'static> Pending<T> {
 	/// # Panics
 	/// If wrapper is not yet filled
 	pub fn unwrap(&self) -> T {
-		self.0.borrow().as_ref().cloned().unwrap()
+		self.0.get().cloned().expect("pending was not filled")
+	}
+	pub fn try_get(&self) -> Option<T> {
+		self.0.get().cloned()
+	}
+}
+
+impl<T: Trace + Clone> ThunkValue for Pending<T> {
+	type Output = T;
+
+	fn get(self: Box<Self>) -> Result<Self::Output> {
+		let Some(value) = self.0.get() else {
+			throw!(InfiniteRecursionDetected);
+		};
+		Ok(value.clone())
 	}
 }
 
@@ -32,3 +52,9 @@
 		Self::new()
 	}
 }
+
+impl<T: Trace + Clone> Into<Thunk<T>> for Pending<T> {
+	fn into(self) -> Thunk<T> {
+		Thunk::new(self)
+	}
+}
--- a/crates/jrsonnet-evaluator/src/function/arglike.rs
+++ b/crates/jrsonnet-evaluator/src/function/arglike.rs
@@ -48,28 +48,22 @@
 where
 	T: Typed + Clone,
 {
-	fn evaluate_arg(&self, _ctx: Context, _tailstrict: bool) -> Result<Thunk<Val>> {
+	fn evaluate_arg(&self, _ctx: Context, tailstrict: bool) -> Result<Thunk<Val>> {
+		if T::provides_lazy() && !tailstrict {
+			return Ok(T::into_lazy_untyped(self.clone()));
+		}
 		let val = T::into_untyped(self.clone())?;
 		Ok(Thunk::evaluated(val))
 	}
 }
 impl<T> OptionalContext for T where T: Typed + Clone {}
 
-impl ArgLike for Thunk<Val> {
-	fn evaluate_arg(&self, _ctx: Context, tailstrict: bool) -> Result<Thunk<Val>> {
-		if tailstrict {
-			self.force()?;
-		}
-		Ok(self.clone())
-	}
-}
-impl OptionalContext for Thunk<Val> {}
-
 #[derive(Clone, Trace)]
 pub enum TlaArg {
 	String(IStr),
 	Code(LocExpr),
 	Val(Val),
+	Lazy(Thunk<Val>),
 }
 impl ArgLike for TlaArg {
 	fn evaluate_arg(&self, ctx: Context, tailstrict: bool) -> Result<Thunk<Val>> {
@@ -84,6 +78,7 @@
 				})
 			}),
 			TlaArg::Val(val) => Ok(Thunk::evaluated(val.clone())),
+			TlaArg::Lazy(lazy) => Ok(lazy.clone()),
 		}
 	}
 }
--- a/crates/jrsonnet-evaluator/src/obj.rs
+++ b/crates/jrsonnet-evaluator/src/obj.rs
@@ -15,7 +15,9 @@
 	function::CallLocation,
 	gc::{GcHashMap, GcHashSet, TraceBox},
 	operator::evaluate_add_op,
-	tb, throw, MaybeUnbound, Result, State, Thunk, Unbound, Val,
+	tb, throw,
+	val::ThunkValue,
+	MaybeUnbound, Result, State, Thunk, Unbound, Val,
 };
 
 #[cfg(not(feature = "exp-preserve-order"))]
--- a/crates/jrsonnet-evaluator/src/typed/conversions.rs
+++ b/crates/jrsonnet-evaluator/src/typed/conversions.rs
@@ -1,6 +1,6 @@
-use std::ops::Deref;
+use std::{collections::BTreeMap, marker::PhantomData, ops::Deref};
 
-use jrsonnet_gcmodule::Cc;
+use jrsonnet_gcmodule::{Cc, Trace};
 use jrsonnet_interner::{IBytes, IStr};
 pub use jrsonnet_macros::Typed;
 use jrsonnet_types::{ComplexValType, ValType};
@@ -11,10 +11,28 @@
 	function::{native::NativeDesc, FuncDesc, FuncVal},
 	throw,
 	typed::CheckType,
-	val::{IndexableVal, StrValue},
-	ObjValue, ObjValueBuilder, Val,
+	val::{IndexableVal, StrValue, ThunkMapper},
+	ObjValue, ObjValueBuilder, Thunk, Val,
 };
 
+#[derive(Trace)]
+struct FromUntyped<K: Trace>(PhantomData<fn() -> K>);
+impl<K> ThunkMapper<Val> for FromUntyped<K>
+where
+	K: Typed + Trace,
+{
+	type Output = K;
+
+	fn map(self, from: Val) -> Result<Self::Output> {
+		K::from_untyped(from)
+	}
+}
+impl<K: Trace> Default for FromUntyped<K> {
+	fn default() -> Self {
+		Self(PhantomData)
+	}
+}
+
 pub trait TypedObj: Typed {
 	fn serialize(self, out: &mut ObjValueBuilder) -> Result<()>;
 	fn parse(obj: &ObjValue) -> Result<Self>;
@@ -28,8 +46,24 @@
 pub trait Typed: Sized {
 	const TYPE: &'static ComplexValType;
 	fn into_untyped(typed: Self) -> Result<Val>;
+	fn into_lazy_untyped(typed: Self) -> Thunk<Val> {
+		Thunk::from(Self::into_untyped(typed))
+	}
 	fn from_untyped(untyped: Val) -> Result<Self>;
+	fn from_lazy_untyped(lazy: Thunk<Val>) -> Result<Self> {
+		Self::from_untyped(lazy.evaluate()?)
+	}
+
+	// Whatever caller should use `into_lazy_untyped` instead of `into_untyped`
+	fn provides_lazy() -> bool {
+		false
+	}
 
+	// Whatever caller should use `from_lazy_untyped` instead of `from_untyped` when possible
+	fn wants_lazy() -> bool {
+		false
+	}
+
 	/// Hack to make builtins be able to return non-result values, and make macros able to convert those values to result
 	/// This method returns identity in impl Typed for Result, and should not be overriden
 	#[doc(hidden)]
@@ -39,6 +73,54 @@
 	}
 }
 
+impl<T> Typed for Thunk<T>
+where
+	T: Typed + Trace + Clone,
+{
+	const TYPE: &'static ComplexValType = &ComplexValType::Lazy(T::TYPE);
+
+	fn into_untyped(typed: Self) -> Result<Val> {
+		T::into_untyped(typed.evaluate()?)
+	}
+
+	fn from_untyped(untyped: Val) -> Result<Self> {
+		Self::from_lazy_untyped(Thunk::evaluated(untyped))
+	}
+
+	fn provides_lazy() -> bool {
+		true
+	}
+
+	fn into_lazy_untyped(inner: Self) -> Thunk<Val> {
+		#[derive(Trace)]
+		struct IntoUntyped<K: Trace>(PhantomData<fn() -> K>);
+		impl<K> ThunkMapper<K> for IntoUntyped<K>
+		where
+			K: Typed + Trace,
+		{
+			type Output = Val;
+
+			fn map(self, from: K) -> Result<Self::Output> {
+				K::into_untyped(from)
+			}
+		}
+		impl<K: Trace> Default for IntoUntyped<K> {
+			fn default() -> Self {
+				Self(PhantomData)
+			}
+		}
+		inner.map(<IntoUntyped<T>>::default())
+	}
+
+	fn wants_lazy() -> bool {
+		true
+	}
+
+	fn from_lazy_untyped(inner: Thunk<Val>) -> Result<Self> {
+		Ok(inner.map(<FromUntyped<T>>::default()))
+	}
+}
+
 const MAX_SAFE_INTEGER: f64 = ((1u64 << (f64::MANTISSA_DIGITS + 1)) - 1) as f64;
 
 macro_rules! impl_int {
--- a/crates/jrsonnet-evaluator/src/typed/mod.rs
+++ b/crates/jrsonnet-evaluator/src/typed/mod.rs
@@ -252,6 +252,7 @@
 				}
 				Ok(())
 			}
+			Self::Lazy(_lazy) => Ok(()),
 		}
 	}
 }
--- a/crates/jrsonnet-evaluator/src/val.rs
+++ b/crates/jrsonnet-evaluator/src/val.rs
@@ -88,6 +88,54 @@
 	}
 }
 
+pub trait ThunkMapper<Input>: Trace {
+	type Output;
+	fn map(self, from: Input) -> Result<Self::Output>;
+}
+impl<Input> Thunk<Input>
+where
+	Input: Trace + Clone,
+{
+	pub fn map<M>(self, mapper: M) -> Thunk<M::Output>
+	where
+		M: ThunkMapper<Input>,
+		M::Output: Trace,
+	{
+		#[derive(Trace)]
+		struct Mapped<Input: Trace, Mapper: Trace> {
+			inner: Thunk<Input>,
+			mapper: Mapper,
+		}
+		impl<Input, Mapper> ThunkValue for Mapped<Input, Mapper>
+		where
+			Input: Trace + Clone,
+			Mapper: ThunkMapper<Input>,
+		{
+			type Output = Mapper::Output;
+
+			fn get(self: Box<Self>) -> Result<Self::Output> {
+				let value = self.inner.evaluate()?;
+				let mapped = self.mapper.map(value)?;
+				Ok(mapped)
+			}
+		}
+
+		Thunk::new(Mapped::<Input, M> {
+			inner: self,
+			mapper,
+		})
+	}
+}
+
+impl<T: Trace> From<Result<T>> for Thunk<T> {
+	fn from(value: Result<T>) -> Self {
+		match value {
+			Ok(o) => Self::evaluated(o),
+			Err(e) => Self::errored(e),
+		}
+	}
+}
+
 type CacheKey = (Option<WeakObjValue>, Option<WeakObjValue>);
 
 #[derive(Trace, Clone)]
@@ -272,6 +320,11 @@
 		Self::Flat(value.into())
 	}
 }
+impl From<IStr> for StrValue {
+	fn from(value: IStr) -> Self {
+		Self::Flat(value)
+	}
+}
 impl Display for StrValue {
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 		match self {
--- a/crates/jrsonnet-types/src/lib.rs
+++ b/crates/jrsonnet-types/src/lib.rs
@@ -128,10 +128,12 @@
 	Array(Box<ComplexValType>),
 	ArrayRef(&'static ComplexValType),
 	ObjectRef(&'static [(&'static str, &'static ComplexValType)]),
+	AttrsOf(&'static ComplexValType),
 	Union(Vec<ComplexValType>),
 	UnionRef(&'static [&'static ComplexValType]),
 	Sum(Vec<ComplexValType>),
 	SumRef(&'static [&'static ComplexValType]),
+	Lazy(&'static ComplexValType),
 }
 
 impl From<ValType> for ComplexValType {
@@ -195,10 +197,18 @@
 				}
 				write!(f, "}}")?;
 			}
+			ComplexValType::AttrsOf(a) => {
+				if matches!(a, ComplexValType::Any) {
+					write!(f, "object")?;
+				} else {
+					write!(f, "AttrsOf<{a}>")?;
+				}
+			}
 			ComplexValType::Union(v) => write_union(f, true, v.iter())?,
 			ComplexValType::UnionRef(v) => write_union(f, true, v.iter().copied())?,
 			ComplexValType::Sum(v) => write_union(f, false, v.iter())?,
 			ComplexValType::SumRef(v) => write_union(f, false, v.iter().copied())?,
+			ComplexValType::Lazy(lazy) => write!(f, "Lazy<{lazy}>")?,
 		};
 		Ok(())
 	}

-- 
gitstuff