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 96 97 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 178 179 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 328 #[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}