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

1use std::rc::Rc;23#[cfg(feature = "exp-object-iteration")]4use jrsonnet_interner::IStr;5use jrsonnet_types::ValType;67use super::{8	destructure::{destruct, evaluate_locals_unbound, fill_letrec_binds},9	evaluate_field_member_static, evaluate_field_member_unbound,10};11use crate::{12	Context, ObjValue, ObjValueBuilder, Result, Thunk, Val,13	analyze::{14		ClosureShape, LArrComp, LBind, LCompSpec, LDestruct, LExpr, LFieldMember, LObjComp,15		LocalSlot,16	},17	arr::ArrValue,18	bail,19	error::ErrorKind::*,20	evaluate::{evaluate, evaluate_trivial},21};2223enum CachedOver {24	Arr(ArrValue),25	#[cfg(feature = "exp-object-iteration")]26	Obj(ObjValue),27}2829trait CompCollector {30	fn reserve(&mut self, _guaranteed: usize) {}31	fn collect(&mut self, ctx: Context) -> Result<()>;32}3334struct EagerArrCollector<'a> {35	out: &'a mut Vec<Val>,36	value_shape: &'a ClosureShape,37	value: &'a LExpr,38}39impl CompCollector for EagerArrCollector<'_> {40	fn reserve(&mut self, size_hint: usize) {41		self.out.reserve(size_hint);42	}43	fn collect(&mut self, ctx: Context) -> Result<()> {44		if let Some(v) = evaluate_trivial(self.value) {45			self.out.push(v);46			return Ok(());47		}48		if let LExpr::Slot(slot) = self.value {49			self.out.push(ctx.slot(*slot).evaluate()?);50			return Ok(());51		}52		let env = Context::enter_using(&ctx, self.value_shape);53		self.out.push(evaluate(env, self.value)?);54		Ok(())55	}56}5758struct LazyArrCollector<'a> {59	out: &'a mut Vec<Thunk<Val>>,60	value_shape: &'a ClosureShape,61	value: &'a Rc<LExpr>,62}63impl CompCollector for LazyArrCollector<'_> {64	fn reserve(&mut self, size_hint: usize) {65		self.out.reserve(size_hint);66	}67	fn collect(&mut self, ctx: Context) -> Result<()> {68		if let Some(v) = evaluate_trivial(self.value) {69			self.out.push(Thunk::evaluated(v));70			return Ok(());71		}72		if let LExpr::Slot(slot) = self.value.as_ref() {73			self.out.push(ctx.slot(*slot));74			return Ok(());75		}76		let env = Context::enter_using(&ctx, self.value_shape);77		let value_expr = self.value.clone();78		self.out.push(Thunk!(move || evaluate(env, &value_expr)));79		Ok(())80	}81}8283struct ObjCompCollectorStatic<'a> {84	builder: &'a mut ObjValueBuilder,85	frame_shape: &'a ClosureShape,86	locals: &'a [LBind],87	field: &'a LFieldMember,88}89impl CompCollector for ObjCompCollectorStatic<'_> {90	fn reserve(&mut self, guaranteed: usize) {91		self.builder.reserve_fields(guaranteed);92	}93	fn collect(&mut self, inner_ctx: Context) -> Result<()> {94		// Build the object's A-frame fresh per iteration: captures from95		// the comp's iter ctx, locals = `this` (slot 0, unfilled in the96		// static path) + member-locals via letrec.97		let value_ctx = inner_ctx98			.pack_captures_sup_this(self.frame_shape)99			.enter(|fill, ctx| {100				fill_letrec_binds(fill, &ctx, self.locals);101			});102		evaluate_field_member_static(self.builder, inner_ctx, value_ctx, self.field)103	}104}105106struct ObjCompCollectorUnbound<'a> {107	builder: &'a mut ObjValueBuilder,108	frame_shape: Rc<ClosureShape>,109	locals: Rc<Vec<LBind>>,110	this_slot: Option<LocalSlot>,111	field: &'a LFieldMember,112}113impl CompCollector for ObjCompCollectorUnbound<'_> {114	fn reserve(&mut self, guaranteed: usize) {115		self.builder.reserve_fields(guaranteed);116	}117	fn collect(&mut self, inner_ctx: Context) -> Result<()> {118		let uctx = evaluate_locals_unbound(119			&inner_ctx,120			&self.frame_shape,121			self.this_slot,122			self.locals.clone(),123		);124		evaluate_field_member_unbound(self.builder, inner_ctx, uctx, self.field)125	}126}127128pub fn evaluate_obj_comp(129	super_obj: Option<ObjValue>,130	ctx: Context,131	comp: &LObjComp,132) -> Result<Val> {133	let mut builder = ObjValueBuilder::new();134	if let Some(super_obj) = super_obj {135		builder.with_super(super_obj);136	}137138	let cached_overs = cache_overs(&ctx, &comp.compspecs)?;139	if comp.this.is_some() || comp.uses_super {140		evaluate_compspecs(141			ctx,142			&comp.compspecs,143			&cached_overs,144			0,145			0,146			&mut ObjCompCollectorUnbound {147				builder: &mut builder,148				frame_shape: comp.frame_shape.clone(),149				locals: comp.locals.clone(),150				this_slot: comp.this,151				field: &comp.field,152			},153		)?;154	} else {155		evaluate_compspecs(156			ctx,157			&comp.compspecs,158			&cached_overs,159			0,160			0,161			&mut ObjCompCollectorStatic {162				builder: &mut builder,163				frame_shape: &comp.frame_shape,164				locals: &comp.locals,165				field: &comp.field,166			},167		)?;168	}169170	Ok(Val::Obj(builder.build()))171}172173pub fn evaluate_arr_comp(ctx: Context, comp: &LArrComp) -> Result<Val> {174	let cached_overs = cache_overs(&ctx, &comp.compspecs)?;175176	// Eager fast-path: when the comp has only `if` and `for { destruct: Full(_) }`177	// specs, allocate one Iter A-frame per for-spec and re-set the slot178	// per iteration as long as the frame's refcount stays at 1.179	'eager: {180		let mut out = Vec::new();181182		if comp.compspecs.iter().all(|c| {183			matches!(184				c,185				LCompSpec::If(_)186					| LCompSpec::For {187						destruct: LDestruct::Full(_),188						..189					}190			)191		}) && evaluate_compspecs_eager(192			ctx.clone(),193			&comp.compspecs,194			&cached_overs,195			0,196			0,197			&mut EagerArrCollector {198				out: &mut out,199				value_shape: &comp.value_shape,200				value: &comp.value,201			},202		)203		.is_err()204		{205			break 'eager;206		}207		return Ok(Val::arr(out));208	}209210	let mut items: Vec<Thunk<Val>> = Vec::new();211	evaluate_compspecs(212		ctx,213		&comp.compspecs,214		&cached_overs,215		0,216		0,217		&mut LazyArrCollector {218			out: &mut items,219			value_shape: &comp.value_shape,220			value: &comp.value,221		},222	)?;223	Ok(Val::arr(items))224}225226fn cache_overs(ctx: &Context, specs: &[LCompSpec]) -> Result<Vec<Option<CachedOver>>> {227	specs228		.iter()229		.map(|spec| {230			Ok(match spec {231				LCompSpec::For {232					over,233					loop_invariant: true,234					..235				} => {236					let val = evaluate(ctx.clone(), over)?;237					let Val::Arr(arr) = val else {238						bail!(InComprehensionCanOnlyIterateOverArray)239					};240					Some(CachedOver::Arr(arr))241				}242				#[cfg(feature = "exp-object-iteration")]243				LCompSpec::ForObj {244					over,245					loop_invariant: true,246					..247				} => {248					let val = evaluate(ctx.clone(), over)?;249					let Val::Obj(obj) = val else {250						bail!(TypeMismatch(251							"object iteration over",252							vec![jrsonnet_types::ValType::Obj],253							val.value_type(),254						))255					};256					Some(CachedOver::Obj(obj))257				}258				_ => None,259			})260		})261		.collect::<Result<_>>()262}263264fn evaluate_compspecs_eager(265	ctx: Context,266	specs: &[LCompSpec],267	cached_overs: &[Option<CachedOver>],268	idx: usize,269	guaranteed_reserve: usize,270	collector: &mut dyn CompCollector,271) -> Result<()> {272	if idx >= specs.len() {273		collector.reserve(guaranteed_reserve);274		return collector.collect(ctx);275	}276	match &specs[idx] {277		LCompSpec::If(cond) => {278			let val = evaluate(ctx.clone(), cond)?;279			let Val::Bool(b) = val else {280				bail!(TypeMismatch(281					"if spec condition",282					vec![ValType::Bool],283					val.value_type()284				))285			};286			if b {287				evaluate_compspecs_eager(ctx, specs, cached_overs, idx + 1, 0, collector)?;288			}289		}290		LCompSpec::For {291			frame_shape,292			destruct,293			over,294			..295		} => {296			let arr = if let Some(CachedOver::Arr(cached)) = &cached_overs[idx] {297				cached.clone()298			} else {299				let arr_val = evaluate(ctx.clone(), over)?;300				let Val::Arr(arr) = arr_val else {301					bail!(InComprehensionCanOnlyIterateOverArray)302				};303				arr304			};305			let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;306			match destruct {307				LDestruct::Full(slot) => {308					Context::enter_iter(&ctx, frame_shape, |it| {309						for (i, item) in arr.iter().enumerate() {310							let item = item?;311							let ctx = it.create(|f| {312								f.set(*slot, Thunk::evaluated(item));313							})?;314							evaluate_compspecs_eager(315								ctx,316								specs,317								cached_overs,318								idx + 1,319								if i == 0 { inner_reserve } else { 0 },320								collector,321							)?;322						}323						Ok(())324					})?;325				}326				// TODO: Should not be eager? CoW won't work here327				#[cfg(feature = "exp-destruct")]328				_ => unreachable!("eager compspecs are not possible with non-full patterns"),329			}330		}331		#[cfg(feature = "exp-object-iteration")]332		LCompSpec::ForObj { .. } => {333			unreachable!("eager compspecs filter rejects ForObj");334		}335	}336	Ok(())337}338339fn evaluate_compspecs(340	ctx: Context,341	specs: &[LCompSpec],342	cached_overs: &[Option<CachedOver>],343	idx: usize,344	guaranteed_reserve: usize,345	collector: &mut dyn CompCollector,346) -> Result<()> {347	if idx >= specs.len() {348		collector.reserve(guaranteed_reserve);349		return collector.collect(ctx);350	}351	match &specs[idx] {352		LCompSpec::If(cond) => {353			let val = evaluate(ctx.clone(), cond)?;354			let Val::Bool(b) = val else {355				bail!(TypeMismatch(356					"if spec condition",357					vec![ValType::Bool],358					val.value_type()359				))360			};361			if b {362				evaluate_compspecs(ctx, specs, cached_overs, idx + 1, 0, collector)?;363			}364		}365		LCompSpec::For {366			frame_shape,367			destruct: dst,368			over,369			..370		} => {371			let arr = if let Some(CachedOver::Arr(cached)) = &cached_overs[idx] {372				cached.clone()373			} else {374				let arr_val = evaluate(ctx.clone(), over)?;375				let Val::Arr(arr) = arr_val else {376					bail!(InComprehensionCanOnlyIterateOverArray)377				};378				arr379			};380			let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;381			for (i, item) in arr.iter().enumerate() {382				let item = item?;383				let inner_ctx = ctx.pack_captures_sup_this(frame_shape).enter(|fill, ctx| {384					destruct(dst, fill, Thunk::evaluated(item), &ctx);385				});386				evaluate_compspecs(387					inner_ctx,388					specs,389					cached_overs,390					idx + 1,391					if i == 0 { inner_reserve } else { 0 },392					collector,393				)?;394			}395		}396		#[cfg(feature = "exp-object-iteration")]397		LCompSpec::ForObj {398			frame_shape,399			key,400			visibility,401			value,402			over,403			..404		} => {405			use jrsonnet_ir::Visibility;406			let obj = if let Some(CachedOver::Obj(cached)) = &cached_overs[idx] {407				cached.clone()408			} else {409				let val = evaluate(ctx.clone(), over)?;410				let Val::Obj(obj) = val else {411					bail!(TypeMismatch(412						"object iteration over",413						vec![ValType::Obj],414						val.value_type(),415					))416				};417				obj418			};419			let fields = obj.fields_with_visibility(420				#[cfg(feature = "exp-preserve-order")]421				false,422			);423			let pairs: Vec<(IStr, Visibility)> = fields424				.into_iter()425				.filter(|(_, v)| match visibility {426					Visibility::Normal => v.is_visible(),427					Visibility::Hidden => !v.is_visible(),428					Visibility::Unhide => true,429				})430				.collect();431			let inner_reserve = guaranteed_reserve.max(1) * pairs.len();432			for (i, (field_name, _)) in pairs.into_iter().enumerate() {433				let key_val = Val::string(field_name.clone());434				let value_thunk = obj435					.get_lazy(field_name.clone())436					.expect("field exists, just enumerated");437				let inner_ctx = ctx.pack_captures_sup_this(frame_shape).enter(|fill, ctx| {438					fill.set(*key, Thunk::evaluated(key_val));439					destruct(value, fill, value_thunk, &ctx);440				});441				evaluate_compspecs(442					inner_ctx,443					specs,444					cached_overs,445					idx + 1,446					if i == 0 { inner_reserve } else { 0 },447					collector,448				)?;449			}450		}451	}452	Ok(())453}