git.delta.rocks / jrsonnet / refs/commits / 4c008687f967

difftreelog

perf lazy slice

Yaroslav Bolyukin2022-04-20parent: #9c0fa01.patch.diff
in: master

5 files changed

modifiedcrates/jrsonnet-evaluator/src/builtin/mod.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/builtin/mod.rs
+++ b/crates/jrsonnet-evaluator/src/builtin/mod.rs
@@ -43,33 +43,45 @@
 
 pub fn std_slice(
 	indexable: IndexableVal,
-	index: Option<usize>,
-	end: Option<usize>,
-	step: Option<usize>,
+	index: Option<BoundedUsize<0, { i32::MAX as usize }>>,
+	end: Option<BoundedUsize<0, { i32::MAX as usize }>>,
+	step: Option<BoundedUsize<1, { i32::MAX as usize }>>,
 ) -> Result<Val> {
-	let index = index.unwrap_or(0);
-	let end = end.unwrap_or_else(|| match &indexable {
-		IndexableVal::Str(_) => usize::MAX,
-		IndexableVal::Arr(v) => v.len(),
-	});
-	let step = step.unwrap_or(1);
 	match &indexable {
-		IndexableVal::Str(s) => Ok(Val::Str(
+		IndexableVal::Str(s) => {
+			let index = index.as_deref().copied().unwrap_or(0);
+			let end = end.as_deref().copied().unwrap_or(usize::MAX);
+			let step = step.as_deref().copied().unwrap_or(1);
+
+			if index >= end {
+				return Ok(Val::Str("".into()));
+			}
+
+			Ok(Val::Str(
 			(s.chars()
 				.skip(index)
 				.take(end - index)
 				.step_by(step)
 				.collect::<String>())
 			.into(),
-		)),
-		IndexableVal::Arr(arr) => Ok(Val::Arr(
-			(arr.iter()
-				.skip(index)
-				.take(end - index)
-				.step_by(step)
-				.collect::<Result<Vec<Val>>>()?)
-			.into(),
-		)),
+			))
+		}
+		IndexableVal::Arr(arr) => {
+			let index = index.as_deref().copied().unwrap_or(0);
+			let end = end.as_deref().copied().unwrap_or(usize::MAX).min(arr.len());
+			let step = step.as_deref().copied().unwrap_or(1);
+
+			if index >= end {
+				return Ok(Val::Arr(ArrValue::new_eager()));
+			}
+
+			Ok(Val::Arr(ArrValue::Slice(Box::new(Slice {
+				inner: arr.clone(),
+				from: index as u32,
+				to: end as u32,
+				step: step as u32,
+			}))))
+		}
 	}
 }
 
@@ -221,9 +233,9 @@
 #[jrsonnet_macros::builtin]
 fn builtin_slice(
 	indexable: IndexableVal,
-	index: Option<usize>,
-	end: Option<usize>,
-	step: Option<usize>,
+	index: Option<BoundedUsize<0, { i32::MAX as usize }>>,
+	end: Option<BoundedUsize<0, { i32::MAX as usize }>>,
+	step: Option<BoundedUsize<1, { i32::MAX as usize }>>,
 ) -> Result<Any> {
 	std_slice(indexable, index, end, step).map(Any)
 }
