1use std::rc::Rc;23use jrsonnet_types::ValType;45use super::{6 destructure::{self, evaluate_locals, evaluate_locals_unbound},7 evaluate_field_member_static, evaluate_field_member_unbound,8};9use crate::{10 Context, ContextBuilder, ObjValue, ObjValueBuilder, Pending, Result, Thunk, Val,11 analyze::{LArrComp, LBind, LCompSpec, LDestruct, LExpr, LFieldMember, LObjComp, LocalId},12 arr::ArrValue,13 bail,14 error::ErrorKind::*,15 evaluate::evaluate,16};1718trait CompCollector {19 fn reserve(&mut self, _guaranteed: usize) {}20 fn collect(&mut self, ctx: Context) -> Result<()>;21}2223struct EagerArrCollector<'a> {24 out: &'a mut Vec<Val>,25 value: &'a LExpr,26}27impl CompCollector for EagerArrCollector<'_> {28 fn reserve(&mut self, size_hint: usize) {29 self.out.reserve(size_hint);30 }31 fn collect(&mut self, ctx: Context) -> Result<()> {32 self.out.push(evaluate(ctx, self.value)?);33 Ok(())34 }35}3637struct LazyArrCollector<'a> {38 out: &'a mut Vec<Thunk<Val>>,39 value: &'a Rc<LExpr>,40}41impl CompCollector for LazyArrCollector<'_> {42 fn reserve(&mut self, size_hint: usize) {43 self.out.reserve(size_hint);44 }45 fn collect(&mut self, ctx: Context) -> Result<()> {46 let value_expr = self.value.clone();47 self.out.push(Thunk!(move || evaluate(ctx, &value_expr)));48 Ok(())49 }50}5152struct ObjCompCollectorStatic<'a> {53 builder: &'a mut ObjValueBuilder,54 locals: &'a [LBind],55 field: &'a LFieldMember,56}57impl CompCollector for ObjCompCollectorStatic<'_> {58 fn reserve(&mut self, guaranteed: usize) {59 self.builder.reserve_fields(guaranteed);60 }61 fn collect(&mut self, inner_ctx: Context) -> Result<()> {62 let value_ctx = evaluate_locals(inner_ctx.clone(), self.locals);63 evaluate_field_member_static(self.builder, inner_ctx, value_ctx, self.field)64 }65}6667struct ObjCompCollectorUnbound<'a> {68 builder: &'a mut ObjValueBuilder,69 locals: Rc<Vec<LBind>>,70 this_id: Option<LocalId>,71 field: &'a LFieldMember,72}73impl CompCollector for ObjCompCollectorUnbound<'_> {74 fn reserve(&mut self, guaranteed: usize) {75 self.builder.reserve_fields(guaranteed);76 }77 fn collect(&mut self, inner_ctx: Context) -> Result<()> {78 let uctx = evaluate_locals_unbound(inner_ctx.clone(), self.locals.clone(), self.this_id);79 evaluate_field_member_unbound(self.builder, inner_ctx, uctx, self.field)80 }81}8283pub fn evaluate_obj_comp(84 super_obj: Option<ObjValue>,85 ctx: Context,86 comp: &LObjComp,87) -> Result<Val> {88 let mut builder = ObjValueBuilder::new();89 if let Some(super_obj) = super_obj {90 builder.with_super(super_obj);91 }9293 let cached_overs = cache_overs(&ctx, &comp.compspecs)?;94 if comp.this.is_some() || comp.uses_super {95 evaluate_compspecs(96 ctx,97 &comp.compspecs,98 &cached_overs,99 0,100 0,101 &mut ObjCompCollectorUnbound {102 builder: &mut builder,103 locals: comp.locals.clone(),104 this_id: comp.this,105 field: &comp.field,106 },107 )?;108 } else {109 evaluate_compspecs(110 ctx,111 &comp.compspecs,112 &cached_overs,113 0,114 0,115 &mut ObjCompCollectorStatic {116 builder: &mut builder,117 locals: &comp.locals,118 field: &comp.field,119 },120 )?;121 }122123 Ok(Val::Obj(builder.build()))124}125126pub fn evaluate_arr_comp(ctx: Context, comp: &LArrComp) -> Result<Val> {127 let cached_overs = cache_overs(&ctx, &comp.compspecs)?;128129 130 'eager: {131 let mut out = Vec::new();132133 if comp.compspecs.iter().all(|c| {134 matches!(135 c,136 LCompSpec::If(_)137 | LCompSpec::For {138 destruct: LDestruct::Full(_),139 ..140 }141 )142 }) && evaluate_compspecs_eager(143 ctx.clone(),144 &comp.compspecs,145 &cached_overs,146 0,147 0,148 &mut EagerArrCollector {149 out: &mut out,150 value: &comp.value,151 },152 )153 .is_err()154 {155 break 'eager;156 }157 return Ok(Val::arr(out));158 }159160 let mut items: Vec<Thunk<Val>> = Vec::new();161 evaluate_compspecs(162 ctx,163 &comp.compspecs,164 &cached_overs,165 0,166 0,167 &mut LazyArrCollector {168 out: &mut items,169 value: &comp.value,170 },171 )?;172 Ok(Val::arr(items))173}174175fn cache_overs(ctx: &Context, specs: &[LCompSpec]) -> Result<Vec<Option<ArrValue>>> {176 specs177 .iter()178 .map(|spec| {179 Ok(match spec {180 LCompSpec::For {181 over,182 loop_invariant: true,183 ..184 } => {185 let val = evaluate(ctx.clone(), over)?;186 let Val::Arr(arr) = val else {187 bail!(InComprehensionCanOnlyIterateOverArray)188 };189 Some(arr)190 }191 _ => None,192 })193 })194 .collect::<Result<_>>()195}196197fn evaluate_compspecs_eager(198 ctx: Context,199 specs: &[LCompSpec],200 cached_overs: &[Option<ArrValue>],201 idx: usize,202 guaranteed_reserve: usize,203 collector: &mut dyn CompCollector,204) -> Result<()> {205 if idx >= specs.len() {206 collector.reserve(guaranteed_reserve);207 return collector.collect(ctx);208 }209 match &specs[idx] {210 LCompSpec::If(cond) => {211 let val = evaluate(ctx.clone(), cond)?;212 let Val::Bool(b) = val else {213 bail!(TypeMismatch(214 "if spec condition",215 vec![ValType::Bool],216 val.value_type()217 ))218 };219 if b {220 evaluate_compspecs_eager(ctx, specs, cached_overs, idx + 1, 0, collector)?;221 }222 }223 LCompSpec::For { destruct, over, .. } => {224 let arr = if let Some(cached) = &cached_overs[idx] {225 cached.clone()226 } else {227 let arr_val = evaluate(ctx.clone(), over)?;228 let Val::Arr(arr) = arr_val else {229 bail!(InComprehensionCanOnlyIterateOverArray)230 };231 arr232 };233 let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;234 match destruct {235 LDestruct::Full(id) => {236 let id = *id;237 let mut inner_ctx = ContextBuilder::extend(ctx, 1).build();238 for (i, item) in arr.iter().enumerate() {239 240 inner_ctx.cow_fill_binding(id, Thunk::evaluated(item?));241 evaluate_compspecs_eager(242 inner_ctx.clone(),243 specs,244 cached_overs,245 idx + 1,246 if i == 0 { inner_reserve } else { 0 },247 collector,248 )?;249 }250 }251 252 #[cfg(feature = "exp-destruct")]253 _ => unreachable!("eager compspecs are not possible with non-full patterns"),254 }255 }256 }257 Ok(())258}259260fn evaluate_compspecs(261 ctx: Context,262 specs: &[LCompSpec],263 cached_overs: &[Option<ArrValue>],264 idx: usize,265 guaranteed_reserve: usize,266 collector: &mut dyn CompCollector,267) -> Result<()> {268 if idx >= specs.len() {269 collector.reserve(guaranteed_reserve);270 return collector.collect(ctx);271 }272 match &specs[idx] {273 LCompSpec::If(cond) => {274 let val = evaluate(ctx.clone(), cond)?;275 let Val::Bool(b) = val else {276 bail!(TypeMismatch(277 "if spec condition",278 vec![ValType::Bool],279 val.value_type()280 ))281 };282 if b {283 evaluate_compspecs(ctx, specs, cached_overs, idx + 1, 0, collector)?;284 }285 }286 LCompSpec::For { destruct, over, .. } => {287 let arr = if let Some(cached) = &cached_overs[idx] {288 cached.clone()289 } else {290 let arr_val = evaluate(ctx.clone(), over)?;291 let Val::Arr(arr) = arr_val else {292 bail!(InComprehensionCanOnlyIterateOverArray)293 };294 arr295 };296 let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;297 for (i, item) in arr.iter().enumerate() {298 let item_val = item?;299 let mut inner_builder = ContextBuilder::extend(ctx.clone(), 1);300 let fctx = Pending::new();301 destructure::destruct(302 destruct,303 Thunk::evaluated(item_val),304 fctx.clone(),305 &mut inner_builder,306 );307 let inner_ctx = inner_builder.build().into_future(fctx);308 evaluate_compspecs(309 inner_ctx,310 specs,311 cached_overs,312 idx + 1,313 if i == 0 { inner_reserve } else { 0 },314 collector,315 )?;316 }317 }318 }319 Ok(())320}