fix(zero-swift): parse, write, compare, and hash deep JSON without recursion

JSONParser accepts 5,000 levels, but parsing recursed and overflowed a
512 KB stack (a macOS test or cooperative thread) at about 300 levels in
a debug build. The parser, JSONSerializer, the query hash's visitValue,
and JSONValue's hash now keep the arrays and objects they are in on an
explicit stack. JSONValue's == recurses for the first 64 levels, which is
faster, and goes on without recursion past them.

Freeing a deep value still recurses (Array's and JSONObject's deinit):
a 5,000-deep value fits in 1 MB in a release build, not in 512 KB.

SweepWireTests no longer needs an 8 MB stack to parse; deepRowValue runs
on 512 KB and frees on 8 MB. JSONDepthTests parses, writes, compares, and
hashes 5,000-deep arrays and objects on a 512 KB thread.

Amp-Thread-ID: https://ampcode.com/threads/T-c513a037-a257-45e0-b9a9-15c15edfe859
Co-authored-by: tmm <tmm@tmm.dev>

Author: ampagent

.changeset/json-parser-depth.md

@@ -0,0 +1,5 @@
+---
+'zero-swift': patch
+---
+
+JSON nested as deep as the parser accepts (5,000 levels) no longer overflows the stack when it is parsed, written, compared, hashed, or hashed as a custom query's arguments, even on a thread with a 512 KB stack, such as a test or cooperative thread on macOS. Before, a debug build overflowed at about 300 levels on such a thread. Freeing such a value still recurses: in a release build, a 5,000-deep value fits in a 1 MB stack, such as the iPhone main thread's, but not in 512 KB.

swift/Sources/Zero/Hash/QueryHash.swift

