implement std.objectRemoveKey as a builtin which retains hidden fields

The Go Jsonnet implementation already implements objectRemoveKey as
a builtin and retains hidden fields, this brings the C++ Jsonnet into
alignment.

See google/go-jsonnet#830

Author: johnbartholomew

core/desugarer.cpp

@@ -36,7 +36,7 @@ struct BuiltinDecl {
     std::vector<UString> params;
 };
 
-static unsigned long max_builtin = 40;
+static unsigned long max_builtin = 41;
 BuiltinDecl jsonnet_builtin_decl(unsigned long builtin)
 {
     switch (builtin) {
@@ -81,6 +81,7 @@ BuiltinDecl jsonnet_builtin_decl(unsigned long builtin)
         case 38: return {U"decodeUTF8", {U"arr"}};
         case 39: return {U"atan2", {U"y", U"x"}};
         case 40: return {U"hypot", {U"a", U"b"}};
+        case 41: return {U"objectRemoveKey", {U"obj", U"key"}};
         default:
             std::cerr << "INTERNAL ERROR: Unrecognized builtin function: " << builtin << std::endl;
             std::abort();

core/state.h

@@ -35,6 +35,7 @@ struct HeapEntity {
         SIMPLE_OBJECT,
         COMPREHENSION_OBJECT,
         EXTENDED_OBJECT,
+        RESTRICTED_OBJECT,
     };
     GarbageCollectionMark mark;
     Type type;
@@ -210,6 +211,20 @@ struct HeapExtendedObject : public HeapObject {
     }
 };
 
+/** Objects created by std.objectRemoveKey(s). */
+struct HeapRestrictedObject : public HeapLeafObject {
+    /** The 'parent' object from which we inherit. */
+    HeapObject *obj;
+
+    /** A 'filter list' of keys retained in the object. */
+    std::map<const Identifier *, ObjectField::Hide> retainedKeys;
+
+    HeapRestrictedObject(HeapObject *obj, const std::map<const Identifier *, ObjectField::Hide> &retained_keys)
+        : HeapLeafObject(RESTRICTED_OBJECT), obj(obj), retainedKeys(retained_keys)
+    {
+    }
+};
+
 /** Objects created by the ObjectComprehensionSimple construct. */
 struct HeapComprehensionObject : public HeapLeafObject {
     /** The captured environment. */
@@ -382,6 +397,12 @@ class Heap {
                         addIfHeapEntity(obj->right, s.children);
                         break;
                     }
+                    case HeapEntity::RESTRICTED_OBJECT: {
+                        assert(dynamic_cast<HeapRestrictedObject *>(curr));
+                        auto *obj = static_cast<HeapRestrictedObject *>(curr);
+                        addIfHeapEntity(obj->obj, s.children);
+                        break;
+                    }
                     case HeapEntity::COMPREHENSION_OBJECT: {
                         assert(dynamic_cast<HeapComprehensionObject *>(curr));
                         auto *obj = static_cast<HeapComprehensionObject *>(curr);

core/vm.cpp

@@ -688,6 +688,20 @@ class Interpreter {
             auto *l = findObject(f, ext->left, start_from, counter);
             if (l)
                 return l;
+        } else if (auto *ext = dynamic_cast<HeapRestrictedObject *>(curr)) {
+            if (counter >= start_from) {
+                ++counter;
+                const auto it = ext->retainedKeys.find(f);
+                if (it != ext->retainedKeys.end()) {
+                    return findObject(f, ext->obj, start_from, counter);
+                } else {
+                    counter += countLeaves(ext->obj);
+                    return nullptr;
+                }
+            } else {
+                ++counter;
+                return findObject(f, ext->obj, start_from, counter);
+            }
         } else {
             if (counter >= start_from) {
                 if (auto *simp = dynamic_cast<HeapSimpleObject *>(curr)) {
@@ -719,6 +733,9 @@ class Interpreter {
                 r[f.first] = f.second.hide;
             }
 
+        } else if (auto *obj = dynamic_cast<const HeapRestrictedObject *>(obj_)) {
+            return obj->retainedKeys;
+
         } else if (auto *obj = dynamic_cast<const HeapExtendedObject *>(obj_)) {
             r = objectFieldsAux(obj->right);
             for (const auto &pair : objectFieldsAux(obj->left)) {
@@ -848,6 +865,8 @@ class Interpreter {
     {
         if (auto *ext = dynamic_cast<HeapExtendedObject *>(obj)) {
             return countLeaves(ext->left) + countLeaves(ext->right);
+        } else if (auto *ext = dynamic_cast<HeapRestrictedObject *>(obj)) {
+            return countLeaves(ext->obj) + 1;
         } else {
             // Must be a HeapLeafObject.
             return 1;
@@ -944,6 +963,7 @@ class Interpreter {
         builtins["objectHasEx"] = &Interpreter::builtinObjectHasEx;
         builtins["length"] = &Interpreter::builtinLength;
         builtins["objectFieldsEx"] = &Interpreter::builtinObjectFieldsEx;
+        builtins["objectRemoveKey"] = &Interpreter::builtinObjectRemoveKey;
         builtins["codepoint"] = &Interpreter::builtinCodepoint;
         builtins["char"] = &Interpreter::builtinChar;
         builtins["log"] = &Interpreter::builtinLog;
@@ -1263,6 +1283,22 @@ class Interpreter {
         return nullptr;
     }
 
+    const AST *builtinObjectRemoveKey(const LocationRange &loc, const std::vector<Value> &args)
+    {
+        validateBuiltinArgs(loc, "objectRemoveKey", args, {Value::OBJECT, Value::STRING});
+        auto *obj = static_cast<HeapObject *>(args[0].v.h);
+        const auto *key = static_cast<HeapString *>(args[1].v.h);
+        const auto *key_id = alloc->makeIdentifier(key->value);
+        auto fields = objectFieldsAux(obj);
+        const auto it = fields.find(key_id);
+        if (it != fields.end()) {
+            fields.erase(it);
+        }
+        // TODO: If the key isn't found perhaps we can evaluate to the existing `obj`?
+        scratch = makeObject<HeapRestrictedObject>(obj, fields);
+        return nullptr;
+    }
+
     const AST *builtinCodepoint(const LocationRange &loc, const std::vector<Value> &args)
     {
         validateBuiltinArgs(loc, "codepoint", args, {Value::STRING});
@@ -1939,6 +1975,8 @@ class Interpreter {
         if (auto *ext = dynamic_cast<HeapExtendedObject *>(curr)) {
             objectInvariants(ext->right, self, counter, thunks);
             objectInvariants(ext->left, self, counter, thunks);
+        } else if (auto *ext = dynamic_cast<HeapRestrictedObject *>(curr)) {
+            objectInvariants(ext->obj, self, counter, thunks);
         } else {
             if (auto *simp = dynamic_cast<HeapSimpleObject *>(curr)) {
                 for (AST *assert : simp->asserts) {

stdlib/std.jsonnet

@@ -1786,12 +1786,6 @@ limitations under the License.
       std.removeAt(arr, indexes[0])
   ,
 
-  objectRemoveKey(obj, key):: {
-    [k]: obj[k],
-    for k in std.objectFields(obj)
-    if k != key
-  },
-
   sha1(str):: go_only_function,
   sha256(str):: go_only_function,
   sha512(str):: go_only_function,

test_suite/stdlib.jsonnet

@@ -1630,6 +1630,12 @@ std.assertEqual(std.remove([1, 2, 3], 2), [1, 3]) &&
 std.assertEqual(std.removeAt([1, 2, 3], 1), [1, 3]) &&
 
 std.assertEqual(std.objectRemoveKey({ foo: 1, bar: 2, baz: 3 }, 'foo'), { bar: 2, baz: 3 }) &&
+// objectRemoveKey should retain hidden fields as hidden fields.
+std.assertEqual(std.objectRemoveKey({ foo: 1, bar: 2, baz:: 3 }, 'foo').baz, 3) &&
+// objectRemoveKey doesn't break inheritance within the provided object.
+std.assertEqual(std.objectRemoveKey({ a: 1 } + { b: super.a }, 'a'), { b: 1 }) &&
+// objectRemoveKey works with inheritance outside of the object.
+std.assertEqual({ a: 1 } + std.objectRemoveKey({ b: super.a }, 'a'), { a: 1, b: 1 }) &&
 
 std.assertEqual(std.trim('already trimmed string'), 'already trimmed string') &&
 std.assertEqual(std.trim('    string with spaces on both ends     '), 'string with spaces on both ends') &&