perf(json): index large replacer key sets

json/json.mbt

@@ -116,15 +116,49 @@ priv enum WriteFrame {
 ///|
 /// A Replacer provides a way to filter and transform JSON object properties during stringification.
 /// 
-/// Replacers contain a function that takes a property key and value, and returns:
+/// A replacer takes a property key and value, and returns:
 /// - `Some(value)` to include the property in the output (possibly transformed)
 /// - `None` to exclude the property from the output
 /// 
 /// Only applies to object properties, not array elements.
+///
+/// The replacer is consulted at every nesting level, including inside objects
+/// it has just kept, matching JavaScript's `JSON.stringify`. So a replacer that
+/// keeps only `"a"` turns `{"a": {"b": 1}}` into `{"a": {}}` — the nested
+/// `"b"` is filtered by the same rule.
 pub struct Replacer {
-  priv f : (String, Json) -> Json?
+  priv kind : ReplacerKind
+} derive(@debug.Debug)
+
+///|
+priv enum ReplacerKind {
+  Custom((String, Json) -> Json?)
+  Keep(ArrayView[StringView])
+  Exclude(ArrayView[StringView])
+  AdaptiveKeep(ReplacerKeyLookup)
+  AdaptiveExclude(ReplacerKeyLookup)
+} derive(@debug.Debug)
+
+///|
+/// Lookup state for a key set large enough to be worth indexing. The index is
+/// built on first heavy use and then reused for the lifetime of the `Replacer`.
+priv struct ReplacerKeyLookup {
+  keys : ArrayView[StringView]
+  mut linear_lookups : Int
+  mut index : Map[StringView, Unit]?
 } derive(@debug.Debug)
 
+///|
+// Linear scanning beats hashing for very small key sets, even when repeated
+// for every property of a large object.
+const REPLACER_INDEX_MIN_KEY_COUNT = 9
+
+///|
+// Building an index costs O(keys). Delaying it until the replacer has actually
+// been consulted this many times keeps one-shot uses allocation-free and bounds
+// the wasted work to REPLACER_INDEX_LOOKUP_THRESHOLD * keys.
+const REPLACER_INDEX_LOOKUP_THRESHOLD = 64
+
 ///|
 /// Create a new Replacer with a custom function.
 ///
@@ -153,12 +187,17 @@ pub struct Replacer {
 /// ```
 #alias(new, deprecated="Use `Replacer()` instead")
 pub fn Replacer::Replacer(f : (String, Json) -> Json?) -> Replacer {
-  { f, }
+  { kind: Custom(f) }
 }
 
 ///|
 /// Create a Replacer that only keeps the specified property keys.
-/// All other properties will be excluded from the output.
+/// All other properties will be excluded from the output, at every nesting
+/// level — `keep(["a"])` turns `{"a": {"b": 1}, "c": 2}` into `{"a": {}}`.
+///
+/// `array` is read while the replacer is in use, and a large key list is
+/// indexed on first heavy use, so mutating `array` afterwards is not
+/// guaranteed to be observed.
 /// 
 /// ## Example
 /// 
@@ -175,12 +214,21 @@ pub fn Replacer::Replacer(f : (String, Json) -> Json?) -> Replacer {
 /// }
 /// ```
 pub fn Replacer::keep(array : ArrayView[StringView]) -> Replacer {
-  { f: (idx, value) => if array.contains(idx) { Some(value) } else { None } }
+  if array.length() < REPLACER_INDEX_MIN_KEY_COUNT {
+    { kind: Keep(array) }
+  } else {
+    { kind: AdaptiveKeep({ keys: array, linear_lookups: 0, index: None }) }
+  }
 }
 
 ///|
 /// Create a Replacer that excludes the specified property keys.
-/// All other properties will be included in the output.
+/// All other properties will be included in the output. Keys are excluded at
+/// every nesting level, not just the top level.
+///
+/// `array` is read while the replacer is in use, and a large key list is
+/// indexed on first heavy use, so mutating `array` afterwards is not
+/// guaranteed to be observed.
 ///
 /// ## Example
 ///
@@ -197,7 +245,42 @@ pub fn Replacer::keep(array : ArrayView[StringView]) -> Replacer {
 /// }
 /// ```
 pub fn Replacer::exclude(array : ArrayView[StringView]) -> Replacer {
-  { f: (idx, value) => if array.contains(idx) { None } else { Some(value) } }
+  if array.length() < REPLACER_INDEX_MIN_KEY_COUNT {
+    { kind: Exclude(array) }
+  } else {
+    { kind: AdaptiveExclude({ keys: array, linear_lookups: 0, index: None }) }
+  }
+}
+
+///|
+fn ReplacerKeyLookup::contains(self : ReplacerKeyLookup, key : String) -> Bool {
+  match self.index {
+    Some(index) => index.contains(key)
+    None =>
+      if self.linear_lookups < REPLACER_INDEX_LOOKUP_THRESHOLD {
+        self.linear_lookups += 1
+        self.keys.contains(key)
+      } else {
+        let index : Map[StringView, Unit] = Map([], capacity=self.keys.length())
+        for indexed_key in self.keys {
+          index[indexed_key] = ()
+        }
+        self.index = Some(index)
+        index.contains(key)
+      }
+  }
+}
+
+///|
+#inline
+fn Replacer::apply(self : Replacer, key : String, value : Json) -> Json? {
+  match self.kind {
+    Custom(f) => f(key, value)
+    Keep(keys) => if keys.contains(key) { Some(value) } else { None }
+    Exclude(keys) => if keys.contains(key) { None } else { Some(value) }
+    AdaptiveKeep(keys) => if keys.contains(key) { Some(value) } else { None }
+    AdaptiveExclude(keys) => if keys.contains(key) { None } else { Some(value) }
+  }
 }
 
 ///|
@@ -348,7 +431,7 @@ pub fn Json::stringify(
               Some((k, v)) => {
                 let mut v2 = v
                 if replacer is Some(replacer) {
-                  if (replacer.f)(k, v) is Some(v) {
+                  if replacer.apply(k, v) is Some(v) {
                     v2 = v
                   } else {
                     continue None
@@ -533,7 +616,7 @@ pub fn Json::transform(self : Self, replacer : Replacer) -> Json {
       .iter()
       .filter_map(pair => {
         let (k, v) = pair
-        if (replacer.f)(k, v) is Some(v2) {
+        if replacer.apply(k, v) is Some(v2) {
           Some((k, v2.transform(replacer)))
         } else {
           None

json/json_test.mbt

@@ -193,6 +193,93 @@ test "stringify with replacer" {
   )
 }
 
+///|
+/// A key set large enough to be indexed must behave exactly like the linear
+/// path it replaces, including for keys absent from the object.
+test "Replacer::keep and exclude agree across the indexing threshold" {
+  let object = Map([])
+  for i in 0..<40 {
+    object["k\{i}"] = Json::number(i.to_double())
+  }
+  let json = Json::object(object)
+  let small : Array[StringView] = ["k1", "k2", "k3"]
+  let large : Array[StringView] = Array::makei(40, i => "k\{i * 2}")
+  // Small sets scan linearly, large sets build an index; both must agree with
+  // an equivalent custom replacer.
+  for keys in [small, large] {
+    let members = Map(keys.map(k => (k.to_owned(), true)))
+    let kept_by_list = json.stringify(replacer=@json.Replacer::keep(keys))
+    let kept_by_fn = json.stringify(
+      replacer=Replacer((k, v) => {
+        if members.contains(k) {
+          Some(v)
+        } else {
+          None
+        }
+      }),
+    )
+    assert_eq(kept_by_list, kept_by_fn)
+    let dropped_by_list = json.stringify(replacer=@json.Replacer::exclude(keys))
+    let dropped_by_fn = json.stringify(
+      replacer=Replacer((k, v) => {
+        if members.contains(k) {
+          None
+        } else {
+          Some(v)
+        }
+      }),
+    )
+    assert_eq(dropped_by_list, dropped_by_fn)
+  }
+}
+
+///|
+/// A replacer is reusable: the key set it was built from must still be honoured
+/// after the index has been built, and `transform` must agree with `stringify`.
+test "Replacer with an indexed key set is stable across reuse" {
+  let keys : Array[StringView] = Array::makei(12, i => "k\{i}")
+  let keep = @json.Replacer::keep(keys)
+  let exclude = @json.Replacer::exclude(keys)
+  let object = Map([])
+  // Comfortably past the point where the index is built.
+  for i in 0..<200 {
+    object["k\{i}"] = Json::number(i.to_double())
+  }
+  let json = Json::object(object)
+  let expected_keep = "{" +
+    Array::makei(12, i => "\"k\{i}\":\{i}").join(",") +
+    "}"
+  for _ in 0..<3 {
+    assert_eq(json.stringify(replacer=keep), expected_keep)
+    assert_eq(json.transform(keep).stringify(), expected_keep)
+    let dropped = json.stringify(replacer=exclude)
+    assert_false(dropped.contains("\"k0\":"))
+    assert_false(dropped.contains("\"k11\":"))
+    assert_true(dropped.contains("\"k12\":12"))
+    assert_true(dropped.contains("\"k199\":199"))
+  }
+}
+
+///|
+/// `keep` and `exclude` filter every nesting level, like JavaScript's
+/// `JSON.stringify(value, keys)`.
+test "Replacer::keep filters nested objects too" {
+  let json : Json = { "a": { "b": 1, "c": 2 }, "d": 3 }
+  inspect(
+    json.stringify(replacer=@json.Replacer::keep(["a"])),
+    content="{\"a\":{}}",
+  )
+  inspect(
+    json.stringify(replacer=@json.Replacer::keep(["a", "b"])),
+    content="{\"a\":{\"b\":1}}",
+  )
+  let nested : Json = { "a": [{ "b": 1, "z": 9 }] }
+  inspect(
+    nested.stringify(replacer=@json.Replacer::keep(["a"])),
+    content="{\"a\":[{}]}",
+  )
+}
+
 ///|
 test "stringify with replace recursively" {
   let json : Json = {

json/replacer_bench_test.mbt

@@ -0,0 +1,69 @@
+// Copyright 2026 International Digital Economy Academy
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+///|
+fn make_replacer_bench_data(fields : Int) -> (Json, Array[StringView]) {
+  let object = Map([])
+  let keys : Array[StringView] = Array::makei(2500, i => {
+    "key-" + (i * 2).to_string()
+  })
+  for i in 0..<fields {
+    object["key-" + i.to_string()] = Json::number(i.to_double())
+  }
+  (Json::object(object), keys)
+}
+
+///|
+test "bench Json::stringify Replacer::keep keys=2500 fields=5" (it : @bench.T) {
+  let (json, keys) = make_replacer_bench_data(5)
+  let replacer = @json.Replacer::keep(keys)
+  it.bench(fn() { it.keep(json.stringify(replacer~).length()) })
+}
+
+///|
+test "bench Json::stringify Replacer::keep keys=2500 fields=50" (it : @bench.T) {
+  let (json, keys) = make_replacer_bench_data(50)
+  let replacer = @json.Replacer::keep(keys)
+  it.bench(fn() { it.keep(json.stringify(replacer~).length()) })
+}
+
+///|
+test "bench Json::stringify Replacer::keep keys=2500 fields=5000" (
+  it : @bench.T,
+) {
+  let (json, keys) = make_replacer_bench_data(5000)
+  let replacer = @json.Replacer::keep(keys)
+  it.bench(fn() { it.keep(json.stringify(replacer~).length()) })
+}
+
+///|
+test "bench Json::stringify Replacer::keep keys=3 fields=5000" (it : @bench.T) {
+  let (json, _) = make_replacer_bench_data(5000)
+  let keys : Array[StringView] = ["key-0", "key-2500", "key-4999"]
+  let replacer = @json.Replacer::keep(keys)
+  it.bench(fn() { it.keep(json.stringify(replacer~).length()) })
+}
+
+///|
+/// A replacer built fresh for each call must not pay for an index it never
+/// amortizes.
+test "bench Json::stringify Replacer::keep fresh keys=2500 fields=5" (
+  it : @bench.T,
+) {
+  let (json, keys) = make_replacer_bench_data(5)
+  it.bench(fn() {
+    let replacer = @json.Replacer::keep(keys)
+    it.keep(json.stringify(replacer~).length())
+  })
+}

json/types_test.mbt

@@ -55,7 +55,7 @@ test "Debug for Replacer" {
   @debug.debug_inspect(
     r,
     content=(
-      #|{ f: <function: ...> }
+      #|{ kind: Custom(<function: ...>) }
     ),
   )
 }