Merge pull request #78 from orlandos-nl/feature/faster-equatable

Faster Document equation

Sources/BSON/Document/Document+Array.swift

@@ -17,6 +17,7 @@ extension Document: ExpressibleByArrayLiteral {
                 return values
             }
 
+            // Advance beyond type identifier
             index += 1
             guard skipKey(at: &index) else {
                 return values
@@ -204,7 +205,7 @@ extension Document: ExpressibleByArrayLiteral {
                 storage.setString(javascript.code, at: offset)
                 offset += utf8Size
 
-                storage.setInteger(0, at: offset, endianness: .little, as: UInt8.self)
+                storage.setInteger(0 as UInt8, at: offset)
                 offset += 1
 
                 storage.setBuffer(scopeBuffer, at: offset)

Sources/BSON/Document/Document+Cache.swift

@@ -151,14 +151,19 @@ extension Document {
 
             guard
                 let codeLength = self.storage.getInteger(at: offset &+ 4, endianness: .little, as: Int32.self),
-                let code = self.storage.getString(at: offset &+ 8, length: numericCast(length) - 1)
+                let code = self.storage.getString(at: offset &+ 8, length: numericCast(codeLength) - 1)
             else {
                 return nil
             }
 
+            let documentOffset = offset + 4 + 4 + Int(codeLength)
+
             guard
-                let documentLength = self.storage.getInteger(at: offset &+ 8 &+ numericCast(codeLength), endianness: .little, as: Int32.self),
-                let slice = self.storage.getSlice(at: offset, length: numericCast(documentLength))
+                // Offset + JSCodeWithScope Length + Code Length + Code UTF8
+                let documentLength = self.storage.getInteger(at: documentOffset, endianness: .little, as: Int32.self),
+                // Length == JSCodeWithScope Length + Code Length + Code UTF8 + Document
+                length == 4 + 4 + Int(codeLength) + Int(documentLength),
+                let slice = self.storage.getSlice(at: documentOffset, length: numericCast(documentLength))
             else {
                 return nil
             }

Sources/BSON/Document/Document+Dictionary.swift

@@ -220,7 +220,7 @@ extension Document: ExpressibleByDictionaryLiteral {
             storage.writeInteger(Int32(primitiveLength), endianness: .little) // header
             storage.writeInteger(Int32(codeLength), endianness: .little) // string (code)
             storage.writeString(javascript.code)
-            storage.writeInteger(0, endianness: .little, as: UInt8.self)
+            storage.writeInteger(0 as UInt8)
             storage.writeBuffer(&scopeBuffer)
         case let bsonData as BSONPrimitiveRepresentable:
             self.appendValue(bsonData.primitive, forKey: key)

Sources/BSON/Document/Document+Equatable.swift

@@ -1,39 +1,179 @@
 import Foundation
 
 extension Document: Equatable {
-    public static func == (lhs: Document, rhs: Document) -> Bool {
-        if lhs.isArray != rhs.isArray {
+    private func equateAsArray(with rhs: Document) -> Bool {
+        var lhsBuffer = self.storage
+        var rhsBuffer = rhs.storage
+
+        guard lhsBuffer.readableBytes > 4, rhsBuffer.readableBytes > 4 else {
             return false
         }
 
-        if lhs.count != rhs.count {
+        lhsBuffer.moveReaderIndex(forwardBy: 4)
+        rhsBuffer.moveReaderIndex(forwardBy: 4)
+
+        while let byte = lhsBuffer.getByte(at: lhsBuffer.readerIndex), byte != 0x00 {
+            // Check the next LHS value type
+            guard
+                let lhsTypeId: UInt8 = lhsBuffer.readInteger(),
+                let lhsType = TypeIdentifier(rawValue: lhsTypeId)
+            else {
+                // Both counts end here
+                if
+                    let rhsTypeId: UInt8 = rhsBuffer.readInteger()
+                {
+                    return rhsTypeId == 0
+                } else {
+                    return true
+                }
+            }
+
+            // Check the next RHS value type
+            guard
+                let rhsTypeId: UInt8 = rhsBuffer.readInteger(),
+                let rhsType = TypeIdentifier(rawValue: rhsTypeId)
+            else {
+                return false
+            }
+
+            // Both types must match
+            guard lhsType == rhsType else {
+                return false
+            }
+
+            // Since they're both the same, this is now our type
+            let type = lhsType
+
+            guard
+                let lhsLength = lhsBuffer.firstRelativeIndexOf(startingAt: 0),
+                lhsLength + 1 < lhsBuffer.readableBytes,
+                let rhsLength = rhsBuffer.firstRelativeIndexOf(startingAt: 0),
+                rhsLength + 1 < rhsBuffer.readableBytes
+            else {
+                // Corrupt buffer
+                return false
+            }
+
+            // For arrays, only care about indices. Not keys
+            // Skip until after the null terminator
+            lhsBuffer.moveReaderIndex(forwardBy: lhsLength + 1)
+            rhsBuffer.moveReaderIndex(forwardBy: rhsLength + 1)
+
+            guard
+                let lhsLength = self.valueLength(forType: lhsType, at: lhsBuffer.readerIndex),
+                let rhsLength = rhs.valueLength(forType: rhsType, at: rhsBuffer.readerIndex),
+                let lhsSlice = lhsBuffer.readSlice(length: Int(lhsLength)),
+                let rhsSlice = rhsBuffer.readSlice(length: Int(rhsLength))
+            else {
+                return false
+            }
+
+            if type == .array || type == .document {
+                let lhsSubDocument = Document(buffer: lhsSlice, isArray: type == .array)
+                let rhsSubDocument = Document(buffer: rhsSlice, isArray: type == .array)
+
+                guard lhsSubDocument == rhsSubDocument else {
+                    return false
+                }
+                // TODO: JSCode wit Scope also has a Document embedded in it
+            } else {
+                guard lhsLength == rhsLength, lhsSlice == rhsSlice else {
+                    return false
+                }
+            }
+        }
+
+        if let byte = rhsBuffer.getByte(at: rhsBuffer.readerIndex), byte != 0x00 {
+            // RHS had more data
             return false
         }
 
-        if lhs.isArray {
-            for i in 0..<lhs.count {
-                let lhsValue = lhs[i]
-                let rhsValue = rhs[i]
+        return true
+    }
+
+    private func equateAsDictionary(with rhs: Document) -> Bool {
+        var lhsBuffer = self.storage
+        let rhsBuffer = rhs.storage
+
+        guard lhsBuffer.readableBytes > 4, rhsBuffer.readableBytes > 4 else {
+            return false
+        }
+
+        lhsBuffer.moveReaderIndex(forwardBy: 4)
+
+        var count = 0
+        let rhsCount = rhs.count
+
+        while let byte = lhsBuffer.getByte(at: lhsBuffer.readerIndex), byte != 0x00 {
+            count += 1
+
+            // Early exit, rhs has less fields than lhs
+            if count > rhsCount {
+                return false
+            }
+
+            // Read the next LHS value type
+            guard
+                let lhsTypeId: UInt8 = lhsBuffer.readInteger(),
+                let lhsType = TypeIdentifier(rawValue: lhsTypeId)
+            else {
+                // Unknown type identifier
+                return false
+            }
+
+            guard
+                let lhsKey = lhsBuffer.readNullTerminatedString(),
+                let (rhsType, rhsOffset) = rhs.typeAndValueOffset(forKey: lhsKey)
+            else {
+                // Corrupt buffer
+                return false
+            }
+
+            // Both types must match
+            guard lhsType == rhsType else {
+                return false
+            }
+
+            // Since they're both the same, this is now our type
+            let type = lhsType
+
+            guard
+                let lhsLength = self.valueLength(forType: lhsType, at: lhsBuffer.readerIndex),
+                let rhsLength = rhs.valueLength(forType: rhsType, at: rhsOffset),
+                let lhsSlice = lhsBuffer.readSlice(length: Int(lhsLength)),
+                let rhsSlice = rhsBuffer.getSlice(at: rhsOffset, length: Int(rhsLength))
+            else {
+                return false
+            }
+
+            if type == .array || type == .document {
+                let lhsSubDocument = Document(buffer: lhsSlice, isArray: type == .array)
+                let rhsSubDocument = Document(buffer: rhsSlice, isArray: type == .array)
 
-                guard
-                    lhsValue.equals(rhsValue)
-                else {
+                guard lhsSubDocument == rhsSubDocument else {
                     return false
                 }
-            }
-        } else {
-            for key in lhs.keys {
-                guard
-                    let lhsValue = lhs[key],
-                    let rhsValue = rhs[key],
-                    lhsValue.equals(rhsValue)
-                else {
+            } else {
+                guard lhsSlice == rhsSlice else {
                     return false
                 }
             }
         }
 
-        return true
+        // Ensure lhs and rhs had all their fields scanned
+        return count == rhsCount
+    }
+
+    public static func == (lhs: Document, rhs: Document) -> Bool {
+        if lhs.isArray != rhs.isArray {
+            return false
+        }
+
+        if lhs.isArray {
+            return lhs.equateAsArray(with: rhs)
+        } else {
+            return lhs.equateAsDictionary(with: rhs)
+        }
     }
 }
 

Sources/BSON/Document/Document+Mutations.swift

@@ -83,11 +83,16 @@ extension Document {
         guard
             let typeId = storage.getInteger(at: index, as: UInt8.self),
             let type = TypeIdentifier(rawValue: typeId)
-        else { return false }
+        else {
+            return false
+        }
 
         index += 1
 
-        guard skipKey(at: &index) else { return false }
+        guard skipKey(at: &index) else {
+            return false
+        }
+
         return skipValue(ofType: type, at: &index)
     }
 
@@ -193,13 +198,13 @@ extension Document {
             // Still need to check the key's size
             return MaxKey()
         case .regex:
-            guard let patternEnd = storage.firstRelativeIndexOf(startingAt: offset), let pattern = storage.getString(at: offset, length: patternEnd - 1) else {
+            guard let patternEnd = storage.firstRelativeIndexOf(startingAt: offset), let pattern = storage.getString(at: offset, length: patternEnd) else {
                 return nil
             }
 
-            let offset = offset + patternEnd
+            let offset = offset + patternEnd + 1
 
-            guard let optionsEnd = storage.firstRelativeIndexOf(startingAt: offset), let options = storage.getString(at: offset, length: optionsEnd - 1) else {
+            guard let optionsEnd = storage.firstRelativeIndexOf(startingAt: offset), let options = storage.getString(at: offset, length: optionsEnd) else {
                 return nil
             }
 
@@ -219,17 +224,22 @@ extension Document {
             }
 
             guard
-                let codeLength = self.storage.getInteger(at: offset &+ 4, endianness: .little, as: Int32.self),
-                let code = self.storage.getString(at: offset &+ 8, length: numericCast(length) - 1)
-                else {
-                    return nil
+                let codeLength = self.storage.getInteger(at: offset + 4, endianness: .little, as: Int32.self),
+                let code = self.storage.getString(at: offset + 8, length: Int(codeLength) - 1)
+            else {
+                return nil
             }
+
+            let documentOffset = offset + 4 + 4 + Int(codeLength)
 
             guard
-                let documentLength = self.storage.getInteger(at: offset &+ 8 &+ numericCast(codeLength), endianness: .little, as: Int32.self),
-                let slice = self.storage.getSlice(at: offset, length: numericCast(documentLength))
-                else {
-                    return nil
+                // Offset + JSCodeWithScope Length + Code Length + Code UTF8
+                let documentLength = self.storage.getInteger(at: documentOffset, endianness: .little, as: Int32.self),
+                // Length == JSCodeWithScope Length + Code Length + Code UTF8 + Document
+                length == 4 + 4 + Int(codeLength) + Int(documentLength),
+                let slice = self.storage.getSlice(at: documentOffset, length: numericCast(documentLength))
+            else {
+                return nil
             }
 
             let scope = Document(

Sources/BSON/Document/Document+Subscripts.swift

@@ -1,34 +1,42 @@
 extension Document {
+    internal func typeAndValueOffset(forKey key: String) -> (TypeIdentifier, Int)? {
+        var offset = 4
+
+        repeat {
+            guard
+                let typeId = storage.getInteger(at: offset, as: UInt8.self),
+                let type = TypeIdentifier(rawValue: typeId)
+            else {
+                return nil
+            }
+
+            offset += 1
+
+            let matches = matchesKey(key, at: offset)
+            guard skipKey(at: &offset) else {
+                return nil
+            }
+
+            if matches {
+                return (type, offset)
+            }
+
+            guard skipValue(ofType: type, at: &offset) else {
+                return nil
+            }
+        } while offset + 1 < storage.readableBytes
+
+        return nil
+    }
+
     /// Extracts any `Primitive` fom the value at key `key`
     public subscript(key: String) -> Primitive? {
         get {
-            var offset = 4
-
-            repeat {
-                guard
-                    let typeId = storage.getInteger(at: offset, as: UInt8.self),
-                    let type = TypeIdentifier(rawValue: typeId)
-                else {
-                    return nil
-                }
-
-                offset += 1
-
-                let matches = matchesKey(key, at: offset)
-                guard skipKey(at: &offset) else {
-                    return nil
-                }
-
-                if matches {
-                    return value(forType: type, at: offset)
-                }
-
-                guard skipValue(ofType: type, at: &offset) else {
-                    return nil
-                }
-            } while offset + 1 < storage.readableBytes
-
-            return nil
+            guard let (type, offset) = typeAndValueOffset(forKey: key) else {
+                return nil
+            }
+
+            return value(forType: type, at: offset)
         }
         set {
             var offset = 4