change(eval): implement std.mapWithKey as built-in

rsjsonnet-lang/src/program/data.rs

@@ -670,6 +670,7 @@ pub(super) enum BuiltInFunc {
     Prune,
     ObjectHasEx,
     ObjectFieldsEx,
+    MapWithKey,
     PrimitiveEquals,
     Equals,
     Compare,

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

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();

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"]);

ui-tests/fail/stdlib/mapWithKey/invalid_arg_1.jsonnet

@@ -0,0 +1 @@
+std.mapWithKey(null, {})

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
+

ui-tests/fail/stdlib/mapWithKey/invalid_arg_2.jsonnet

@@ -0,0 +1 @@
+std.mapWithKey(function(k, v) null, null)

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
+

ui-tests/fail/stdlib/mapWithKey/object_assert_failed.jsonnet

@@ -0,0 +1 @@
+std.mapWithKey(function(k, v) null, { assert false })

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
+

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