difftreelog
fix object iteration in comps
in: master
1 file changed
crates/jrsonnet-evaluator/src/evaluate/compspec.rsdiffbeforeafterboth1use 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 if !comp.compspecs.iter().all(|c| {181 matches!(182 c,183 LCompSpec::If(_)184 | LCompSpec::For {185 destruct: LDestruct::Full(_),186 ..187 }188 )189 }) {190 break 'eager;191 }192 let mut out = Vec::new();193 if evaluate_compspecs_eager(194 ctx.clone(),195 &comp.compspecs,196 &cached_overs,197 0,198 0,199 &mut EagerArrCollector {200 out: &mut out,201 value_shape: &comp.value_shape,202 value: &comp.value,203 },204 )205 .is_err()206 {207 break 'eager;208 }209 return Ok(Val::arr(out));210 }211212 let mut items: Vec<Thunk<Val>> = Vec::new();213 evaluate_compspecs(214 ctx,215 &comp.compspecs,216 &cached_overs,217 0,218 0,219 &mut LazyArrCollector {220 out: &mut items,221 value_shape: &comp.value_shape,222 value: &comp.value,223 },224 )?;225 Ok(Val::arr(items))226}227228fn cache_overs(ctx: &Context, specs: &[LCompSpec]) -> Result<Vec<Option<CachedOver>>> {229 specs230 .iter()231 .map(|spec| {232 Ok(match spec {233 LCompSpec::For {234 over,235 loop_invariant: true,236 ..237 } => {238 let val = evaluate(ctx.clone(), over)?;239 let Val::Arr(arr) = val else {240 bail!(InComprehensionCanOnlyIterateOverArray)241 };242 Some(CachedOver::Arr(arr))243 }244 #[cfg(feature = "exp-object-iteration")]245 LCompSpec::ForObj {246 over,247 loop_invariant: true,248 ..249 } => {250 let val = evaluate(ctx.clone(), over)?;251 let Val::Obj(obj) = val else {252 bail!(TypeMismatch(253 "object iteration over",254 vec![jrsonnet_types::ValType::Obj],255 val.value_type(),256 ))257 };258 Some(CachedOver::Obj(obj))259 }260 _ => None,261 })262 })263 .collect::<Result<_>>()264}265266fn evaluate_compspecs_eager(267 ctx: Context,268 specs: &[LCompSpec],269 cached_overs: &[Option<CachedOver>],270 idx: usize,271 guaranteed_reserve: usize,272 collector: &mut dyn CompCollector,273) -> Result<()> {274 if idx >= specs.len() {275 collector.reserve(guaranteed_reserve);276 return collector.collect(ctx);277 }278 match &specs[idx] {279 LCompSpec::If(cond) => {280 let val = evaluate(ctx.clone(), cond)?;281 let Val::Bool(b) = val else {282 bail!(TypeMismatch(283 "if spec condition",284 vec![ValType::Bool],285 val.value_type()286 ))287 };288 if b {289 evaluate_compspecs_eager(ctx, specs, cached_overs, idx + 1, 0, collector)?;290 }291 }292 LCompSpec::For {293 frame_shape,294 destruct,295 over,296 ..297 } => {298 let arr = if let Some(CachedOver::Arr(cached)) = &cached_overs[idx] {299 cached.clone()300 } else {301 let arr_val = evaluate(ctx.clone(), over)?;302 let Val::Arr(arr) = arr_val else {303 bail!(InComprehensionCanOnlyIterateOverArray)304 };305 arr306 };307 let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;308 match destruct {309 LDestruct::Full(slot) => {310 Context::enter_iter(&ctx, frame_shape, |it| {311 for (i, item) in arr.iter().enumerate() {312 let item = item?;313 let ctx = it.create(|f| {314 f.set(*slot, Thunk::evaluated(item));315 })?;316 evaluate_compspecs_eager(317 ctx,318 specs,319 cached_overs,320 idx + 1,321 if i == 0 { inner_reserve } else { 0 },322 collector,323 )?;324 }325 Ok(())326 })?;327 }328 // TODO: Should not be eager? CoW won't work here329 #[cfg(feature = "exp-destruct")]330 _ => unreachable!("eager compspecs are not possible with non-full patterns"),331 }332 }333 #[cfg(feature = "exp-object-iteration")]334 LCompSpec::ForObj { .. } => {335 unreachable!("eager compspecs filter rejects ForObj");336 }337 }338 Ok(())339}340341#[allow(clippy::too_many_lines)]342fn evaluate_compspecs(343 ctx: Context,344 specs: &[LCompSpec],345 cached_overs: &[Option<CachedOver>],346 idx: usize,347 guaranteed_reserve: usize,348 collector: &mut dyn CompCollector,349) -> Result<()> {350 if idx >= specs.len() {351 collector.reserve(guaranteed_reserve);352 return collector.collect(ctx);353 }354 match &specs[idx] {355 LCompSpec::If(cond) => {356 let val = evaluate(ctx.clone(), cond)?;357 let Val::Bool(b) = val else {358 bail!(TypeMismatch(359 "if spec condition",360 vec![ValType::Bool],361 val.value_type()362 ))363 };364 if b {365 evaluate_compspecs(ctx, specs, cached_overs, idx + 1, 0, collector)?;366 }367 }368 LCompSpec::For {369 frame_shape,370 destruct: dst,371 over,372 ..373 } => {374 let arr = if let Some(CachedOver::Arr(cached)) = &cached_overs[idx] {375 cached.clone()376 } else {377 let arr_val = evaluate(ctx.clone(), over)?;378 let Val::Arr(arr) = arr_val else {379 bail!(InComprehensionCanOnlyIterateOverArray)380 };381 arr382 };383 let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;384 for (i, item) in arr.iter().enumerate() {385 let item = item?;386 let inner_ctx = ctx.pack_captures_sup_this(frame_shape).enter(|fill, ctx| {387 destruct(dst, fill, Thunk::evaluated(item), ctx);388 });389 evaluate_compspecs(390 inner_ctx,391 specs,392 cached_overs,393 idx + 1,394 if i == 0 { inner_reserve } else { 0 },395 collector,396 )?;397 }398 }399 #[cfg(feature = "exp-object-iteration")]400 LCompSpec::ForObj {401 frame_shape,402 key,403 visibility,404 value,405 over,406 ..407 } => {408 use jrsonnet_ir::Visibility;409 let obj = if let Some(CachedOver::Obj(cached)) = &cached_overs[idx] {410 cached.clone()411 } else {412 let val = evaluate(ctx.clone(), over)?;413 let Val::Obj(obj) = val else {414 bail!(TypeMismatch(415 "object iteration over",416 vec![ValType::Obj],417 val.value_type(),418 ))419 };420 obj421 };422 let fields = obj.fields_with_visibility(423 #[cfg(feature = "exp-preserve-order")]424 false,425 );426 let pairs: Vec<(IStr, Visibility)> = fields427 .into_iter()428 .filter(|(_, v)| match visibility {429 Visibility::Normal => v.is_visible(),430 Visibility::Hidden => !v.is_visible(),431 Visibility::Unhide => true,432 })433 .collect();434 let inner_reserve = guaranteed_reserve.max(1) * pairs.len();435 for (i, (field_name, _)) in pairs.into_iter().enumerate() {436 let key_val = Val::string(field_name.clone());437 let value_thunk = obj438 .get_lazy(field_name.clone())439 .expect("field exists, just enumerated");440 let inner_ctx = ctx.pack_captures_sup_this(frame_shape).enter(|fill, ctx| {441 fill.set(*key, Thunk::evaluated(key_val));442 destruct(value, fill, value_thunk, &ctx);443 });444 evaluate_compspecs(445 inner_ctx,446 specs,447 cached_overs,448 idx + 1,449 if i == 0 { inner_reserve } else { 0 },450 collector,451 )?;452 }453 }454 }455 Ok(())456}