From 40657e6e29a5f21ce1d9f0b54020e763c3bfd67d Mon Sep 17 00:00:00 2001
From: Joannis Orlandos <joannis@orlandos.nl>
Date: Wed, 17 May 2017 16:25:25 +0200
Subject: [PATCH] Attempting to fix the issue again. #3

---
 Sources/BSON/Document+Subscripts.swift | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/Sources/BSON/Document+Subscripts.swift b/Sources/BSON/Document+Subscripts.swift
index d3c632a..52021fb 100644
--- a/Sources/BSON/Document+Subscripts.swift
+++ b/Sources/BSON/Document+Subscripts.swift
@@ -103,26 +103,29 @@ extension Document {
                         let len = getLengthOfElement(withDataPosition: meta.dataPosition, type: meta.type)
                         let dataEndPosition = meta.dataPosition+len
                         
-                        storage.removeSubrange(meta.dataPosition..<dataEndPosition)
-                        
                         let relativeLength: Int
                         
                         if let newValue = newValue {
+                            storage.removeSubrange(meta.dataPosition..<dataEndPosition)
                             let oldLength = dataEndPosition - meta.dataPosition
                             let newBinary = newValue.makeBinary()
                             storage.insert(contentsOf: newBinary, at: meta.dataPosition)
                             storage[meta.elementTypePosition] = newValue.typeIdentifier
                             relativeLength = newBinary.count - oldLength
+                            
+                            for (key, startPosition) in searchTree where startPosition > meta.elementTypePosition {
+                                searchTree[key] = startPosition + relativeLength
+                            }
                         } else {
                             storage.removeSubrange(meta.elementTypePosition..<(meta.elementTypePosition + part.bytes.count + 2 + len))
                             // key + null terminator + type
                             relativeLength = -((part.bytes.count + 2) + len)
                             
                             searchTree[part] = nil
-                        }
-                        
-                        for (key, startPosition) in searchTree where startPosition > meta.elementTypePosition {
-                            searchTree[key] = startPosition + relativeLength
+                            
+                            for (key, startPosition) in searchTree where startPosition > meta.dataPosition + relativeLength {
+                                searchTree[key] = startPosition + relativeLength
+                            }
                         }
                         
                         updateDocumentHeader()