modifiedcrates/jrsonnet-evaluator/src/error.rsdiffbeforeafterboth
192 };192 };
193}193}
194
195#[macro_export]
196macro_rules! throw_runtime {
197 ($($tt:tt)*) => {
198 return Err($crate::error::Error::RuntimeError(format!($($tt)*).into()).into())
199 };
200}
194201
modifiedcrates/jrsonnet-evaluator/src/evaluate/mod.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/evaluate/mod.rs
+++ b/crates/jrsonnet-evaluator/src/evaluate/mod.rs
@@ -665,23 +665,28 @@
 		}
 		Slice(value, desc) => {
 			let indexable = evaluate(context.clone(), value)?;
+			let loc = CallLocation::new(loc);
 
-			fn parse_num(
+			fn parse_idx<const MIN: usize>(
+				loc: CallLocation,
 				context: &Context,
-				expr: Option<&LocExpr>,
+				expr: &Option<LocExpr>,
 				desc: &'static str,
-			) -> Result<Option<usize>> {
-				Ok(match expr {
-					Some(s) => evaluate(context.clone(), s)?
-						.try_cast_nullable_num(desc)?
-						.map(|v| v as usize),
-					None => None,
-				})
+			) -> Result<Option<BoundedUsize<MIN, { i32::MAX as usize }>>> {
+				if let Some(value) = expr {
+					Ok(Some(push_frame(
+						loc,
+						|| format!("slice {}", desc),
+						|| Ok(evaluate(context.clone(), value)?.try_into()?),
+					)?))
+				} else {
+					Ok(None)
+				}
 			}
 
-			let start = parse_num(&context, desc.start.as_ref(), "start")?;
-			let end = parse_num(&context, desc.end.as_ref(), "end")?;
-			let step = parse_num(&context, desc.step.as_ref(), "step")?;
+			let start = parse_idx(loc, &context, &desc.start, "start")?;
+			let end = parse_idx(loc, &context, &desc.end, "end")?;
+			let step = parse_idx(loc, &context, &desc.step, "step")?;
 
 			std_slice(indexable.into_indexable()?, start, end, step)?
 		}
modifiedcrates/jrsonnet-evaluator/src/typed/conversions.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/typed/conversions.rs
+++ b/crates/jrsonnet-evaluator/src/typed/conversions.rs
@@ -1,5 +1,6 @@
 use std::{
 	convert::{TryFrom, TryInto},
+	ops::Deref,
 	rc::Rc,
 };
 
@@ -12,7 +13,8 @@
 	error::{Error::*, LocError, Result},
 	throw,
 	typed::CheckType,
-	ArrValue, FuncDesc, FuncVal, IndexableVal, ObjValue, ObjValueBuilder, Val,
+	val::{ArrValue, FuncDesc, FuncVal, IndexableVal},
+	ObjValue, ObjValueBuilder, Val,
 };
 
 pub trait TypedObj: Typed {
@@ -69,6 +71,76 @@
 
 impl_int!(i8 u8 i16 u16 i32 u32);
 
+macro_rules! impl_bounded_int {
+	($($name:ident = $ty:ty)*) => {$(
+		#[derive(Clone, Copy)]
+		pub struct $name<const MIN: $ty, const MAX: $ty>($ty);
+		impl<const MIN: $ty, const MAX: $ty> $name<MIN, MAX> {
+			pub const fn new(value: $ty) -> Option<$name<MIN, MAX>> {
+				if value >= MIN && value <= MAX {
+					Some(Self(value))
+				} else {
+					None
+				}
+			}
+			pub const fn value(self) -> $ty {
+				self.0
+			}
+		}
+		impl<const MIN: $ty, const MAX: $ty> Deref for $name<MIN, MAX> {
+			type Target = $ty;
+			fn deref(&self) -> &Self::Target {
+				&self.0
+			}
+		}
+
+		impl<const MIN: $ty, const MAX: $ty> Typed for $name<MIN, MAX> {
+			const TYPE: &'static ComplexValType =
+				&ComplexValType::BoundedNumber(
+					Some(MIN as f64),
+					Some(MAX as f64),
+				);
+		}
+		impl<const MIN: $ty, const MAX: $ty> TryFrom<Val> for $name<MIN, MAX> {
+			type Error = LocError;
+
+			fn try_from(value: Val) -> Result<Self> {
+				<Self as Typed>::TYPE.check(&value)?;
+				match value {
+					Val::Num(n) => {
+						if n.trunc() != n {
+							throw!(RuntimeError(
+								format!(
+									"cannot convert number with fractional part to {}",
+									stringify!($ty)
+								)
+								.into()
+							))
+						}
+						Ok(Self(n as $ty))
+					}
+					_ => unreachable!(),
+				}
+			}
+		}
+		impl<const MIN: $ty, const MAX: $ty> TryFrom<$name<MIN, MAX>> for Val {
+			type Error = LocError;
+
+			fn try_from(value: $name<MIN, MAX>) -> Result<Self> {
+				Ok(Self::Num(value.0 as f64))
+			}
+		}
+	)*};
+}
+
+impl_bounded_int!(
+	BoundedI8 = i8
+	BoundedI16 = i16
+	BoundedI32 = i32
+	BoundedI64 = i64
+	BoundedUsize = usize
+);
+
 impl Typed for f64 {
 	const TYPE: &'static ComplexValType = &ComplexValType::Simple(ValType::Num);
 }
modifiedcrates/jrsonnet-evaluator/src/val.rsdiffbeforeafterboth
--- a/crates/jrsonnet-evaluator/src/val.rs
+++ b/crates/jrsonnet-evaluator/src/val.rs
@@ -172,6 +172,37 @@
 }
 
 #[derive(Debug, Clone, Trace)]
