change(eval): implement std.mergePatch as a built-in function

eduardosm · Nov 3, 2024 · b1077c9 · b1077c9
1 parent b736e51
commit b1077c9
Show file tree

Hide file tree

Showing 12 changed files with 167 additions and 39 deletions.
diff --git a/rsjsonnet-lang/src/program/data.rs b/rsjsonnet-lang/src/program/data.rs
@@ -421,6 +421,23 @@ impl GcTrace for ObjectData<'_> {
 }
 
 impl<'p> ObjectData<'p> {
+    #[inline]
+    pub(super) fn new_empty() -> Self {
+        Self {
+            self_layer: ObjectLayer {
+                is_top: false,
+                locals: &[],
+                base_env: None,
+                env: OnceCell::new(),
+                fields: FHashMap::default(),
+                asserts: &[],
+            },
+            super_layers: Vec::new(),
+            fields_order: OnceCell::new(),
+            asserts_checked: Cell::new(true),
+        }
+    }
+
     #[inline]
     pub(super) fn new_simple(fields: FHashMap<InternedStr<'p>, ObjectField<'p>>) -> Self {
         Self {
@@ -503,7 +520,7 @@ impl<'p> ObjectData<'p> {
     #[inline]
     pub(super) fn get_visible_fields_order(
         &self,
-    ) -> impl DoubleEndedIterator<Item = InternedStr<'p>> + '_ {
+    ) -> impl DoubleEndedIterator<Item = InternedStr<'p>> + Clone + '_ {
         self.get_fields_order()
             .iter()
             .filter_map(|&(name, visible)| visible.then_some(name))
@@ -742,12 +759,14 @@ pub(super) enum BuiltInFunc {
     Base64DecodeBytes,
     Base64Decode,
     Md5,
+    // JSON Merge Patch
+    MergePatch,
+    // Other
+    Mod,
     // Native Functions
     Native,
     // Debugging
     Trace,
-    // Other
-    Mod,
 }
 
 pub(super) struct ThunkEnv<'p> {

diff --git a/rsjsonnet-lang/src/program/eval/call.rs b/rsjsonnet-lang/src/program/eval/call.rs
@@ -767,6 +767,18 @@ impl<'p> Evaluator<'_, 'p> {
                 self.state_stack.push(State::StdMd5);
                 self.state_stack.push(State::DoThunk(arg.view()));
             }
+            BuiltInFunc::MergePatch => {
+                let [arg0, arg1] = check_num_args(args);
+                self.state_stack.push(State::StdMergePatchValue);
+                self.state_stack.push(State::DoThunk(arg1.view()));
+                self.state_stack.push(State::DoThunk(arg0.view()));
+            }
+            BuiltInFunc::Mod => {
+                let [arg0, arg1] = check_num_args(args);
+                self.state_stack.push(State::StdMod);
+                self.state_stack.push(State::DoThunk(arg1.view()));
+                self.state_stack.push(State::DoThunk(arg0.view()));
+            }
             BuiltInFunc::Native => {
                 let [arg] = check_num_args(args);
                 self.state_stack.push(State::StdNative);
@@ -778,12 +790,6 @@ impl<'p> Evaluator<'_, 'p> {
                 self.state_stack.push(State::DoThunk(arg0.view()));
                 self.state_stack.push(State::DoThunk(arg1.view()));
             }
-            BuiltInFunc::Mod => {
-                let [arg0, arg1] = check_num_args(args);
-                self.state_stack.push(State::StdMod);
-                self.state_stack.push(State::DoThunk(arg1.view()));
-                self.state_stack.push(State::DoThunk(arg0.view()));
-            }
         }
     }
 

diff --git a/rsjsonnet-lang/src/program/eval/mod.rs b/rsjsonnet-lang/src/program/eval/mod.rs
@@ -1551,9 +1551,11 @@ impl<'p, 'a> Evaluator<'a, 'p> {
                 State::StdBase64DecodeBytes => self.do_std_base64_decode_bytes()?,
                 State::StdBase64Decode => self.do_std_base64_decode()?,
                 State::StdMd5 => self.do_std_md5()?,
+                State::StdMergePatchValue => self.do_std_merge_patch_value(),
+                State::StdMergePatchField { name } => self.do_std_merge_patch_field(name),
+                State::StdMod => self.do_std_mod()?,
                 State::StdNative => self.do_std_native()?,
                 State::StdTrace => self.do_std_trace()?,
-                State::StdMod => self.do_std_mod()?,
             }
 
             if self.stack_trace_len > self.program.max_stack {

diff --git a/rsjsonnet-lang/src/program/eval/state.rs b/rsjsonnet-lang/src/program/eval/state.rs
@@ -521,7 +521,11 @@ pub(super) enum State<'p> {
     StdBase64DecodeBytes,
     StdBase64Decode,
     StdMd5,
+    StdMergePatchValue,
+    StdMergePatchField {
+        name: InternedStr<'p>,
+    },
+    StdMod,
     StdNative,
     StdTrace,
-    StdMod,
 }
diff --git 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, FHashMap};
+use crate::{ast, FHashSet};
 
 impl<'p> Evaluator<'_, 'p> {
     pub(super) fn do_std_ext_var(&mut self) -> EvalResult<()> {
@@ -127,8 +127,7 @@ impl<'p> Evaluator<'_, 'p> {
                 let object = object.view();
                 let visible_fields = object.get_visible_fields_order();
 
-                self.object_stack
-                    .push(ObjectData::new_simple(FHashMap::default()));
+                self.object_stack.push(ObjectData::new_empty());
                 self.state_stack.push(State::ObjectToValue);
 
                 for field_name in visible_fields.rev() {
@@ -2911,6 +2910,94 @@ impl<'p> Evaluator<'_, 'p> {
         Ok(())
     }
 
+    pub(super) fn do_std_merge_patch_value(&mut self) {
+        let patch = self.value_stack.pop().unwrap();
+        let target = self.value_stack.pop().unwrap();
+
+        if let ValueData::Object(patch) = patch {
+            let patch = patch.view();
+            let target = if let ValueData::Object(target) = target {
+                Some(target.view())
+            } else {
+                None
+            };
+
+            let patch_fields_order = patch.get_visible_fields_order();
+            let patch_fields: FHashSet<_> = patch_fields_order.clone().collect();
+            let mut target_fields = FHashSet::default();
+
+            let mut new_object = ObjectData::new_empty();
+            if let Some(ref target) = target {
+                target_fields = target.get_visible_fields_order().collect();
+                for &field_name in target_fields.iter() {
+                    if !patch_fields.contains(&field_name) {
+                        let field_thunk = self
+                            .program
+                            .find_object_field_thunk(target, 0, field_name)
+                            .unwrap();
+                        new_object.self_layer.fields.insert(
+                            field_name,
+                            ObjectField {
+                                base_env: None,
+                                visibility: ast::Visibility::Default,
+                                expr: None,
+                                thunk: OnceCell::from(Gc::from(&field_thunk)),
+                            },
+                        );
+                    }
+                }
+            }
+
+            self.state_stack.push(State::ObjectToValue);
+            self.object_stack.push(new_object);
+
+            for field_name in patch_fields_order.rev() {
+                let patch_field_thunk = self
+                    .program
+                    .find_object_field_thunk(&patch, 0, field_name)
+                    .unwrap();
+
+                self.state_stack
+                    .push(State::StdMergePatchField { name: field_name });
+                self.state_stack.push(State::StdMergePatchValue);
+                self.state_stack.push(State::DoThunk(patch_field_thunk));
+                if target_fields.contains(&field_name) {
+                    let target_field_thunk = self
+                        .program
+                        .find_object_field_thunk(target.as_ref().unwrap(), 0, field_name)
+                        .unwrap();
+
+                    self.state_stack.push(State::DoThunk(target_field_thunk));
+                } else {
+                    self.value_stack.push(ValueData::Null);
+                }
+            }
+
+            self.check_object_asserts(&patch);
+            if let Some(ref target) = target {
+                self.check_object_asserts(target);
+            }
+        } else {
+            self.value_stack.push(patch);
+        }
+    }
+
+    pub(super) fn do_std_merge_patch_field(&mut self, name: InternedStr<'p>) {
+        let value = self.value_stack.pop().unwrap();
+        let object = self.object_stack.last_mut().unwrap();
+        if !matches!(value, ValueData::Null) {
+            object.self_layer.fields.insert(
+                name,
+                ObjectField {
+                    base_env: None,
+                    visibility: ast::Visibility::Default,
+                    expr: None,
+                    thunk: OnceCell::from(self.program.gc_alloc(ThunkData::new_done(value))),
+                },
+            );
+        }
+    }
+
     pub(super) fn do_std_mod(&mut self) -> EvalResult<()> {
         let rhs = self.value_stack.pop().unwrap();
         let lhs = self.value_stack.pop().unwrap();

diff --git a/rsjsonnet-lang/src/program/std.jsonnet b/rsjsonnet-lang/src/program/std.jsonnet
@@ -129,30 +129,6 @@ limitations under the License.
 
   manifestJsonMinified(value):: std.manifestJsonEx(value, '', '', ':'),
 
-  mergePatch(target, patch)::
-    if std.isObject(patch) then
-      local target_object =
-        if std.isObject(target) then target else {};
-
-      local target_fields =
-        if std.isObject(target_object) then std.objectFields(target_object) else [];
-
-      local null_fields = [k for k in std.objectFields(patch) if patch[k] == null];
-      local both_fields = std.setUnion(target_fields, std.objectFields(patch));
-
-      {
-        [k]:
-          if !std.objectHas(patch, k) then
-            target_object[k]
-          else if !std.objectHas(target_object, k) then
-            std.mergePatch(null, patch[k]) tailstrict
-          else
-            std.mergePatch(target_object[k], patch[k]) tailstrict
-        for k in std.setDiff(both_fields, null_fields)
-      }
-    else
-      patch,
-
   get(o, f, default=null, inc_hidden=true)::
     if std.objectHasEx(o, f, inc_hidden) then o[f] else default,
 

diff --git a/rsjsonnet-lang/src/program/stdlib.rs b/rsjsonnet-lang/src/program/stdlib.rs
@@ -200,9 +200,10 @@ impl<'p> Program<'p> {
         );
         add_simple("base64Decode", BuiltInFunc::Base64Decode, &["str"]);
         add_simple("md5", BuiltInFunc::Md5, &["str"]);
+        add_simple("mergePatch", BuiltInFunc::MergePatch, &["target", "patch"]);
+        add_simple("mod", BuiltInFunc::Mod, &["a", "b"]);
         add_simple("native", BuiltInFunc::Native, &["name"]);
         add_simple("trace", BuiltInFunc::Trace, &["str", "rest"]);
-        add_simple("mod", BuiltInFunc::Mod, &["a", "b"]);
 
         let mut add_with_defaults =
             |name: &str, kind: BuiltInFunc, params: &[(&str, Option<&'p ir::Expr<'p>>)]| {

diff --git a/ui-tests/fail/stdlib/mergePatch/object_assert_failed_1.jsonnet b/ui-tests/fail/stdlib/mergePatch/object_assert_failed_1.jsonnet
@@ -0,0 +1 @@
+std.mergePatch({ assert false }, {})
diff --git a/ui-tests/fail/stdlib/mergePatch/object_assert_failed_1.jsonnet.stderr b/ui-tests/fail/stdlib/mergePatch/object_assert_failed_1.jsonnet.stderr
@@ -0,0 +1,12 @@
+error: assertion failed
+ --> object_assert_failed_1.jsonnet:1:18
+  |
+1 | std.mergePatch({ assert false }, {})
+  |                  ^^^^^^^^^^^^
+note: while evaluating call to `mergePatch`
+ --> object_assert_failed_1.jsonnet:1:1
+  |
+1 | std.mergePatch({ assert false }, {})
+  | ------------------------------------
+note: during top-level value evaluation
+
diff --git a/ui-tests/fail/stdlib/mergePatch/object_assert_failed_2.jsonnet b/ui-tests/fail/stdlib/mergePatch/object_assert_failed_2.jsonnet
@@ -0,0 +1 @@
+std.mergePatch({}, { assert false })
diff --git a/ui-tests/fail/stdlib/mergePatch/object_assert_failed_2.jsonnet.stderr b/ui-tests/fail/stdlib/mergePatch/object_assert_failed_2.jsonnet.stderr
@@ -0,0 +1,12 @@
+error: assertion failed
+ --> object_assert_failed_2.jsonnet:1:22
+  |
+1 | std.mergePatch({}, { assert false })
+  |                      ^^^^^^^^^^^^
+note: while evaluating call to `mergePatch`
+ --> object_assert_failed_2.jsonnet:1:1
+  |
+1 | std.mergePatch({}, { assert false })
+  | ------------------------------------
+note: during top-level value evaluation
+
diff --git a/ui-tests/pass/stdlib/mergePatch.jsonnet b/ui-tests/pass/stdlib/mergePatch.jsonnet
@@ -6,6 +6,13 @@ std.assertEqual(std.mergePatch({ a: 1 }, { a: {} }), { a: {} }) &&
 std.assertEqual(std.mergePatch({ a: 1 }, { a: null }), {}) &&
 std.assertEqual(std.mergePatch({ a: { x: 1, y: 2 }, b: -1 }, { a: { x: 3 } }), { a: { x: 3, y: 2 }, b: -1 }) &&
 std.assertEqual(std.mergePatch({ a: { x: 1, y: 2 }, b: -1 }, { a: { x: null, y: null } }), { a: {}, b: -1 }) &&
+std.assertEqual(std.mergePatch({ a: null }, {}), { a: null }) &&
+std.assertEqual(std.mergePatch({ a: null }, { a: null }), {}) &&
+std.assertEqual(std.mergePatch({ a: null }, { b: null }), { a: null }) &&
+std.assertEqual(std.mergePatch(false, true), true) &&
+std.assertEqual(std.mergePatch(false, {}), {}) &&
+std.assertEqual(std.mergePatch(false, { a: 1 }), { a: 1 }) &&
+std.assertEqual(std.mergePatch(false, { a: null }), {}) &&
 
 // Examples from RFC 7396 (https://datatracker.ietf.org/doc/html/rfc7396)
 std.assertEqual(std.mergePatch({"a":"b"}, {"a":"c"}), {"a":"c"}) &&