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

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};9use crate::{10	Context, ObjValue, ObjValueBuilder, Result, Thunk, Val,11	analyze::{12		ClosureShape, LArrComp, LBind, LCompSpec, LDestruct, LExpr, LFieldMember, LObjComp,13		LocalSlot,14	},15	arr::ArrValue,16	bail,17	error::ErrorKind::*,18	evaluate::{evaluate, evaluate_trivial},19};2021trait CompCollector {22	fn reserve(&mut self, _guaranteed: usize) {}23	fn collect(&mut self, ctx: Context) -> Result<()>;24}2526struct EagerArrCollector<'a> {27	out: &'a mut Vec<Val>,28	value_shape: &'a ClosureShape,29	value: &'a LExpr,30}31impl CompCollector for EagerArrCollector<'_> {32	fn reserve(&mut self, size_hint: usize) {33		self.out.reserve(size_hint);34	}35	fn collect(&mut self, ctx: Context) -> Result<()> {36		if let Some(v) = evaluate_trivial(self.value) {37			self.out.push(v);38			return Ok(());39		}40		if let LExpr::Slot(slot) = self.value {41			self.out.push(ctx.slot(*slot).evaluate()?);42			return Ok(());43		}44		let env = Context::enter_using(&ctx, self.value_shape);45		self.out.push(evaluate(env, self.value)?);46		Ok(())47	}48}4950struct LazyArrCollector<'a> {51	out: &'a mut Vec<Thunk<Val>>,52	value_shape: &'a ClosureShape,53	value: &'a Rc<LExpr>,54}55impl CompCollector for LazyArrCollector<'_> {56	fn reserve(&mut self, size_hint: usize) {57		self.out.reserve(size_hint);58	}59	fn collect(&mut self, ctx: Context) -> Result<()> {60		if let Some(v) = evaluate_trivial(self.value) {61			self.out.push(Thunk::evaluated(v));62			return Ok(());63		}64		if let LExpr::Slot(slot) = self.value.as_ref() {65			self.out.push(ctx.slot(*slot));66			return Ok(());67		}68		let env = Context::enter_using(&ctx, self.value_shape);69		let value_expr = self.value.clone();70		self.out.push(Thunk!(move || evaluate(env, &value_expr)));71		Ok(())72	}73}7475struct ObjCompCollectorStatic<'a> {76	builder: &'a mut ObjValueBuilder,77	frame_shape: &'a ClosureShape,78	locals: &'a [LBind],79	field: &'a LFieldMember,80}81impl CompCollector for ObjCompCollectorStatic<'_> {82	fn reserve(&mut self, guaranteed: usize) {83		self.builder.reserve_fields(guaranteed);84	}85	fn collect(&mut self, inner_ctx: Context) -> Result<()> {86		// Build the object's A-frame fresh per iteration: captures from87		// the comp's iter ctx, locals = `this` (slot 0, unfilled in the88		// static path) + member-locals via letrec.89		let value_ctx = inner_ctx90			.pack_captures_sup_this(self.frame_shape)91			.enter(|fill, ctx| {92				fill_letrec_binds(fill, &ctx, self.locals);93			});94		evaluate_field_member_static(self.builder, inner_ctx, value_ctx, self.field)95	}96}9798struct ObjCompCollectorUnbound<'a> {99	builder: &'a mut ObjValueBuilder,100	frame_shape: Rc<ClosureShape>,101	locals: Rc<Vec<LBind>>,102	this_slot: Option<LocalSlot>,103	field: &'a LFieldMember,104}105impl CompCollector for ObjCompCollectorUnbound<'_> {106	fn reserve(&mut self, guaranteed: usize) {107		self.builder.reserve_fields(guaranteed);108	}109	fn collect(&mut self, inner_ctx: Context) -> Result<()> {110		let uctx = evaluate_locals_unbound(111			&inner_ctx,112			&self.frame_shape,113			self.this_slot,114			self.locals.clone(),115		);116		evaluate_field_member_unbound(self.builder, inner_ctx, uctx, self.field)117	}118}119120pub fn evaluate_obj_comp(121	super_obj: Option<ObjValue>,122	ctx: Context,123	comp: &LObjComp,124) -> Result<Val> {125	let mut builder = ObjValueBuilder::new();126	if let Some(super_obj) = super_obj {127		builder.with_super(super_obj);128	}129130	let cached_overs = cache_overs(&ctx, &comp.compspecs)?;131	if comp.this.is_some() || comp.uses_super {132		evaluate_compspecs(133			ctx,134			&comp.compspecs,135			&cached_overs,136			0,137			0,138			&mut ObjCompCollectorUnbound {139				builder: &mut builder,140				frame_shape: comp.frame_shape.clone(),141				locals: comp.locals.clone(),142				this_slot: comp.this,143				field: &comp.field,144			},145		)?;146	} else {147		evaluate_compspecs(148			ctx,149			&comp.compspecs,150			&cached_overs,151			0,152			0,153			&mut ObjCompCollectorStatic {154				builder: &mut builder,155				frame_shape: &comp.frame_shape,156				locals: &comp.locals,157				field: &comp.field,158			},159		)?;160	}161162	Ok(Val::Obj(builder.build()))163}164165pub fn evaluate_arr_comp(ctx: Context, comp: &LArrComp) -> Result<Val> {166	let cached_overs = cache_overs(&ctx, &comp.compspecs)?;167168	// Eager fast-path: when the comp has only `if` and `for { destruct: Full(_) }`169	// specs, allocate one Iter A-frame per for-spec and re-set the slot170	// per iteration as long as the frame's refcount stays at 1.171	'eager: {172		let mut out = Vec::new();173174		if comp.compspecs.iter().all(|c| {175			matches!(176				c,177				LCompSpec::If(_)178					| LCompSpec::For {179						destruct: LDestruct::Full(_),180						..181					}182			)183		}) && evaluate_compspecs_eager(184			ctx.clone(),185			&comp.compspecs,186			&cached_overs,187			0,188			0,189			&mut EagerArrCollector {190				out: &mut out,191				value_shape: &comp.value_shape,192				value: &comp.value,193			},194		)195		.is_err()196		{197			break 'eager;198		}199		return Ok(Val::arr(out));200	}201202	let mut items: Vec<Thunk<Val>> = Vec::new();203	evaluate_compspecs(204		ctx,205		&comp.compspecs,206		&cached_overs,207		0,208		0,209		&mut LazyArrCollector {210			out: &mut items,211			value_shape: &comp.value_shape,212			value: &comp.value,213		},214	)?;215	Ok(Val::arr(items))216}217218fn cache_overs(ctx: &Context, specs: &[LCompSpec]) -> Result<Vec<Option<ArrValue>>> {219	specs220		.iter()221		.map(|spec| {222			Ok(match spec {223				LCompSpec::For {224					over,225					loop_invariant: true,226					..227				} => {228					let val = evaluate(ctx.clone(), over)?;229					let Val::Arr(arr) = val else {230						bail!(InComprehensionCanOnlyIterateOverArray)231					};232					Some(arr)233				}234				_ => None,235			})236		})237		.collect::<Result<_>>()238}239240fn evaluate_compspecs_eager(241	ctx: Context,242	specs: &[LCompSpec],243	cached_overs: &[Option<ArrValue>],244	idx: usize,245	guaranteed_reserve: usize,246	collector: &mut dyn CompCollector,247) -> Result<()> {248	if idx >= specs.len() {249		collector.reserve(guaranteed_reserve);250		return collector.collect(ctx);251	}252	match &specs[idx] {253		LCompSpec::If(cond) => {254			let val = evaluate(ctx.clone(), cond)?;255			let Val::Bool(b) = val else {256				bail!(TypeMismatch(257					"if spec condition",258					vec![ValType::Bool],259					val.value_type()260				))261			};262			if b {263				evaluate_compspecs_eager(ctx, specs, cached_overs, idx + 1, 0, collector)?;264			}265		}266		LCompSpec::For {267			frame_shape,268			destruct,269			over,270			..271		} => {272			let arr = if let Some(cached) = &cached_overs[idx] {273				cached.clone()274			} else {275				let arr_val = evaluate(ctx.clone(), over)?;276				let Val::Arr(arr) = arr_val else {277					bail!(InComprehensionCanOnlyIterateOverArray)278				};279				arr280			};281			let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;282			match destruct {283				LDestruct::Full(slot) => {284					Context::enter_iter(&ctx, frame_shape, |it| {285						for (i, item) in arr.iter().enumerate() {286							let item = item?;287							let ctx = it.create(|f| {288								f.set(*slot, Thunk::evaluated(item));289							})?;290							evaluate_compspecs_eager(291								ctx,292								specs,293								cached_overs,294								idx + 1,295								if i == 0 { inner_reserve } else { 0 },296								collector,297							)?;298						}299						Ok(())300					})?;301				}302				// TODO: Should not be eager? CoW won't work here303				#[cfg(feature = "exp-destruct")]304				_ => unreachable!("eager compspecs are not possible with non-full patterns"),305			}306		}307	}308	Ok(())309}310311fn evaluate_compspecs(312	ctx: Context,313	specs: &[LCompSpec],314	cached_overs: &[Option<ArrValue>],315	idx: usize,316	guaranteed_reserve: usize,317	collector: &mut dyn CompCollector,318) -> Result<()> {319	if idx >= specs.len() {320		collector.reserve(guaranteed_reserve);321		return collector.collect(ctx);322	}323	match &specs[idx] {324		LCompSpec::If(cond) => {325			let val = evaluate(ctx.clone(), cond)?;326			let Val::Bool(b) = val else {327				bail!(TypeMismatch(328					"if spec condition",329					vec![ValType::Bool],330					val.value_type()331				))332			};333			if b {334				evaluate_compspecs(ctx, specs, cached_overs, idx + 1, 0, collector)?;335			}336		}337		LCompSpec::For {338			frame_shape,339			destruct: dst,340			over,341			..342		} => {343			let arr = if let Some(cached) = &cached_overs[idx] {344				cached.clone()345			} else {346				let arr_val = evaluate(ctx.clone(), over)?;347				let Val::Arr(arr) = arr_val else {348					bail!(InComprehensionCanOnlyIterateOverArray)349				};350				arr351			};352			let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;353			for (i, item) in arr.iter().enumerate() {354				let item = item?;355				let inner_ctx = ctx.pack_captures_sup_this(frame_shape).enter(|fill, ctx| {356					destruct(dst, fill, Thunk::evaluated(item), &ctx);357				});358				evaluate_compspecs(359					inner_ctx,360					specs,361					cached_overs,362					idx + 1,363					if i == 0 { inner_reserve } else { 0 },364					collector,365				)?;366			}367		}368	}369	Ok(())370}