+pub struct Slice {
+	pub(crate) inner: ArrValue,
+	pub(crate) from: u32,
+	pub(crate) to: u32,
+	pub(crate) step: u32,
+}
+impl Slice {
+	fn from(&self) -> usize {
+		self.from as usize
+	}
+	fn to(&self) -> usize {
+		self.to as usize
+	}
+	fn step(&self) -> usize {
+		self.step as usize
+	}
+	fn len(&self) -> usize {
+		// TODO: use div_ceil
+		let diff = self.to() - self.from();
+		let rem = diff % self.step();
+		let div = diff / self.step();
+
+		if rem != 0 {
+			div + 1
+		} else {
+			div
+		}
+	}
+}
+
+#[derive(Debug, Clone, Trace)]
 #[force_tracking]
 pub enum ArrValue {
 	Bytes(#[skip_trace] Rc<[u8]>),
@@ -179,6 +210,7 @@
 	Eager(Cc<Vec<Val>>),
 	Extended(Box<(Self, Self)>),
 	Range(i32, i32),
+	Slice(Box<Slice>),
 	Reversed(Box<Self>),
 }
 impl ArrValue {
@@ -190,6 +222,22 @@
 		Self::Range(a, b)
 	}
 
+	pub fn slice(self, from: Option<usize>, to: Option<usize>, step: Option<usize>) -> Self {
+		let len = self.len();
+		let from = from.unwrap_or(0);
+		let to = to.unwrap_or(len).min(len);
+		let step = step.unwrap_or(1);
+		assert!(from < to);
+		assert!(step > 0);
+
+		Self::Slice(Box::new(Slice {
+			inner: self,
+			from: from as u32,
+			to: to as u32,
+			step: step as u32,
+		}))
+	}
+
 	pub fn len(&self) -> usize {
 		match self {
 			Self::Bytes(i) => i.len(),
@@ -198,6 +246,7 @@
 			Self::Extended(v) => v.0.len() + v.1.len(),
 			Self::Range(a, b) => a.abs_diff(*b) as usize,
 			Self::Reversed(i) => i.len(),
+			Self::Slice(s) => s.len(),
 		}
 	}
 
@@ -239,6 +288,13 @@
 				}
 				v.get(len - index - 1)
 			}
+			Self::Slice(s) => {
+				let index = s.from() + index * s.step();
+				if index >= s.to() {
+					return Ok(None);
+				}
+				s.inner.get(index as usize)
+			}
 		}
 	}
 
@@ -272,6 +328,13 @@
 				}
 				v.get_lazy(len - index - 1)
 			}
+			Self::Slice(s) => {
+				let index = s.from() + index * s.step();
+				if index >= s.to() {
+					return None;
+				}
+				s.inner.get_lazy(index as usize)
+			}
 		}
 	}
 
@@ -311,33 +374,43 @@
 				Cc::update_with(&mut r, |v| v.reverse());
 				r
 			}
+			Self::Slice(v) => {
+				let mut out = Vec::with_capacity(v.inner.len());
+				for v in v
+					.inner
+					.iter_lazy()
+					.skip(v.from())
+					.take(v.to() - v.from())
+					.step_by(v.step())
+				{
+					out.push(v.evaluate()?)
+				}
+				Cc::new(out)
+			}
 		})
 	}
 
 	pub fn iter(&self) -> impl DoubleEndedIterator<Item = Result<Val>> + '_ {
-		// if let Self::Reversed(v) = self {
-		// 	return v.iter().rev();
-		// }
-		let len = self.len();
-		(0..len).map(move |idx| match self {
+		(0..self.len()).map(move |idx| match self {
 			Self::Bytes(b) => Ok(Val::Num(b[idx] as f64)),
 			Self::Lazy(l) => l[idx].evaluate(),
 			Self::Eager(e) => Ok(e[idx].clone()),
 			Self::Extended(_) => self.get(idx).map(|e| e.unwrap()),
 			Self::Range(..) => self.get(idx).map(|e| e.unwrap()),
-			Self::Reversed(..) => self.get(len - idx - 1).map(|e| e.unwrap()),
+			Self::Reversed(..) => self.get(idx).map(|e| e.unwrap()),
+			Self::Slice(..) => self.get(idx).map(|e| e.unwrap()),
 		})
 	}
 
 	pub fn iter_lazy(&self) -> impl DoubleEndedIterator<Item = LazyVal> + '_ {
-		let len = self.len();
-		(0..len).map(move |idx| match self {
+		(0..self.len()).map(move |idx| match self {
 			Self::Bytes(b) => LazyVal::new_resolved(Val::Num(b[idx] as f64)),
 			Self::Lazy(l) => l[idx].clone(),
 			Self::Eager(e) => LazyVal::new_resolved(e[idx].clone()),
 			Self::Extended(_) => self.get_lazy(idx).unwrap(),
 			Self::Range(..) => self.get_lazy(idx).unwrap(),
-			Self::Reversed(..) => self.get_lazy(len - idx - 1).unwrap(),
+			Self::Reversed(..) => self.get_lazy(idx).unwrap(),
+			Self::Slice(..) => self.get_lazy(idx).unwrap(),
 		})
 	}
 
@@ -459,17 +532,6 @@
 		}
 	}
 
-	pub fn try_cast_nullable_num(self, context: &'static str) -> Result<Option<f64>> {
-		Ok(match self {
-			Val::Null => None,
-			Val::Num(num) => Some(num),
-			_ => throw!(TypeMismatch(
-				context,
-				vec![ValType::Null, ValType::Num],
-				self.value_type()
-			)),
-		})
-	}
 	pub const fn value_type(&self) -> ValType {
 		match self {
 			Self::Str(..) => ValType::Str,