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

difftreelog

source

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