diff --git a/rsjsonnet-lang/src/program/data.rs b/rsjsonnet-lang/src/program/data.rs index b04c3d0..9137cde 100644 --- a/rsjsonnet-lang/src/program/data.rs +++ b/rsjsonnet-lang/src/program/data.rs @@ -670,6 +670,7 @@ pub(super) enum BuiltInFunc { Prune, ObjectHasEx, ObjectFieldsEx, + MapWithKey, PrimitiveEquals, Equals, Compare, diff --git a/rsjsonnet-lang/src/program/eval/call.rs b/rsjsonnet-lang/src/program/eval/call.rs index b51b6d3..47f5d6d 100644 --- a/rsjsonnet-lang/src/program/eval/call.rs +++ b/rsjsonnet-lang/src/program/eval/call.rs @@ -314,6 +314,13 @@ impl<'p> Evaluator<'_, 'p> { self.state_stack.push(State::DoThunk(arg1.view())); self.state_stack.push(State::DoThunk(arg0.view())); } + BuiltInFunc::MapWithKey => { + let [arg0, arg1] = check_num_args(args); + self.state_stack + .push(State::FnFallible(Self::do_std_map_with_key)); + self.state_stack.push(State::DoThunk(arg1.view())); + self.state_stack.push(State::DoThunk(arg0.view())); + } BuiltInFunc::PrimitiveEquals => { let [arg0, arg1] = check_num_args(args); self.state_stack diff --git a/rsjsonnet-lang/src/program/eval/stdlib.rs b/rsjsonnet-lang/src/program/eval/stdlib.rs index 26b69bc..5cca944 100644 --- a/rsjsonnet-lang/src/program/eval/stdlib.rs +++ b/rsjsonnet-lang/src/program/eval/stdlib.rs @@ -10,7 +10,7 @@ use super::{ }; use crate::gc::{Gc, GcView}; use crate::interner::InternedStr; -use crate::{ast, FHashSet}; +use crate::{ast, FHashMap, FHashSet}; impl<'p> Evaluator<'_, 'p> { pub(super) fn do_std_ext_var(&mut self) -> EvalResult<()> { @@ -235,6 +235,49 @@ impl<'p> Evaluator<'_, 'p> { Ok(()) } + pub(super) fn do_std_map_with_key(&mut self) -> EvalResult<()> { + let object = self.value_stack.pop().unwrap(); + let func = self.value_stack.pop().unwrap(); + + let func = self.expect_std_func_arg_func(func, "mapWithKey", 0)?; + let object = self.expect_std_func_arg_object(object, "mapWithKey", 1)?; + + let mut mapped_fields = FHashMap::default(); + for field_name in object.get_visible_fields_order() { + let field_thunk = self + .program + .find_object_field_thunk(&object, 0, field_name) + .unwrap(); + + let args_thunks = Box::new([ + self.program.gc_alloc(ThunkData::new_done(ValueData::String( + field_name.value().into(), + ))), + Gc::from(&field_thunk), + ]); + + let mapped_field = ObjectField { + base_env: None, + visibility: ast::Visibility::Default, + expr: None, + thunk: OnceCell::from( + self.program + .gc_alloc(ThunkData::new_pending_call(Gc::from(&func), args_thunks)), + ), + }; + + mapped_fields.insert(field_name, mapped_field); + } + + let mapped_object = ObjectData::new_simple(mapped_fields); + self.value_stack + .push(ValueData::Object(self.program.gc_alloc(mapped_object))); + + self.check_object_asserts(&object); + + Ok(()) + } + pub(super) fn do_std_primitive_equals(&mut self) -> EvalResult<()> { let rhs = self.value_stack.pop().unwrap(); let lhs = self.value_stack.pop().unwrap(); diff --git a/rsjsonnet-lang/src/program/stdlib.rs b/rsjsonnet-lang/src/program/stdlib.rs index 57a6010..6264165 100644 --- a/rsjsonnet-lang/src/program/stdlib.rs +++ b/rsjsonnet-lang/src/program/stdlib.rs @@ -85,6 +85,7 @@ impl<'p> Program<'p> { BuiltInFunc::ObjectFieldsEx, &["obj", "inc_hidden"], ); + add_simple("mapWithKey", BuiltInFunc::MapWithKey, &["func", "obj"]); add_simple("primitiveEquals", BuiltInFunc::PrimitiveEquals, &["a", "b"]); add_simple("equals", BuiltInFunc::Equals, &["a", "b"]); add_simple("__compare", BuiltInFunc::Compare, &["v1", "v2"]); diff --git a/ui-tests/fail/stdlib/mapWithKey/invalid_arg_1.jsonnet b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_1.jsonnet new file mode 100644 index 0000000..dcc245b --- /dev/null +++ b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_1.jsonnet @@ -0,0 +1 @@ +std.mapWithKey(null, {}) diff --git a/ui-tests/fail/stdlib/mapWithKey/invalid_arg_1.jsonnet.stderr b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_1.jsonnet.stderr new file mode 100644 index 0000000..0d4939f --- /dev/null +++ b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_1.jsonnet.stderr @@ -0,0 +1,8 @@ +error: first argument of `std.mapWithKey` is expected to be function, got null +note: while evaluating call to `mapWithKey` + --> invalid_arg_1.jsonnet:1:1 + | +1 | std.mapWithKey(null, {}) + | ------------------------ +note: during top-level value evaluation + diff --git a/ui-tests/fail/stdlib/mapWithKey/invalid_arg_2.jsonnet b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_2.jsonnet new file mode 100644 index 0000000..55cabe0 --- /dev/null +++ b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_2.jsonnet @@ -0,0 +1 @@ +std.mapWithKey(function(k, v) null, null) diff --git a/ui-tests/fail/stdlib/mapWithKey/invalid_arg_2.jsonnet.stderr b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_2.jsonnet.stderr new file mode 100644 index 0000000..df1407a --- /dev/null +++ b/ui-tests/fail/stdlib/mapWithKey/invalid_arg_2.jsonnet.stderr @@ -0,0 +1,8 @@ +error: second argument of `std.mapWithKey` is expected to be object, got null +note: while evaluating call to `mapWithKey` + --> invalid_arg_2.jsonnet:1:1 + | +1 | std.mapWithKey(function(k, v) null, null) + | ----------------------------------------- +note: during top-level value evaluation + diff --git a/ui-tests/fail/stdlib/mapWithKey/object_assert_failed.jsonnet b/ui-tests/fail/stdlib/mapWithKey/object_assert_failed.jsonnet new file mode 100644 index 0000000..02ad29c --- /dev/null +++ b/ui-tests/fail/stdlib/mapWithKey/object_assert_failed.jsonnet @@ -0,0 +1 @@ +std.mapWithKey(function(k, v) null, { assert false }) diff --git a/ui-tests/fail/stdlib/mapWithKey/object_assert_failed.jsonnet.stderr b/ui-tests/fail/stdlib/mapWithKey/object_assert_failed.jsonnet.stderr new file mode 100644 index 0000000..ae20004 --- /dev/null +++ b/ui-tests/fail/stdlib/mapWithKey/object_assert_failed.jsonnet.stderr @@ -0,0 +1,12 @@ +error: assertion failed + --> object_assert_failed.jsonnet:1:39 + | +1 | std.mapWithKey(function(k, v) null, { assert false }) + | ^^^^^^^^^^^^ +note: while evaluating call to `mapWithKey` + --> object_assert_failed.jsonnet:1:1 + | +1 | std.mapWithKey(function(k, v) null, { assert false }) + | ----------------------------------------------------- +note: during top-level value evaluation + diff --git a/ui-tests/pass/stdlib/mapWithKey.jsonnet b/ui-tests/pass/stdlib/mapWithKey.jsonnet new file mode 100644 index 0000000..7f64b61 --- /dev/null +++ b/ui-tests/pass/stdlib/mapWithKey.jsonnet @@ -0,0 +1,15 @@ +std.assertEqual(std.mapWithKey(function(k, v) error "err", {}), {}) && +std.assertEqual( + std.mapWithKey(function(k, v) k + ":" + v, { a: 1, b: 2 }), + { a: "a:1", b: "b:2" }, +) && +std.assertEqual( + std.mapWithKey(std.format, { "%i": 9, "%x": 10 }), + { "%i": "9", "%x": "a" }, +) && +std.assertEqual( + std.objectFieldsAll(std.mapWithKey(function(k, v) k + ":" + v, { a: 1, b: 2, c:: 3 })), + ["a", "b"], +) && + +true