git.delta.rocks / jrsonnet / refs/commits / b82846440add

--- a/crates/jrsonnet-evaluator/src/arr/mod.rs
+++ b/crates/jrsonnet-evaluator/src/arr/mod.rs
@@ -9,7 +9,7 @@
 use jrsonnet_interner::IBytes;
 use jrsonnet_ir::Expr;
 
-use crate::{function::NativeFn, Context, Result, Thunk, Val};
+use crate::{function::NativeFn, typed::IntoUntyped, Context, Result, Thunk, Val};
 
 mod spec;
 pub use spec::{ArrayLike, *};
@@ -71,16 +71,28 @@
 		Self::new(<MappedArray>::new(self, ArrayMapper::WithIndex(mapper)))
 	}
 
-	pub fn filter(self, filter: impl Fn(&Val) -> Result<bool>) -> Result<Self> {
+	pub fn filter(self, filter: NativeFn!((Thunk<Val>) -> bool)) -> Result<Self> {
 		// TODO: ArrValue::Picked(inner, indexes) for large arrays
+		'eager: {
+			let mut out = Vec::new();
+			for i in self.iter() {
+				let Ok(i) = i else {
+					break 'eager;
+				};
+				if filter.call(IntoUntyped::into_lazy_untyped(i.clone()))? {
+					out.push(i);
+				}
+			}
+			return Ok(Self::eager(out));
+		};
+
 		let mut out = Vec::new();
-		for i in self.iter() {
-			let i = i?;
-			if filter(&i)? {
+		for i in self.iter_lazy() {
+			if filter.call(i.clone())? {
 				out.push(i);
 			}
 		}
-		Ok(Self::eager(out))
+		Ok(Self::lazy(out))
 	}
 
 	pub fn extended(a: Self, b: Self) -> Self {

126	}126	}
127}127}
128128
129type FilterFunc = NativeFn!((Val) -> bool);129type FilterFunc = NativeFn!((Thunk<Val>) -> bool);
130130
131#[builtin]131#[builtin]
132pub fn builtin_filter(func: FilterFunc, arr: ArrValue) -> Result<ArrValue> {132pub fn builtin_filter(func: FilterFunc, arr: ArrValue) -> Result<ArrValue> {
133	arr.filter(|val| func.call(val.clone()))133	arr.filter(func)
134}134}
135135
136#[builtin]136#[builtin]
139	map_func: NativeFn!((Val) -> Val),139	map_func: NativeFn!((Val) -> Val),
140	arr: ArrValue,140	arr: ArrValue,
141) -> Result<ArrValue> {141) -> Result<ArrValue> {
142	Ok(builtin_filter(filter_func, arr)?.map(map_func))142	Ok(arr.filter(filter_func)?.map(map_func))
143}143}
144144
145#[builtin]145#[builtin]