Improve count implementation

Author: kylef

RxQueryKit/RxQueryKit.swift

@@ -1,6 +1,7 @@
 import CoreData
 import RxSwift
 import QueryKit
+import Foundation
 
 
 /// Extension to QuerySet to provide observables
@@ -16,8 +17,29 @@ extension QuerySet {
   /// Performs a query for the count of all objects matching the set predicate.
   /// Emits an Int containing the amount of objects matching the predicate and updates when the managed object context is changed.
   func count() throws -> Observable<Int> {
-    return context.qk_objectsDidChange().map { [unowned self] notification in
-      return try self.count()
-    }.startWith(try self.count())
+    var count:Int = try self.count()
+
+    return AnonymousObservable { observer in
+      observer.on(.Next(count))
+
+      return self.context.qk_objectsDidChange().subscribeNext { notification in
+        let insertedObjects = notification.insertedObjects.filter {
+          $0.entity.name == self.entityName
+        }
+        let deletedObjects = notification.deletedObjects.filter {
+          $0.entity.name == self.entityName
+        }
+
+        if let predicate = self.predicate {
+          count += (insertedObjects as NSArray).filteredArrayUsingPredicate(predicate).count
+          count += (deletedObjects as NSArray).filteredArrayUsingPredicate(predicate).count
+        } else {
+          count -= insertedObjects.count
+          count -= deletedObjects.count
+        }
+
+        observer.on(.Next(count))
+      }
+    }
   }
 }