SWIFT-1070 Improve performance of BSONDecoder and BSONEncoder (#54)

patrickfreed · web-flow · commit ac0e0ee3f7c5 · 2021-01-22T12:33:22.000-05:00
diff --git a/Sources/SwiftBSON/Array+BSONValue.swift b/Sources/SwiftBSON/Array+BSONValue.swift
@@ -46,7 +46,12 @@ extension Array: BSONValue where Element == BSON {
         guard let doc = try BSONDocument.read(from: &buffer).documentValue else {
             throw BSONError.InternalError(message: "BSON Array cannot be read, failed to get documentValue")
         }
-        return .array(doc.values)
+        var values: [BSON] = []
+        let it = doc.makeIterator()
+        while let (_, val) = try it.nextThrowing() {
+            values.append(val)
+        }
+        return .array(values)
     }
 
     internal func write(to buffer: inout ByteBuffer) {
@@ -56,4 +61,10 @@ extension Array: BSONValue where Element == BSON {
         }
         array.write(to: &buffer)
     }
+
+    internal func validate() throws {
+        for v in self {
+            try v.bsonValue.validate()
+        }
+    }
 }
diff --git a/Sources/SwiftBSON/BSON.swift b/Sources/SwiftBSON/BSON.swift
@@ -446,7 +446,7 @@ extension BSON: ExpressibleByBooleanLiteral {
 
 extension BSON: ExpressibleByDictionaryLiteral {
     public init(dictionaryLiteral elements: (String, BSON)...) {
-        self = .document(BSONDocument(keyValuePairs: elements))
+        self = .document(BSONDocument(dictionaryLiteral: elements))
     }
 }
 
diff --git a/Sources/SwiftBSON/BSONCode.swift b/Sources/SwiftBSON/BSONCode.swift
@@ -186,6 +186,10 @@ extension BSONCodeWithScope: BSONValue {
         return .codeWithScope(BSONCodeWithScope(code: code, scope: scope))
     }
 
+    internal func validate() throws {
+        try self.scope.validate()
+    }
+
     internal func write(to buffer: inout ByteBuffer) {
         let writer = buffer.writerIndex
         buffer.writeInteger(0, endianness: .little, as: Int32.self) // reserve space
diff --git a/Sources/SwiftBSON/BSONDecoder.swift b/Sources/SwiftBSON/BSONDecoder.swift
@@ -511,16 +511,20 @@ private struct _BSONKeyedDecodingContainer<K: CodingKey>: KeyedDecodingContainer
     /// A reference to the decoder we're reading from.
     private let decoder: _BSONDecoder
 
-    /// A reference to the container we're reading from.
-    fileprivate let container: BSONDocument
+    /// An unordered copy of the container we're reading from.
+    fileprivate let container: [String: BSON]
 
     /// The path of coding keys taken to get to this point in decoding.
     public private(set) var codingPath: [CodingKey]
 
     /// Initializes `self`, referencing the given decoder and container.
-    fileprivate init(referencing decoder: _BSONDecoder, wrapping container: BSONDocument) {
+    fileprivate init(referencing decoder: _BSONDecoder, wrapping doc: BSONDocument) {
         self.decoder = decoder
-        self.container = container
+        var map: [String: BSON] = [:]
+        for (k, v) in doc {
+            map[k] = v
+        }
+        self.container = map
         self.codingPath = decoder.codingPath
     }
 
@@ -531,7 +535,7 @@ private struct _BSONKeyedDecodingContainer<K: CodingKey>: KeyedDecodingContainer
 
     /// Returns a Boolean value indicating whether the decoder contains a value associated with the given key.
     public func contains(_ key: Key) -> Bool {
-        self.container.hasKey(key.stringValue)
+        self.container[key.stringValue] != nil
     }
 
     /// A string description of a CodingKey, for use in error messages.
diff --git a/Sources/SwiftBSON/BSONDocument+Sequence.swift b/Sources/SwiftBSON/BSONDocument+Sequence.swift
@@ -15,7 +15,7 @@ extension BSONDocument: Sequence {
     public typealias SubSequence = BSONDocument
 
     /// Returns a `Bool` indicating whether the document is empty.
-    public var isEmpty: Bool { self.keySet.isEmpty }
+    public var isEmpty: Bool { self.storage.encodedLength == BSON_MIN_SIZE }
 
     /// Returns a `DocumentIterator` over the values in this `Document`.
     public func makeIterator() -> BSONDocumentIterator {
@@ -29,7 +29,11 @@ extension BSONDocument: Sequence {
     public func map<T>(
         _ transform: (Element) throws -> T
     ) rethrows -> [T] {
-        try AnySequence(self).map(transform)
+        var values: [T] = []
+        for (k, v) in self {
+            values.append(try transform((key: k, value: v)))
+        }
+        return values
     }
 
     /**
diff --git a/Sources/SwiftBSON/BSONDocument.swift b/Sources/SwiftBSON/BSONDocument.swift
@@ -15,28 +15,15 @@ public struct BSONDocument {
     /// The element type of a document: a tuple containing an individual key-value pair.
     public typealias KeyValuePair = (key: String, value: BSON)
 
-    private var storage: BSONDocumentStorage
-
-    /// An unordered set containing the keys in this document.
-    internal private(set) var keySet: Set<String>
+    internal var storage: BSONDocumentStorage
 
     internal init(_ elements: [BSON]) {
         self = BSONDocument(keyValuePairs: elements.enumerated().map { i, element in (String(i), element) })
     }
 
     internal init(keyValuePairs: [(String, BSON)]) {
-        self.keySet = Set(keyValuePairs.map { $0.0 })
-        guard self.keySet.count == keyValuePairs.count else {
-            fatalError("Dictionary \(keyValuePairs) contains duplicate keys")
-        }
-
         self.storage = BSONDocumentStorage()
 
-        guard !self.keySet.isEmpty else {
-            self = BSONDocument()
-            return
-        }
-
         let start = self.storage.buffer.writerIndex
 
         // reserve space for our byteLength that will be calculated
@@ -58,7 +45,6 @@ public struct BSONDocument {
 
     /// Initializes a new, empty `BSONDocument`.
     public init() {
-        self.keySet = Set()
         self.storage = BSONDocumentStorage()
         self.storage.buffer.writeInteger(5, endianness: .little, as: Int32.self)
         self.storage.buffer.writeBytes([0])
@@ -89,13 +75,12 @@ public struct BSONDocument {
      */
     public init(fromBSON bson: ByteBuffer) throws {
         let storage = BSONDocumentStorage(bson)
-        let keys = try storage.validateAndRetrieveKeys()
-        self = BSONDocument(fromUnsafeBSON: storage, keys: keys)
+        try storage.validate()
+        self = BSONDocument(fromUnsafeBSON: storage)
     }
 
-    internal init(fromUnsafeBSON storage: BSONDocumentStorage, keys: Set<String>) {
+    internal init(fromUnsafeBSON storage: BSONDocumentStorage) {
         self.storage = storage
-        self.keySet = keys
     }
 
     /**
@@ -155,7 +140,13 @@ public struct BSONDocument {
     public var values: [BSON] { self.map { _, val in val } }
 
     /// The number of (key, value) pairs stored at the top level of this document.
-    public var count: Int { self.keySet.count }
+    public var count: Int {
+        do {
+            return try BSONDocumentIterator.getKeys(from: self.storage.buffer).count
+        } catch {
+            return 0
+        }
+    }
 
     /// A copy of the `ByteBuffer` backing this document, containing raw BSON data. As `ByteBuffer`s implement
     /// copy-on-write, this copy will share byte storage with this document until either the document or the returned
@@ -166,7 +157,9 @@ public struct BSONDocument {
     public func toData() -> Data { Data(self.storage.buffer.readableBytesView) }
 
     /// Returns a `Boolean` indicating whether this `BSONDocument` contains the provided key.
-    public func hasKey(_ key: String) -> Bool { self.keySet.contains(key) }
+    public func hasKey(_ key: String) -> Bool {
+        (try? BSONDocumentIterator.find(key: key, in: self)) != nil
+    }
 
     /**
      * Allows getting and setting values on the document via subscript syntax.
@@ -236,7 +229,7 @@ public struct BSONDocument {
      * - SeeAlso: https://docs.mongodb.com/manual/core/document/#the-id-field
      */
     public func withID() throws -> BSONDocument {
-        guard !self.keySet.contains("_id") else {
+        guard !self.hasKey("_id") else {
             return self
         }
 
@@ -259,37 +252,35 @@ public struct BSONDocument {
         }
         newStorage.buffer.writeBytes(suffix)
 
-        var newKeys = self.keySet
-        newKeys.insert("_id")
-        var document = BSONDocument(fromUnsafeBSON: newStorage, keys: newKeys)
+        var document = BSONDocument(fromUnsafeBSON: newStorage)
         document.storage.encodedLength = newSize
         return document
     }
 
+    /// Appends the provided key value pair without checking to see if the key already exists.
+    /// Warning: appending two of the same key may result in errors or undefined behavior.
+    internal mutating func append(key: String, value: BSON) {
+        // setup to overwrite null terminator
+        self.storage.buffer.moveWriterIndex(to: self.storage.encodedLength - 1)
+        let size = self.storage.append(key: key, value: value)
+        self.storage.buffer.writeInteger(0, endianness: .little, as: UInt8.self) // add back in our null terminator
+        self.storage.encodedLength += size
+    }
+
     /**
      * Sets a BSON element with the corresponding key
      * if element.value is nil the element is deleted from the BSON
      */
     internal mutating func set(key: String, to value: BSON?) throws {
-        if !self.keySet.contains(key) {
+        guard let range = try BSONDocumentIterator.findByteRange(for: key, in: self) else {
             guard let value = value else {
                 // no-op: key does not exist and the value is nil
                 return
             }
-            // appending new key
-            self.keySet.insert(key)
-            // setup to overwrite null terminator
-            self.storage.buffer.moveWriterIndex(to: self.storage.encodedLength - 1)
-            let size = self.storage.append(key: key, value: value)
-            self.storage.buffer.writeInteger(0, endianness: .little, as: UInt8.self) // add back in our null terminator
-            self.storage.encodedLength += size
+            self.append(key: key, value: value)
             return
         }
 
-        guard let range = try BSONDocumentIterator.findByteRange(for: key, in: self) else {
-            throw BSONError.InternalError(message: "Cannot find \(key) to delete")
-        }
-
         let prefixLength = range.startIndex
         let suffixLength = self.storage.encodedLength - range.endIndex
 
@@ -316,9 +307,6 @@ public struct BSONDocument {
             guard newSize <= BSON_MAX_SIZE else {
                 throw BSONError.DocumentTooLargeError(value: value.bsonValue, forKey: key)
             }
-        } else {
-            // Deleting
-            self.keySet.remove(key)
         }
 
         newStorage.buffer.writeBytes(suffix)
@@ -401,8 +389,7 @@ public struct BSONDocument {
             return totalBytes
         }
 
-        @discardableResult
-        internal func validateAndRetrieveKeys() throws -> Set<String> {
+        internal func validate() throws {
             // Pull apart the underlying binary into [KeyValuePair], should reveal issues
             guard let encodedLength = self.buffer.getInteger(at: 0, endianness: .little, as: Int32.self) else {
                 throw BSONError.InvalidArgumentError(message: "Validation Failed: Cannot read encoded length")
@@ -423,7 +410,6 @@ public struct BSONDocument {
 
             var keySet = Set<String>()
             let iter = BSONDocumentIterator(over: self.buffer)
-            // Implicitly validate with iterator
             do {
                 while let (key, value) = try iter.nextThrowing() {
                     let (inserted, _) = keySet.insert(key)
@@ -432,26 +418,13 @@ public struct BSONDocument {
                             message: "Validation Failed: BSON contains multiple values for key \(key)"
                         )
                     }
-                    switch value {
-                    case let .document(doc):
-                        try doc.storage.validateAndRetrieveKeys()
-                    case let .array(array):
-                        for item in array {
-                            if let doc = item.documentValue {
-                                try doc.storage.validateAndRetrieveKeys()
-                            }
-                        }
-                    default:
-                        continue
-                    }
+                    try value.bsonValue.validate()
                 }
             } catch let error as BSONError.InternalError {
                 throw BSONError.InvalidArgumentError(
                     message: "Validation Failed: \(error.message)"
                 )
             }
-
-            return keySet
         }
     }
 }
@@ -469,6 +442,13 @@ extension BSONDocument: ExpressibleByDictionaryLiteral {
      * - Returns: a new `BSONDocument`
      */
     public init(dictionaryLiteral keyValuePairs: (String, BSON)...) {
+        self.init(dictionaryLiteral: Array(keyValuePairs))
+    }
+
+    internal init(dictionaryLiteral keyValuePairs: [(String, BSON)]) {
+        guard Set(keyValuePairs.map { $0.0 }).count == keyValuePairs.count else {
+            fatalError("Dictionary \(keyValuePairs) contains duplicate keys")
+        }
         self.init(keyValuePairs: keyValuePairs)
     }
 }
@@ -548,13 +528,16 @@ extension BSONDocument: BSONValue {
             throw BSONError.InternalError(message: "Cannot read document contents")
         }
 
-        let keys = try BSONDocumentIterator.getKeySet(from: bytes)
-        return .document(BSONDocument(fromUnsafeBSON: BSONDocument.BSONDocumentStorage(bytes), keys: keys))
+        return .document(BSONDocument(fromUnsafeBSON: BSONDocument.BSONDocumentStorage(bytes)))
     }
 
     internal func write(to buffer: inout ByteBuffer) {
         buffer.writeBytes(self.storage.buffer.readableBytesView)
     }
+
+    internal func validate() throws {
+        try self.storage.validate()
+    }
 }
 
 extension BSONDocument: CustomStringConvertible {
@@ -575,7 +558,7 @@ extension BSONDocument {
      * - Returns: a `Bool` indicating whether the two documents are equal.
      */
     public func equalsIgnoreKeyOrder(_ other: BSONDocument) -> Bool {
-        guard self.count == other.count else {
+        guard self.storage.encodedLength == other.storage.encodedLength else {
             return false
         }
 
diff --git a/Sources/SwiftBSON/BSONEncoder.swift b/Sources/SwiftBSON/BSONEncoder.swift
diff --git a/Sources/SwiftBSON/BSONValue.swift b/Sources/SwiftBSON/BSONValue.swift
diff --git a/Sources/SwiftBSON/ExtendedJSONDecoder.swift b/Sources/SwiftBSON/ExtendedJSONDecoder.swift

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,12 @@ extension Array: BSONValue where Element == BSON {`
`46`	`46`	`guard let doc = try BSONDocument.read(from: &buffer).documentValue else {`
`47`	`47`	`throw BSONError.InternalError(message: "BSON Array cannot be read, failed to get documentValue")`
`48`	`48`	`}`
`49`		`- return .array(doc.values)`
	`49`	`+ var values: [BSON] = []`
	`50`	`+ let it = doc.makeIterator()`
	`51`	`+ while let (_, val) = try it.nextThrowing() {`
	`52`	`+ values.append(val)`
	`53`	`+ }`
	`54`	`+ return .array(values)`
`50`	`55`	`}`
`51`	`56`
`52`	`57`	`internal func write(to buffer: inout ByteBuffer) {`
`@@ -56,4 +61,10 @@ extension Array: BSONValue where Element == BSON {`
`56`	`61`	`}`
`57`	`62`	`array.write(to: &buffer)`
`58`	`63`	`}`
	`64`	`+`
	`65`	`+ internal func validate() throws {`
	`66`	`+ for v in self {`
	`67`	`+ try v.bsonValue.validate()`
	`68`	`+ }`
	`69`	`+ }`
`59`	`70`	`}`
Original file line number	Diff line number	Diff line change
`@@ -446,7 +446,7 @@ extension BSON: ExpressibleByBooleanLiteral {`
`446`	`446`
`447`	`447`	`extension BSON: ExpressibleByDictionaryLiteral {`
`448`	`448`	`public init(dictionaryLiteral elements: (String, BSON)...) {`
`449`		`- self = .document(BSONDocument(keyValuePairs: elements))`
	`449`	`+ self = .document(BSONDocument(dictionaryLiteral: elements))`
`450`	`450`	`}`
`451`	`451`	`}`
`452`	`452`