chore: Keep hidden field to std.objectRemoveKey and support inheritan… (#586)

port google/jsonnet#1269
refs: google/go-jsonnet#837
refs: google/go-jsonnet#830

Author: He-Pin

sjsonnet/src/sjsonnet/Util.scala

@@ -1,6 +1,8 @@
 package sjsonnet
 
 object Util {
+  // A special empty LinkedHashMap instance to avoid creating multiple empty maps
+  private val EMPTY_LINKED_HASHMAP = new java.util.LinkedHashMap[Any, Any]()
   private val hashMapDefaultLoadFactor = 0.75f
   def prettyIndex(lineStarts: Array[Int], index: Int): String = {
     // Returns (-insertionPoint - 1) when the value is not found, where
@@ -190,10 +192,26 @@ object Util {
     new java.util.LinkedHashMap[K, V](hashMapCapacity(expectedElems), hashMapDefaultLoadFactor)
   }
 
+  /**
+   * Returns a shared empty Java HashMap instance, please do not modify it!
+   */
+  def emptyJavaLinkedHashMap[K, V]: java.util.LinkedHashMap[K, V] =
+    EMPTY_LINKED_HASHMAP.asInstanceOf[java.util.LinkedHashMap[K, V]]
+
   def preSizedJavaHashMap[K, V](expectedElems: Int): java.util.HashMap[K, V] = {
     new java.util.HashMap[K, V](hashMapCapacity(expectedElems), hashMapDefaultLoadFactor)
   }
 
+  def preSizedJavaHashSet[E](expectedElems: Int): java.util.HashSet[E] = {
+    new java.util.HashSet[E](hashMapCapacity(expectedElems), hashMapDefaultLoadFactor)
+  }
+
+  def intersect[E](a: java.util.Set[E], b: java.util.Set[E]): java.util.Set[E] = {
+    val result = new java.util.HashSet[E](a)
+    result.retainAll(b)
+    result
+  }
+
   private def hashMapCapacity(expectedElems: Int): Int = {
     if (expectedElems < 3) {
       expectedElems + 1
@@ -204,4 +222,8 @@ object Util {
       (expectedElems / hashMapDefaultLoadFactor).toInt + 1
     }
   }
+
+  def isNotEmpty(collection: java.util.Collection[_]): Boolean = {
+    (collection ne null) && !collection.isEmpty
+  }
 }

sjsonnet/src/sjsonnet/Val.scala

@@ -4,7 +4,7 @@ import java.util
 import sjsonnet.Expr.Member.Visibility
 import sjsonnet.Expr.Params
 
-import scala.annotation.tailrec
+import scala.annotation.{nowarn, tailrec}
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
@@ -271,7 +271,8 @@ object Val {
       private val triggerAsserts: (Val.Obj, Val.Obj) => Unit,
       `super`: Obj,
       valueCache: util.HashMap[Any, Val] = new util.HashMap[Any, Val](),
-      private var allKeys: util.LinkedHashMap[String, java.lang.Boolean] = null)
+      private var allKeys: util.LinkedHashMap[String, java.lang.Boolean] = null,
+      private val excludedKeys: java.util.Set[String] = null)
       extends Literal
       with Expr.ObjBody {
     private var asserting: Boolean = false
@@ -314,36 +315,144 @@ object Val {
     }
 
     def addSuper(pos: Position, lhs: Val.Obj): Val.Obj = {
-      val objs = mutable.ArrayBuffer(this)
+      // Single traversal: collect chain in this-first order
+      val builder = new mutable.ArrayBuilder.ofRef[Val.Obj]
       var current = this
-      while (current.getSuper != null) {
-        objs += current.getSuper
+      while (current != null) {
+        builder += current
         current = current.getSuper
       }
+      val chain = builder.result()
 
+      // Pre-collect all keys defined in this chain once (only needed if any obj has excludedKeys)
+      lazy val keysInThisChain: java.util.Set[String] = {
+        val set = Util.preSizedJavaHashSet[String](chain.length * 4)
+        for (s <- chain) set.addAll(s.getValue0.keySet())
+        set
+      }
+
+      // Iterate root-first (reverse of collection order) to build the new super chain
       current = lhs
-      for (s <- objs.reverse) {
-        current = new Val.Obj(s.pos, s.getValue0, false, s.triggerAsserts, current)
+      var i = chain.length - 1
+      while (i >= 0) {
+        val s = chain(i)
+        val filteredExcludedKeys = if (s.excludedKeys != null) {
+          Util.intersect(s.excludedKeys, keysInThisChain)
+        } else null
+        current = new Val.Obj(
+          s.pos,
+          s.getValue0,
+          false,
+          s.triggerAsserts,
+          current,
+          new util.HashMap[Any, Val](),
+          null,
+          filteredExcludedKeys
+        )
+        i -= 1
       }
       current
     }
 
+    /**
+     * Create a new object that removes the specified keys from this object.
+     *
+     * The implementation preserves both internal and external inheritance:
+     *   1. Internal: For `objectRemoveKey({ a: 1 } + { b: super.a }, 'a')`, the original object's
+     *      internal super chain is preserved, so `b: super.a` can still access `a`.
+     *   2. External: For `{ a: 1 } + objectRemoveKey({ b: super.a }, 'a')`, the result can
+     *      participate in a new inheritance chain, where `super.a` accesses the new super.
+     *
+     * The approach is to create a thin wrapper object with the original object as super, and mark
+     * the key as excluded via the excludedKeys set. The excluded key won't appear in
+     * allKeyNames/visibleKeyNames, but super.key can still access the value.
+     */
+    @nowarn("cat=deprecation")
+    def removeKeys(pos: Position, keys: String*): Val.Obj = {
+      val excluded =
+        if (keys.length == 1)
+          java.util.Collections.singleton(keys.head)
+        else {
+          import scala.collection.JavaConverters._
+          new util.HashSet[String](keys.asJavaCollection)
+        }
+
+      new Val.Obj(
+        pos,
+        Util.emptyJavaLinkedHashMap[String, Obj.Member],
+        false,
+        null, // No asserts in wrapper; original object's asserts are triggered via super chain
+        this,
+        new util.HashMap[Any, Val](), // NOTE: Must be a dedicated new value cache.
+        null,
+        excluded
+      )
+    }
+
     def prettyName = "object"
     override def asObj: Val.Obj = this
 
     private def gatherKeys(mapping: util.LinkedHashMap[String, java.lang.Boolean]): Unit = {
-      val objs = mutable.ArrayBuffer(this)
+      // Fast path: no super chain — just copy this object's keys directly
+      if (this.getSuper == null) {
+        gatherKeysForSingle(this, null, mapping)
+        return
+      }
+
+      // Single traversal: collect chain in this-first order using ArrayBuilder
+      val builder = new mutable.ArrayBuilder.ofRef[Val.Obj]
       var current = this
-      while (current.getSuper != null) {
-        objs += current.getSuper
+      while (current != null) {
+        builder += current
         current = current.getSuper
       }
+      val chain = builder.result()
+      val chainLength = chain.length
+
+      // Collect all excluded keys, reusing the set directly when only one source has exclusions
+      var exclusionSet: java.util.Set[String] = null
+      var multipleExclusions = false
+      for (s <- chain) {
+        val keys = s.excludedKeys
+        if (Util.isNotEmpty(keys)) {
+          if (exclusionSet == null) {
+            exclusionSet = keys
+          } else {
+            if (!multipleExclusions) {
+              val merged = new util.HashSet[String](exclusionSet.size + keys.size)
+              merged.addAll(exclusionSet)
+              exclusionSet = merged
+              multipleExclusions = true
+            }
+            exclusionSet.asInstanceOf[util.HashSet[String]].addAll(keys)
+          }
+        }
+      }
 
-      for (s <- objs.reverse) {
-        if (s.static) {
-          mapping.putAll(s.allKeys)
-        } else {
-          s.getValue0.forEach { (k, m) =>
+      // Iterate root-first (reverse of collection order) and populate the mapping
+      var i = chainLength - 1
+      while (i >= 0) {
+        gatherKeysForSingle(chain(i), exclusionSet, mapping)
+        i -= 1
+      }
+    }
+
+    /** Gather keys from a single object into the mapping, filtering by exclusions. */
+    private def gatherKeysForSingle(
+        obj: Val.Obj,
+        exclusionSet: java.util.Set[String],
+        mapping: util.LinkedHashMap[String, java.lang.Boolean]): Unit = {
+      if (obj.static) {
+        obj.allKeys
+          .keySet()
+          .forEach(key => {
+            if (exclusionSet == null || !exclusionSet.contains(key)) {
+              mapping.put(key, false)
+            }
+          })
+      } else {
+        obj.getValue0.forEach { (k, m) =>
+          if (exclusionSet == null || !exclusionSet.contains(k)) {
             val vis = m.visibility
             if (!mapping.containsKey(k)) mapping.put(k, vis == Visibility.Hidden)
             else if (vis == Visibility.Hidden) mapping.put(k, true)
@@ -411,6 +520,11 @@ object Val {
           case x    => x
         }
       } else {
+        // Check if the key is excluded (used by objectRemoveKey)
+        // When self != this, we need to check if the key exists in self's visible keys
+        if ((self eq this) && excludedKeys != null && excludedKeys.contains(k)) {
+          Error.fail("Field does not exist: " + k, pos)
+        }
         val cacheKey = if (self eq this) k else (k, self)
         val cachedValue = valueCache.get(cacheKey)
         if (cachedValue != null) {

sjsonnet/src/sjsonnet/stdlib/ObjectModule.scala

@@ -200,11 +200,7 @@ object ObjectModule extends AbstractFunctionModule {
       )
     },
     builtin("objectRemoveKey", "obj", "key") { (pos, ev, o: Val.Obj, key: String) =>
-      val bindings: Array[(String, Val.Obj.Member)] = for {
-        k <- o.visibleKeyNames if k != key
-        v = o.value(k, pos.fileScope.noOffsetPos)(ev)
-      } yield (k, new Val.Obj.ConstMember(false, Visibility.Normal, v))
-      Val.Obj.mk(pos, bindings)
+      o.removeKeys(pos, key)
     },
     builtin("mergePatch", "target", "patch") { (pos, ev, target: Val, patch: Val) =>
       val mergePosition = pos

sjsonnet/test/resources/go_test_suite/builtinObjectRemoveKey_hidden.jsonnet

@@ -0,0 +1,7 @@
+local o = std.objectRemoveKey({foo: 1, bar:: 2, baz: 3}, "foo");
+{
+    fields: std.objectFields(o),
+    all_fields: std.objectFieldsAll(o),
+    bar: o.bar,
+    object: o,
+}

sjsonnet/test/resources/go_test_suite/builtinObjectRemoveKey_hidden.jsonnet.golden

@@ -0,0 +1,13 @@
+{
+   "all_fields": [
+      "bar",
+      "baz"
+   ],
+   "bar": 2,
+   "fields": [
+      "baz"
+   ],
+   "object": {
+      "baz": 3
+   }
+}

sjsonnet/test/resources/go_test_suite/builtinObjectRemoveKey_super.jsonnet

@@ -0,0 +1,6 @@
+local o1 = {
+  assert self.x : 'x',
+  assert super.y : 'y',
+  b: super.a,
+};
+{ a: 'begin' } + std.objectRemoveKey({ y: true } + o1 + { a: 'end', x: true }, 'y')

sjsonnet/test/resources/go_test_suite/builtinObjectRemoveKey_super.jsonnet.golden

@@ -0,0 +1,5 @@
+{
+   "a": "end",
+   "b": "begin",
+   "x": true
+}

sjsonnet/test/resources/go_test_suite/builtinObjectRemoveKey_super_assert.jsonnet

@@ -0,0 +1,11 @@
+local o1 = {
+  assert self.x : 'x',
+  assert super.y : 'y',
+  b: super.a,
+};
+local o2 = {
+  x: true,
+  y: true,
+  a: 'one',
+};
+{ a: 'begin' } + std.objectRemoveKey({ y: std.isString('yyy') } + o1 + o2, 'x')

sjsonnet/test/resources/go_test_suite/builtinObjectRemoveKey_super_assert.jsonnet.golden

@@ -0,0 +1,3 @@
+sjsonnet.Error: Field does not exist: x
+    at [Select x].(builtinObjectRemoveKey_super_assert.jsonnet:2:14)
+

sjsonnet/test/resources/test_suite/stdlib.jsonnet

@@ -1630,6 +1630,20 @@ 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 }) &&
+// Referential transparency still works.
+std.assertEqual(local o1 = { b: super.a }; std.objectRemoveKey({ a: 1 } + o1, 'a'), { b: 1 }) &&
+std.assertEqual(local o1 = { b: super.a }; { a: 1 } + std.objectRemoveKey(o1, 'a'), { a: 1, b: 1 }) &&
+// Hidden fields still work.
+std.assertEqual(std.objectRemoveKey({ a: 1 } + { b:: super.a }, 'a'), {}) &&
+std.assertEqual(std.objectRemoveKey({ a: 1 } + { b:: super.a }, 'a').b, 1) &&
+std.assertEqual(({ a: 1 } + std.objectRemoveKey({ b:: super.a }, 'a')), { a: 1 }) &&
+std.assertEqual(({ a: 1 } + std.objectRemoveKey({ b:: super.a }, 'a')).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') &&