Optimize set operations with O(n+m) merge-walk algorithm

This replaces the previous O(n log n) and O(n * log m) implementations
of setUnion, setInter, and setDiff with efficient O(n+m) merge-walk
algorithms that exploit the fact that sets are already sorted.

Changes:
- setUnion: Changed from concat + sort + dedup to merge-walk
  (was O((n+m) log(n+m)), now O(n+m))
- setInter: Changed from iterate + binary search to merge-walk
  (was O(min(n,m) * log(max(n,m))), now O(n+m))
- setDiff: Changed from iterate + binary search to merge-walk
  (was O(n * log m), now O(n+m))
- Added applyKeyF helper function for cleaner key function handling
- Added validateSet calls to setUnion for consistency with other ops
- Pre-sized ArrayBuffer allocations for better memory efficiency

The setMember function still uses binary search which is optimal for
single-element lookups.

These optimizations address the same performance issue as upstream
PR databricks#574, but avoid the O(n²) bug in that PR's uniqArr implementation
(which calls ArrayBuilder.result() on each iteration).

All changes respect the throwErrorForInvalidSets setting.

Author: claude

sjsonnet/src/sjsonnet/Std.scala

@@ -1707,15 +1707,53 @@ class Std(
     builtin(Set_),
     builtinWithDefaults("setUnion", "a" -> null, "b" -> null, "keyF" -> Val.False(dummyPos)) {
       (args, pos, ev) =>
+        val keyF = args(2)
+        validateSet(ev, pos, keyF, args(0))
+        validateSet(ev, pos, keyF, args(1))
+
         val a = toSetArr(args, 0, pos, ev)
         val b = toSetArr(args, 1, pos, ev)
+
         if (a.isEmpty) {
-          uniqArr(pos, ev, sortArr(pos, ev, args(1), args(2)), args(2))
+          args(1)
         } else if (b.isEmpty) {
-          uniqArr(pos, ev, sortArr(pos, ev, args(0), args(2)), args(2))
+          args(0)
         } else {
-          val concat = Val.Arr(pos, a ++ b)
-          uniqArr(pos, ev, sortArr(pos, ev, concat, args(2)), args(2))
+          // Use merge-walk algorithm: O(n + m) instead of O((n+m) log(n+m))
+          val out = new mutable.ArrayBuffer[Lazy](a.length + b.length)
+          var idxA = 0
+          var idxB = 0
+
+          while (idxA < a.length && idxB < b.length) {
+            val elemA = a(idxA).force
+            val elemB = b(idxB).force
+            val keyA = applyKeyF(elemA, keyF, pos, ev)
+            val keyB = applyKeyF(elemB, keyF, pos, ev)
+
+            val cmp = ev.compare(keyA, keyB)
+            if (cmp < 0) {
+              out.append(a(idxA))
+              idxA += 1
+            } else if (cmp > 0) {
+              out.append(b(idxB))
+              idxB += 1
+            } else {
+              // Equal keys: take from a, skip both
+              out.append(a(idxA))
+              idxA += 1
+              idxB += 1
+            }
+          }
+          // Append remaining elements
+          while (idxA < a.length) {
+            out.append(a(idxA))
+            idxA += 1
+          }
+          while (idxB < b.length) {
+            out.append(b(idxB))
+            idxB += 1
+          }
+          Val.Arr(pos, out.toArray)
         }
     },
     builtinWithDefaults("setInter", "a" -> null, "b" -> null, "keyF" -> Val.False(dummyPos)) {
@@ -1727,13 +1765,27 @@ class Std(
         val a = toSetArr(args, 0, pos, ev)
         val b = toSetArr(args, 1, pos, ev)
 
-        val out = new mutable.ArrayBuffer[Lazy]
-
-        // The intersection will always be, at most, the size of the smallest set.
-        val sets = if (b.length < a.length) (b, a) else (a, b)
-        for (v <- sets._1) {
-          if (existsInSet(ev, pos, keyF, sets._2, v.force)) {
-            out.append(v)
+        // Use merge-walk algorithm: O(n + m) instead of O(min(n,m) * log(max(n,m)))
+        val out = new mutable.ArrayBuffer[Lazy](math.min(a.length, b.length))
+        var idxA = 0
+        var idxB = 0
+
+        while (idxA < a.length && idxB < b.length) {
+          val elemA = a(idxA).force
+          val elemB = b(idxB).force
+          val keyA = applyKeyF(elemA, keyF, pos, ev)
+          val keyB = applyKeyF(elemB, keyF, pos, ev)
+
+          val cmp = ev.compare(keyA, keyB)
+          if (cmp < 0) {
+            idxA += 1
+          } else if (cmp > 0) {
+            idxB += 1
+          } else {
+            // Equal keys: found intersection element
+            out.append(a(idxA))
+            idxA += 1
+            idxB += 1
           }
         }
         Val.Arr(pos, out.toArray)
@@ -1746,14 +1798,45 @@ class Std(
 
         val a = toSetArr(args, 0, pos, ev)
         val b = toSetArr(args, 1, pos, ev)
-        val out = new mutable.ArrayBuffer[Lazy]
 
-        for (v <- a) {
-          if (!existsInSet(ev, pos, keyF, b, v.force)) {
-            out.append(v)
+        if (b.isEmpty) {
+          args(0)
+        } else {
+          // Use merge-walk algorithm: O(n + m) instead of O(n * log(m))
+          val out = new mutable.ArrayBuffer[Lazy](a.length)
+          var idxA = 0
+          var idxB = 0
+
+          while (idxA < a.length) {
+            val elemA = a(idxA).force
+            val keyA = applyKeyF(elemA, keyF, pos, ev)
+
+            // Advance idxB past elements smaller than keyA
+            while (idxB < b.length && {
+              val keyB = applyKeyF(b(idxB).force, keyF, pos, ev)
+              ev.compare(keyA, keyB) > 0
+            }) {
+              idxB += 1
+            }
+
+            // Check if current b element matches keyA
+            if (idxB < b.length) {
+              val keyB = applyKeyF(b(idxB).force, keyF, pos, ev)
+              if (ev.compare(keyA, keyB) != 0) {
+                // keyA not found in b, add to output
+                out.append(a(idxA))
+              } else {
+                // Found match, skip this element and advance b
+                idxB += 1
+              }
+            } else {
+              // Exhausted b, all remaining elements of a are in the diff
+              out.append(a(idxA))
+            }
+            idxA += 1
           }
+          Val.Arr(pos, out.toArray)
         }
-        Val.Arr(pos, out.toArray)
     },
     builtinWithDefaults("setMember", "x" -> null, "arr" -> null, "keyF" -> Val.False(dummyPos)) {
       (args, pos, ev) =>
@@ -1972,6 +2055,14 @@ class Std(
     }
   }
 
+  private def applyKeyF(elem: Val, keyF: Val, pos: Position, ev: EvalScope): Val = {
+    keyF match {
+      case keyFFunc: Val.Func =>
+        keyFFunc.apply1(elem, pos.noOffset)(ev, TailstrictModeDisabled)
+      case _ => elem
+    }
+  }
+
   private def existsInSet(
       ev: EvalScope,
       pos: Position,