Support MutableCollection

orlandos-nl · Nov 1, 2022 · d93cb0b · d93cb0b
1 parent 0e1b3f5
commit d93cb0b
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 21 deletions.
diff --git a/Sources/BSON/Document/Document+Collection.swift b/Sources/BSON/Document/Document+Collection.swift
@@ -1,26 +1,57 @@
-extension Document: RandomAccessCollection {
+extension Document: RandomAccessCollection, MutableCollection {
     public subscript(position: DocumentIndex) -> (String, Primitive) {
-        var offset = 4
-        for _ in 0..<position.offset {
-            guard self.skipKeyValuePair(at: &offset) else {
-                fatalError("DocumentIndex exceeded Document bounds")
+        get {
+            var offset = 4
+            for _ in 0..<position.offset {
+                guard self.skipKeyValuePair(at: &offset) else {
+                    fatalError("DocumentIndex exceeded Document bounds")
+                }
+            }
+
+            let type = TypeIdentifier(rawValue: storage.getByte(at: offset)!)!
+            offset += 1
+
+            let length = storage.firstRelativeIndexOf(startingAt: offset)!
+            let key = storage.getString(at: offset, length: length)!
+            offset += length + 1
+
+            let value = self.value(forType: type, at: offset)!
+
+            return (key, value)
+        }
+        set {
+            guard let key = self.getKey(at: position.offset) else {
+                fatalError("Attempting to change a value out of bounds")
+            }
+
+            if key == newValue.0 {
+                // Same key, just change the value
+                self[key] = newValue.1
+            } else {
+                self[key] = nil
+                self[newValue.0] = newValue.1
             }
         }
-
-        let type = TypeIdentifier(rawValue: storage.getByte(at: offset)!)!
-        offset += 1
-
-        let length = storage.firstRelativeIndexOf(startingAt: offset)!
-        let key = storage.getString(at: offset, length: length)!
-        offset += length + 1
-
-        let value = self.value(forType: type, at: offset)!
-
-        return (key, value)
     }
 
     public subscript(bounds: Range<DocumentIndex>) -> DocumentSlice {
-        DocumentSlice(document: self, startIndex: bounds.lowerBound, endIndex: bounds.upperBound)
+        get {
+            DocumentSlice(document: self, startIndex: bounds.lowerBound, endIndex: bounds.upperBound)
+        }
+        set {
+            let start = bounds.lowerBound.offset
+            let end = bounds.upperBound.offset
+            var index = end
+
+            while index != start {
+                index -= 1
+                remove(at: index)
+            }
+
+            for (key, value) in newValue.document {
+                self[key] = value
+            }
+        }
     }
 
     public typealias Iterator = DocumentIterator

diff --git a/Sources/BSON/Document/DocumentSlice.swift b/Sources/BSON/Document/DocumentSlice.swift
@@ -1,4 +1,4 @@
-public struct DocumentSlice: RandomAccessCollection {
+public struct DocumentSlice: RandomAccessCollection, MutableCollection {
     public func index(before i: DocumentIndex) -> DocumentIndex {
         return DocumentIndex(offset: i.offset - 1)
     }
@@ -8,17 +8,35 @@ public struct DocumentSlice: RandomAccessCollection {
     }
 
     public subscript(position: DocumentIndex) -> (String, Primitive) {
-        return document.pair(atIndex: position)!
+        get { document[position] }
+        set {
+            guard let key = document.getKey(at: position.offset) else {
+                fatalError("Attempting to change a value out of bounds")
+            }
+
+            if key == newValue.0 {
+                // Same key, just change the value
+                document[key] = newValue.1
+            } else {
+                document[key] = nil
+                document[newValue.0] = newValue.1
+            }
+        }
     }
 
     public subscript(bounds: Range<DocumentIndex>) -> DocumentSlice {
-        DocumentSlice(document: document, startIndex: bounds.lowerBound, endIndex: bounds.upperBound)
+        get {
+            DocumentSlice(document: document, startIndex: bounds.lowerBound, endIndex: bounds.upperBound)
+        }
+        set {
+            document[bounds] = newValue
+        }
     }
 
     public typealias Element = (String, Primitive)
     public typealias SubSequence = DocumentSlice
 
-    let document: Document
+    var document: Document
     public let startIndex: DocumentIndex
     public let endIndex: DocumentIndex
 }