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 analyze::{LArrComp, LBind, LCompSpec, LDestruct, LExpr, LFieldMember, LObjComp, LocalId},11 arr::ArrValue,12 bail,13 error::ErrorKind::*,14 evaluate::evaluate,15 Context, ContextBuilder, ObjValue, ObjValueBuilder, Pending, Result, Thunk, Val,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 evaluate_compspecs_eager(134 ctx.clone(),135 &comp.compspecs,136 &cached_overs,137 0,138 0,139 &mut EagerArrCollector {140 out: &mut out,141 value: &comp.value,142 },143 )144 .is_err()145 {146 break 'eager;147 }148 return Ok(Val::arr(out));149 }150151 let mut items: Vec<Thunk<Val>> = Vec::new();152 evaluate_compspecs(153 ctx,154 &comp.compspecs,155 &cached_overs,156 0,157 0,158 &mut LazyArrCollector {159 out: &mut items,160 value: &comp.value,161 },162 )?;163 Ok(Val::arr(items))164}165166fn cache_overs(ctx: &Context, specs: &[LCompSpec]) -> Result<Vec<Option<ArrValue>>> {167 specs168 .iter()169 .map(|spec| {170 Ok(match spec {171 LCompSpec::For {172 over,173 loop_invariant: true,174 ..175 } => {176 let val = evaluate(ctx.clone(), over)?;177 let Val::Arr(arr) = val else {178 bail!(InComprehensionCanOnlyIterateOverArray)179 };180 Some(arr)181 }182 _ => None,183 })184 })185 .collect::<Result<_>>()186}187188fn evaluate_compspecs_eager(189 ctx: Context,190 specs: &[LCompSpec],191 cached_overs: &[Option<ArrValue>],192 idx: usize,193 guaranteed_reserve: usize,194 collector: &mut dyn CompCollector,195) -> Result<()> {196 if idx >= specs.len() {197 collector.reserve(guaranteed_reserve);198 return collector.collect(ctx.clone());199 }200 match &specs[idx] {201 LCompSpec::If(cond) => {202 let val = evaluate(ctx.clone(), cond)?;203 let Val::Bool(b) = val else {204 bail!(TypeMismatch(205 "if spec condition",206 vec![ValType::Bool],207 val.value_type()208 ))209 };210 if b {211 evaluate_compspecs_eager(ctx, specs, cached_overs, idx + 1, 0, collector)?;212 }213 }214 LCompSpec::For { destruct, over, .. } => {215 let arr = if let Some(cached) = &cached_overs[idx] {216 cached.clone()217 } else {218 let arr_val = evaluate(ctx.clone(), over)?;219 let Val::Arr(arr) = arr_val else {220 bail!(InComprehensionCanOnlyIterateOverArray)221 };222 arr223 };224 let inner_reserve = guaranteed_reserve.max(1) * arr.len() as usize;225 match destruct {226 LDestruct::Full(id) => {227 let id = *id;228 let mut inner_ctx = ContextBuilder::extend(ctx, 1).build();229 for (i, item) in arr.iter().enumerate() {230 231 inner_ctx.cow_fill_binding(id, Thunk::evaluated(item?));232 evaluate_compspecs_eager(233 inner_ctx.clone(),234 specs,235 cached_overs,236 idx + 1,237 if i == 0 { inner_reserve } else { 0 },238 collector,239 )?;240 }241 }242 #[cfg(feature = "exp-destruct")]243 _ => {244 for (i, item) in arr.iter().enumerate() {245 let item_val = item?;246 let mut inner_builder = ContextBuilder::extend(ctx.clone(), 1);247 destructure::destruct(248 destruct,249 Thunk::evaluated(item_val),250 None,251 &mut inner_builder,252 );253 let inner_ctx = inner_builder.build();254 evaluate_compspecs_eager(255 inner_ctx,256 specs,257 cached_overs,258 idx + 1,259 if i == 0 { inner_reserve } else { 0 },260 collector,261 )?;262 }263 }264 }265 }266 }267 Ok(())268}269270fn evaluate_compspecs(271 ctx: Context,272 specs: &[LCompSpec],273 cached_overs: &[Option<ArrValue>],274 idx: usize,275 guaranteed_reserve: usize,276 collector: &mut dyn CompCollector,277) -> Result<()> {278 if idx >= specs.len() {279 collector.reserve(guaranteed_reserve);280 return collector.collect(ctx);281 }282 match &specs[idx] {283 LCompSpec::If(cond) => {284 let val = evaluate(ctx.clone(), cond)?;285 let Val::Bool(b) = val else {286 bail!(TypeMismatch(287 "if spec condition",288 vec![ValType::Bool],289 val.value_type()290 ))291 };292 if b {293 evaluate_compspecs(ctx, specs, cached_overs, idx + 1, 0, collector)?;294 }295 }296 LCompSpec::For { destruct, over, .. } => {297 let arr = if let Some(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 for (i, item) in arr.iter().enumerate() {308 let item_val = item?;309 let mut inner_builder = ContextBuilder::extend(ctx.clone(), 1);310 let fctx = Pending::new();311 destructure::destruct(312 destruct,313 Thunk::evaluated(item_val),314 fctx.clone(),315 &mut inner_builder,316 );317 let inner_ctx = inner_builder.build().into_future(fctx);318 evaluate_compspecs(319 inner_ctx,320 specs,321 cached_overs,322 idx + 1,323 if i == 0 { inner_reserve } else { 0 },324 collector,325 )?;326 }327 }328 }329 Ok(())330}