1use std::{2 any::Any,3 cell::RefCell,4 fmt::{self, Debug},5 mem::replace,6 rc::Rc,7};89use jrsonnet_gcmodule::{Cc, Trace};10use jrsonnet_interner::{IBytes, IStr};11use jrsonnet_ir::Expr;1213use super::ArrValue;14use crate::{15 Context, Error, ObjValue, Result, Thunk, Val,16 error::ErrorKind::InfiniteRecursionDetected,17 evaluate,18 function::NativeFn,19 typed::{IntoUntyped, Typed},20 val::ThunkValue,21};2223pub trait ArrayLike: Any + Trace + Debug {24 fn len(&self) -> usize;25 fn is_empty(&self) -> bool {26 self.len() == 027 }28 fn get(&self, index: usize) -> Result<Option<Val>>;29 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>>;3031 fn is_cheap(&self) -> bool {32 false33 }34}35trait ArrayCheap {36 fn get(&self, index: usize) -> Option<Val>;37 fn len(&self) -> usize;38}39impl<T> ArrayLike for T40where41 T: Any + Trace + Debug + ArrayCheap,42{43 fn len(&self) -> usize {44 <T as ArrayCheap>::len(self)45 }4647 fn get(&self, index: usize) -> Result<Option<Val>> {48 Ok(<T as ArrayCheap>::get(self, index))49 }5051 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {52 <T as ArrayCheap>::get(self, index).map(Thunk::evaluated)53 }5455 fn is_cheap(&self) -> bool {56 true57 }58}5960impl ArrayCheap for () {61 fn len(&self) -> usize {62 063 }64 fn get(&self, _index: usize) -> Option<Val> {65 None66 }67}6869#[derive(Debug, Trace)]70pub struct SliceArray {71 pub(crate) inner: ArrValue,72 pub(crate) from: u32,73 pub(crate) to: u32,74 pub(crate) step: u32,75}7677impl SliceArray {78 fn map_idx(&self, index: usize) -> usize {79 self.from as usize + self.step as usize * index80 }81}82impl ArrayLike for SliceArray {83 fn len(&self) -> usize {84 (self.to - self.from).div_ceil(self.step) as usize85 }8687 fn get(&self, index: usize) -> Result<Option<Val>> {88 self.inner.get(self.map_idx(index))89 }9091 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {92 self.inner.get_lazy(self.map_idx(index))93 }9495 fn is_cheap(&self) -> bool {96 self.inner.is_cheap()97 }98}99100impl ArrayCheap for IBytes {101 fn len(&self) -> usize {102 self.as_slice().len()103 }104 fn get(&self, index: usize) -> Option<Val> {105 self.as_slice().get(index).map(|v| Val::Num((*v).into()))106 }107}108109#[derive(Debug, Trace, Clone)]110enum ArrayThunk {111 Computed(Val),112 Errored(Error),113 Waiting,114 Pending,115}116117#[derive(Debug, Trace, Clone)]118pub struct ExprArray {119 ctx: Context,120 src: Rc<Vec<Expr>>,121 cached: Cc<RefCell<Vec<ArrayThunk>>>,122}123impl ExprArray {124 pub fn new(ctx: Context, src: Rc<Vec<Expr>>) -> Self {125 Self {126 ctx,127 cached: Cc::new(RefCell::new(vec![ArrayThunk::Waiting; src.len()])),128 src,129 }130 }131}132impl ArrayLike for ExprArray {133 fn len(&self) -> usize {134 self.cached.borrow().len()135 }136 fn get(&self, index: usize) -> Result<Option<Val>> {137 if index >= self.len() {138 return Ok(None);139 }140 match &self.cached.borrow()[index] {141 ArrayThunk::Computed(c) => return Ok(Some(c.clone())),142 ArrayThunk::Errored(e) => return Err(e.clone()),143 ArrayThunk::Pending => return Err(InfiniteRecursionDetected.into()),144 ArrayThunk::Waiting => {}145 }146147 let ArrayThunk::Waiting =148 replace(&mut self.cached.borrow_mut()[index], ArrayThunk::Pending)149 else {150 unreachable!()151 };152153 let new_value = match evaluate(self.ctx.clone(), &self.src[index]) {154 Ok(v) => v,155 Err(e) => {156 self.cached.borrow_mut()[index] = ArrayThunk::Errored(e.clone());157 return Err(e);158 }159 };160 self.cached.borrow_mut()[index] = ArrayThunk::Computed(new_value.clone());161 Ok(Some(new_value))162 }163 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {164 #[derive(Trace)]165 struct ExprArrThunk {166 expr: ExprArray,167 index: usize,168 }169 impl ThunkValue for ExprArrThunk {170 type Output = Val;171172 fn get(&self) -> Result<Self::Output> {173 self.expr174 .get(self.index)175 .transpose()176 .expect("index checked")177 }178 }179180 if index >= self.len() {181 return None;182 }183 match &self.cached.borrow()[index] {184 ArrayThunk::Computed(c) => return Some(Thunk::evaluated(c.clone())),185 ArrayThunk::Errored(e) => return Some(Thunk::errored(e.clone())),186 ArrayThunk::Waiting | ArrayThunk::Pending => {}187 }188189 Some(Thunk::new(ExprArrThunk {190 expr: self.clone(),191 index,192 }))193 }194 fn is_cheap(&self) -> bool {195 false196 }197}198199#[derive(Trace, Debug)]200pub struct ExtendedArray {201 pub a: ArrValue,202 pub b: ArrValue,203 split: usize,204 len: usize,205}206impl ExtendedArray {207 pub fn new(a: ArrValue, b: ArrValue) -> Self {208 let a_len = a.len();209 let b_len = b.len();210 Self {211 a,212 b,213 split: a_len,214 len: a_len.checked_add(b_len).expect("too large array value"),215 }216 }217}218219struct WithExactSize<I>(I, usize);220impl<I, T> Iterator for WithExactSize<I>221where222 I: Iterator<Item = T>,223{224 type Item = T;225226 fn next(&mut self) -> Option<Self::Item> {227 self.0.next()228 }229 fn nth(&mut self, n: usize) -> Option<Self::Item> {230 self.0.nth(n)231 }232 fn size_hint(&self) -> (usize, Option<usize>) {233 (self.1, Some(self.1))234 }235}236impl<I> DoubleEndedIterator for WithExactSize<I>237where238 I: DoubleEndedIterator,239{240 fn next_back(&mut self) -> Option<Self::Item> {241 self.0.next_back()242 }243 fn nth_back(&mut self, n: usize) -> Option<Self::Item> {244 self.0.nth_back(n)245 }246}247impl<I> ExactSizeIterator for WithExactSize<I>248where249 I: Iterator,250{251 fn len(&self) -> usize {252 self.1253 }254}255impl ArrayLike for ExtendedArray {256 fn get(&self, index: usize) -> Result<Option<Val>> {257 if self.split > index {258 self.a.get(index)259 } else {260 self.b.get(index - self.split)261 }262 }263 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {264 if self.split > index {265 self.a.get_lazy(index)266 } else {267 self.b.get_lazy(index - self.split)268 }269 }270271 fn len(&self) -> usize {272 self.len273 }274275 fn is_cheap(&self) -> bool {276 self.a.is_cheap() && self.b.is_cheap()277 }278}279280impl<T> ArrayLike for Vec<T>281where282 T: IntoUntyped + Trace + fmt::Debug,283 for<'a> &'a T: IntoUntyped,284{285 fn len(&self) -> usize {286 self.as_slice().len()287 }288289 fn get(&self, index: usize) -> Result<Option<Val>> {290 let Some(elem) = self.as_slice().get(index) else {291 return Ok(None);292 };293 IntoUntyped::into_untyped(elem).map(Some)294 }295296 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {297 let elem = self.as_slice().get(index)?;298 Some(IntoUntyped::into_lazy_untyped(elem))299 }300301 fn is_cheap(&self) -> bool {302 !T::provides_lazy()303 }304}305306307#[derive(Debug, Trace, PartialEq, Eq)]308pub struct RangeArray {309 start: i32,310 end: i32,311}312impl RangeArray {313 pub fn empty() -> Self {314 Self::new_exclusive(0, 0)315 }316 pub fn new_exclusive(start: i32, end: i32) -> Self {317 end.checked_sub(1)318 .map_or_else(Self::empty, |end| Self { start, end })319 }320 pub fn new_inclusive(start: i32, end: i32) -> Self {321 Self { start, end }322 }323 #[expect(324 clippy::cast_sign_loss,325 reason = "the math is valid with wrapping, sign loss works as intended"326 )]327 fn size(&self) -> usize {328 (self.end as usize)329 .wrapping_sub(self.start as usize)330 .wrapping_add(1)331 }332 fn range(&self) -> impl ExactSizeIterator<Item = i32> + DoubleEndedIterator {333 WithExactSize(self.start..=self.end, self.size())334 }335}336impl ArrayCheap for RangeArray {337 fn get(&self, index: usize) -> Option<Val> {338 self.range().nth(index).map(|i| Val::Num(i.into()))339 }340 fn len(&self) -> usize {341 self.size()342 }343}344345#[derive(Debug, Trace)]346pub struct ReverseArray(pub ArrValue);347impl ArrayLike for ReverseArray {348 fn len(&self) -> usize {349 self.0.len()350 }351352 fn get(&self, index: usize) -> Result<Option<Val>> {353 self.0.get(self.0.len() - index - 1)354 }355356 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {357 self.0.get_lazy(self.0.len() - index - 1)358 }359360 fn is_cheap(&self) -> bool {361 self.0.is_cheap()362 }363}364365#[derive(Trace, Clone, Debug)]366pub enum ArrayMapper {367 Plain(NativeFn!((Val) -> Val)),368 WithIndex(NativeFn!((u32, Val) -> Val)),369}370371#[derive(Trace, Debug, Clone)]372pub struct MappedArray {373 inner: ArrValue,374 cached: Cc<RefCell<Vec<ArrayThunk>>>,375 mapper: ArrayMapper,376}377impl MappedArray {378 pub fn new(inner: ArrValue, mapper: ArrayMapper) -> Self {379 let len = inner.len();380 Self {381 inner,382 cached: Cc::new(RefCell::new(vec![ArrayThunk::Waiting; len])),383 mapper,384 }385 }386 fn evaluate(&self, index: usize, value: Val) -> Result<Val> {387 match &self.mapper {388 ArrayMapper::Plain(f) => f.call(value),389 #[expect(390 clippy::cast_possible_truncation,391 reason = "array len is limited to u31"392 )]393 ArrayMapper::WithIndex(f) => f.call(index as u32, value),394 }395 }396}397impl ArrayLike for MappedArray {398 fn len(&self) -> usize {399 self.cached.borrow().len()400 }401402 fn get(&self, index: usize) -> Result<Option<Val>> {403 if index >= self.len() {404 return Ok(None);405 }406 match &self.cached.borrow()[index] {407 ArrayThunk::Computed(c) => return Ok(Some(c.clone())),408 ArrayThunk::Errored(e) => return Err(e.clone()),409 ArrayThunk::Pending => return Err(InfiniteRecursionDetected.into()),410 ArrayThunk::Waiting => {}411 }412413 let ArrayThunk::Waiting =414 replace(&mut self.cached.borrow_mut()[index], ArrayThunk::Pending)415 else {416 unreachable!()417 };418419 let val = self420 .inner421 .get(index)422 .transpose()423 .expect("index checked")424 .and_then(|r| self.evaluate(index, r));425426 let new_value = match val {427 Ok(v) => v,428 Err(e) => {429 self.cached.borrow_mut()[index] = ArrayThunk::Errored(e.clone());430 return Err(e);431 }432 };433 self.cached.borrow_mut()[index] = ArrayThunk::Computed(new_value.clone());434 Ok(Some(new_value))435 }436 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {437 #[derive(Trace)]438 struct MappedArrayThunk {439 arr: MappedArray,440 index: usize,441 }442 impl ThunkValue for MappedArrayThunk {443 type Output = Val;444445 fn get(&self) -> Result<Self::Output> {446 self.arr.get(self.index).transpose().expect("index checked")447 }448 }449450 if index >= self.len() {451 return None;452 }453 match &self.cached.borrow()[index] {454 ArrayThunk::Computed(c) => return Some(Thunk::evaluated(c.clone())),455 ArrayThunk::Errored(e) => return Some(Thunk::errored(e.clone())),456 ArrayThunk::Waiting | ArrayThunk::Pending => {}457 }458459 Some(Thunk::new(MappedArrayThunk {460 arr: self.clone(),461 index,462 }))463 }464}465466#[derive(Trace, Debug)]467pub struct RepeatedArray {468 data: ArrValue,469 repeats: usize,470 total_len: usize,471}472impl RepeatedArray {473 pub fn new(data: ArrValue, repeats: usize) -> Option<Self> {474 let total_len = data.len().checked_mul(repeats)?;475 Some(Self {476 data,477 repeats,478 total_len,479 })480 }481 fn map_idx(&self, index: usize) -> Option<usize> {482 if index > self.total_len {483 return None;484 }485 Some(index % self.data.len())486 }487}488489impl ArrayLike for RepeatedArray {490 fn len(&self) -> usize {491 self.total_len492 }493494 fn get(&self, index: usize) -> Result<Option<Val>> {495 let Some(idx) = self.map_idx(index) else {496 return Ok(None);497 };498 self.data.get(idx)499 }500501 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {502 let idx = self.map_idx(index)?;503 self.data.get_lazy(idx)504 }505506 fn is_cheap(&self) -> bool {507 self.data.is_cheap()508 }509}510511#[derive(Trace, Debug)]512pub struct PickObjectValues {513 obj: ObjValue,514 keys: Vec<IStr>,515}516517impl PickObjectValues {518 pub fn new(obj: ObjValue, keys: Vec<IStr>) -> Self {519 Self { obj, keys }520 }521}522523impl ArrayLike for PickObjectValues {524 fn len(&self) -> usize {525 self.keys.len()526 }527528 fn get(&self, index: usize) -> Result<Option<Val>> {529 let Some(key) = self.keys.as_slice().get(index) else {530 return Ok(None);531 };532 Ok(Some(self.obj.get_or_bail(key.clone())?))533 }534535 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {536 let key = self.keys.as_slice().get(index)?;537 Some(self.obj.get_lazy_or_bail(key.clone()))538 }539540 fn is_cheap(&self) -> bool {541 false542 }543}544545#[derive(Trace, Debug)]546pub struct PickObjectKeyValues {547 obj: ObjValue,548 keys: Vec<IStr>,549}550551impl PickObjectKeyValues {552 pub fn new(obj: ObjValue, keys: Vec<IStr>) -> Self {553 Self { obj, keys }554 }555}556557#[derive(Typed, IntoUntyped)]558pub struct KeyValue {559 key: IStr,560 value: Thunk<Val>,561}562563impl ArrayLike for PickObjectKeyValues {564 fn len(&self) -> usize {565 self.keys.len()566 }567568 fn get(&self, index: usize) -> Result<Option<Val>> {569 let Some(key) = self.keys.as_slice().get(index) else {570 return Ok(None);571 };572 Ok(Some(573 KeyValue::into_untyped(KeyValue {574 key: key.clone(),575 value: Thunk::evaluated(self.obj.get_or_bail(key.clone())?),576 })577 .expect("convertible"),578 ))579 }580581 fn get_lazy(&self, index: usize) -> Option<Thunk<Val>> {582 let key = self.keys.as_slice().get(index)?;583 584 585 Some(Thunk::evaluated(586 KeyValue::into_untyped(KeyValue {587 key: key.clone(),588 value: self.obj.get_lazy_or_bail(key.clone()),589 })590 .expect("convertible"),591 ))592 }593594 fn is_cheap(&self) -> bool {595 false596 }597}