@@ -133,35 +133,83 @@ package enum QueryHash {
 extension QueryHash.State {
   typealias H = QueryHash
 
+  /// Visits `v` without recursion: the arrays and objects it is in are on
+  /// a stack, so nesting takes heap, not stack.
   mutating func visitValue(_ v: JSONValue) {
+    if visitScalar(v) { return }
+    // An array, or an object's entries in `Object.keys(v).sort()` order,
+    // and the index of the next value.
+    var containers: [(entries: Entries, next: Int)] = []
+    var value = v
+    nested: while true {
+      switch value {
+      case .array(let a):
+        mix(H.tagArr)
+        containers.append((.array(a), 0))
+      case .object(let o):
+        mix(H.tagObj)
+        // By UTF-16 code unit, whatever the order.
+        containers.append((.object(o.sorted { StringCompare.utf16Less($0.key, $1.key) }), 0))
+      default:
+        preconditionFailure("not an array or object")
+      }
+      // Visits values until one is an array or object, ending the
+      // containers that are done.
+      while var container = containers.popLast() {
+        while container.next < container.entries.count {
+          let e: JSONValue
+          switch container.entries {
+          case .array(let a):
+            e = a[container.next]
+          case .object(let entries):
+            mixString(entries[container.next].key)
+            e = entries[container.next].value
+          }
+          container.next += 1
+          if visitScalar(e) { continue }
+          containers.append(container)
+          value = e
+          continue nested
+        }
+        mix(H.tagEnd)
+      }
+      return
+    }
+  }
+
+  enum Entries {
+    case array([JSONValue])
+    case object([(key: String, value: JSONValue)])
+
+    var count: Int {
+      switch self {
+      case .array(let a): a.count
+      case .object(let o): o.count
+      }
+    }
+  }
+
+  /// Visits `v` unless it is an array or object; returns whether it did.
+  private mutating func visitScalar(_ v: JSONValue) -> Bool {
     switch v {
     case .string(let s):
       mixString(s)
     case .number(let n):
       // NaN and Infinity are null on the wire (JSON.stringify writes null).
       if !n.isFinite {
         mix(H.tagNull)
-        return
+      } else {
+        mix(JSNumber.isInt32(n) ? H.tagInt : H.tagFloat)
+        mixNumber(n)
       }
-      mix(JSNumber.isInt32(n) ? H.tagInt : H.tagFloat)
-      mixNumber(n)
     case .bool(let b):
       mix(b ? H.tagTrue : H.tagFalse)
     case .null:
       mix(H.tagNull)
-    case .array(let a):
-      mix(H.tagArr)
-      for e in a { visitValue(e) }
-      mix(H.tagEnd)
-    case .object(let o):
-      mix(H.tagObj)
-      // `Object.keys(v).sort()`: by UTF-16 code unit, whatever the order.
-      for (k, value) in o.sorted(by: { StringCompare.utf16Less($0.key, $1.key) }) {
-        mixString(k)
-        visitValue(value)
-      }
-      mix(H.tagEnd)
+    case .array, .object:
+      return false
     }
+    return true
   }
 
   mutating func visitOptionalString(_ s: String?) {

swift/Sources/Zero/JSON/JSONObject.swift

@@ -403,33 +403,19 @@ extension JSONObject: ExpressibleByDictionaryLiteral {
 extension JSONObject: Hashable {
   /// The same keys, by code unit, with equal values, in any order.
   public static func == (a: JSONObject, b: JSONObject) -> Bool {
-    if a.storage === b.storage { return true }
-    guard a.count == b.count else { return false }
-    for i in 0..<a.count {
-      let (key, value) = a[i]
-      // Usually the keys are in the same order.
-      if StringCompare.equal(key, b.storage.key(at: i)) {
-        if value != b.storage.value(at: i) { return false }
-      } else {
-        guard let j = b.index(forKey: key), value == b.storage.value(at: j) else {
-          return false
-        }
-      }
-    }
-    return true
+    JSONValue.object(a) == .object(b)
   }
 
+  /// Whether the two share their storage, so are equal.
+  func hasSameStorage(as other: JSONObject) -> Bool { storage === other.storage }
+
+  func key(at i: Int) -> String { storage.key(at: i) }
+
+  func value(at i: Int) -> JSONValue { storage.value(at: i) }
+
   /// Independent of the order, as `==` is.
   public func hash(into hasher: inout Hasher) {
-    var acc = 0
-    for (key, value) in self {
-      var h = Hasher()
-      h.combine(JSKey(key))
-      h.combine(value)
-      acc ^= h.finalize()
-    }
-    hasher.combine(count)
-    hasher.combine(acc)
+    JSONValue.object(self).hash(into: &hasher)
   }
 }
 

swift/Sources/Zero/JSON/JSONParser.swift

@@ -30,7 +30,7 @@ package struct JSONParser {
   package static func parse(_ buffer: UnsafeBufferPointer<UInt8>) throws(Error) -> JSONValue {
     var p = JSONParser(buffer)
     p.skipWhitespace()
-    let value = try p.parseValue(depth: 0)
+    let value = try p.parseValue()
     p.skipWhitespace()
     if p.i != p.end { throw p.error("Unexpected non-whitespace character after JSON") }
     return value
@@ -57,89 +57,128 @@ package struct JSONParser {
     }
   }
 
-  private mutating func parseValue(depth: Int) throws(Error) -> JSONValue {
-    guard i < end else { throw error("Unexpected end of JSON input") }
-    switch b[i] {
-    case UInt8(ascii: "{"): return try parseObject(depth: depth)
-    case UInt8(ascii: "["): return try parseArray(depth: depth)
-    case UInt8(ascii: "\""): return .string(try parseString())
-    case UInt8(ascii: "t"): try expectLiteral("true"); return .bool(true)
-    case UInt8(ascii: "f"): try expectLiteral("false"); return .bool(false)
-    case UInt8(ascii: "n"): try expectLiteral("null"); return .null
-    case UInt8(ascii: "-"), UInt8(ascii: "0")...UInt8(ascii: "9"): return .number(try parseNumber())
-    default: throw error("Unexpected token")
-    }
+  /// An array or object that a value being parsed is in.
+  private enum Container {
+    case array([JSONValue])
+    /// An object and the key of the value being parsed.
+    case object(JSONObject, key: String)
   }
 
-  private mutating func expectLiteral(_ literal: StaticString) throws(Error) {
-    let n = literal.utf8CodeUnitCount
-    guard end - i >= n else { throw error("Unexpected end of JSON input") }
-    let p = literal.utf8Start
-    for k in 0..<n where b[i + k] != p[k] {
-      throw error("Unexpected token")
-    }
-    i += n
-  }
-
-  private mutating func parseObject(depth: Int) throws(Error) -> JSONValue {
-    if depth >= Self.maxDepth { throw error("JSON nested too deeply") }
-    i += 1  // {
-    var object = JSONObject()
-    skipWhitespace()
-    if i < end && b[i] == UInt8(ascii: "}") {
-      i += 1
-      return .object(object)
-    }
+  /// Parses a value without recursion, so nesting takes heap, not stack.
+  /// The innermost open array or object is `array` or `object` (with
+  /// `key`); the ones it is in are on `parents`. They move between the two
+  /// (`take()`), so each stays uniquely referenced and adding to it doesn't
+  /// copy it.
+  private mutating func parseValue() throws(Error) -> JSONValue {
+    var parents: [Container] = []
+    var array: [JSONValue]?
+    var object: JSONObject?
+    var key = ""
     while true {
-      skipWhitespace()
       guard i < end else { throw error("Unexpected end of JSON input") }
-      guard b[i] == UInt8(ascii: "\"") else { throw error("Expected property name") }
-      let key = try parseString()
-      skipWhitespace()
-      guard i < end, b[i] == UInt8(ascii: ":") else {
-        throw error("Expected ':' after property name")
-      }
-      i += 1
-      skipWhitespace()
-      object[key] = try parseValue(depth: depth + 1)
-      skipWhitespace()
-      guard i < end else { throw error("Unexpected end of JSON input") }
-      if b[i] == UInt8(ascii: ",") {
+      var value: JSONValue
+      switch b[i] {
+      case UInt8(ascii: "{"), UInt8(ascii: "["):
+        let isObject = b[i] == UInt8(ascii: "{")
+        if parents.count + (array == nil && object == nil ? 0 : 1) >= Self.maxDepth {
+          throw error("JSON nested too deeply")
+        }
         i += 1
+        skipWhitespace()
+        if i < end && b[i] == (isObject ? UInt8(ascii: "}") : UInt8(ascii: "]")) {
+          i += 1
+          value = isObject ? .object(JSONObject()) : .array([])
+          break
+        }
+        if let a = array.take() {
+          parents.append(.array(a))
+        } else if let o = object.take() {
+          parents.append(.object(o, key: key))
+        }
+        if isObject {
+          object = JSONObject()
+          key = try parseKey()
+        } else {
+          array = []
+        }
         continue
+      case UInt8(ascii: "\""): value = .string(try parseString())
+      case UInt8(ascii: "t"): try expectLiteral("true"); value = .bool(true)
+      case UInt8(ascii: "f"): try expectLiteral("false"); value = .bool(false)
+      case UInt8(ascii: "n"): try expectLiteral("null"); value = .null
+      case UInt8(ascii: "-"), UInt8(ascii: "0")...UInt8(ascii: "9"):
+        value = .number(try parseNumber())
+      default: throw error("Unexpected token")
       }
-      if b[i] == UInt8(ascii: "}") {
-        i += 1
-        return .object(object)
+      // Adds the value to its container, closing the containers that end
+      // after it, until one has another value.
+      while true {
+        skipWhitespace()
+        if array != nil {
+          array!.append(value)
+          guard i < end else { throw error("Unexpected end of JSON input") }
+          if b[i] == UInt8(ascii: ",") {
+            i += 1
+            skipWhitespace()
+            break
+          }
+          guard b[i] == UInt8(ascii: "]") else {
+            throw error("Expected ',' or ']' after array element")
+          }
+          i += 1
+          value = .array(array.take()!)
+        } else if object != nil {
+          // A repeated key keeps its first place and its last value.
+          object![key] = value
+          guard i < end else { throw error("Unexpected end of JSON input") }
+          if b[i] == UInt8(ascii: ",") {
+            i += 1
+            key = try parseKey()
+            break
+          }
+          guard b[i] == UInt8(ascii: "}") else {
+            throw error("Expected ',' or '}' after property value")
+          }
+          i += 1
+          value = .object(object.take()!)
+        } else {
+          return value
+        }
+        guard let parent = parents.popLast() else { continue }
+        switch consume parent {
+        case .array(let a):
+          array = consume a
+        case .object(let o, let k):
+          object = consume o
+          key = consume k
+        }
       }
-      throw error("Expected ',' or '}' after property value")
     }
   }
 
-  private mutating func parseArray(depth: Int) throws(Error) -> JSONValue {
-    if depth >= Self.maxDepth { throw error("JSON nested too deeply") }
-    i += 1  // [
-    var array: [JSONValue] = []
+  /// A property name and its colon, and the whitespace around them.
+  private mutating func parseKey() throws(Error) -> String {
     skipWhitespace()
-    if i < end && b[i] == UInt8(ascii: "]") {
-      i += 1
-      return .array(array)
+    guard i < end else { throw error("Unexpected end of JSON input") }
+    guard b[i] == UInt8(ascii: "\"") else { throw error("Expected property name") }
+    let key = try parseString()
+    skipWhitespace()
+    guard i < end, b[i] == UInt8(ascii: ":") else {
+      throw error("Expected ':' after property name")
     }
-    while true {
-      skipWhitespace()
-      array.append(try parseValue(depth: depth + 1))
-      skipWhitespace()
-      guard i < end else { throw error("Unexpected end of JSON input") }
-      if b[i] == UInt8(ascii: ",") {
-        i += 1
-        continue
-      }
-      if b[i] == UInt8(ascii: "]") {
-        i += 1
-        return .array(array)
-      }
-      throw error("Expected ',' or ']' after array element")
+    i += 1
+    skipWhitespace()
+    return key
+  }
+
+  private mutating func expectLiteral(_ literal: StaticString) throws(Error) {
+    let n = literal.utf8CodeUnitCount
+    guard end - i >= n else { throw error("Unexpected end of JSON input") }
+    let p = literal.utf8Start
+    for k in 0..<n where b[i + k] != p[k] {
+      throw error("Unexpected token")
     }
+    i += n
   }
 
   private mutating func parseString() throws(Error) -> String {