NSFetchedResultsController replacement

FetchedResultsController.swift

//
//  SectionedFetchedResults.swift
//  DungeonMaster
//
//  Created by Scott James Remnant on 3/14/16.
//  Copyright © 2016 Scott James Remnant. All rights reserved.
//

import CoreData
import Foundation

// Might be a way here to have our cake and eat it with that protocol struct... basically comes down to defining a struct around it with the requirements as <T> rather than associatedtype, and then directly using that in the init

// TODO test using ObjectIdentifier as the Section, wrapping an object that we place in a separate member of the SectionInfo
// TODO remove Comparable requirement from Section by providing an alternate .sort() for the sections list
// ^ trying to diverge based on this is a red herring, since ObjectIdentifier : Comparable so we wouldn't realize anyway

// could have an init which pre-defines Section to something like String, and another that makes it ObjectIdentifier ??

// TODO static sections?
// TODO placeholder entry in section (for insert)

public class FetchedResultsController<Section : protocol<Hashable, Comparable>, Entity : NSManagedObject> : NSObject {
    
    /// The fetch request used to do the fetching.
    ///
    /// Changes to the fetch request are not automatically reflected in the fetched results. You must call `performFetch()` to refresh the fetched results.
    public let fetchRequest: NSFetchRequest
    
    /// The managed object context used to fetch objects.
    ///
    /// The controller registers to listen to change notifications on this context in order to update the fetched results.
    public let managedObjectContext: NSManagedObjectContext

    public let sectionForObject: (Entity) -> Section
    public let sectionKeys: Set<String>?
    
    /// The results of the fetch.
    ///
    /// The value is `nil` until `performFetch()` is called.
    public private(set) var fetchedObjects: [Entity]?
    
    public private(set) var sections: [FetchedResultsSectionInfo<Section, Entity>]?
    
    public let handleChanges: (([FetchedResultsChange<Section, Entity>]) -> Void)?

    public init(fetchRequest: NSFetchRequest, managedObjectContext: NSManagedObjectContext, sectionForObject: (Entity) -> Section, sectionKeys: Set<String>?, handleChanges: (([FetchedResultsChange<Section, Entity>]) -> Void)?) {
        self.fetchRequest = fetchRequest
        self.managedObjectContext = managedObjectContext
        
        self.sectionForObject = sectionForObject
        self.sectionKeys = sectionKeys
        
        self.handleChanges = handleChanges
        
        super.init()

        NSNotificationCenter.defaultCenter().addObserver(self, selector: #selector(managedObjectContextObjectsDidChange(_:)), name: NSManagedObjectContextObjectsDidChangeNotification, object: self.managedObjectContext)
    }
    
    deinit {
        NSNotificationCenter.defaultCenter().removeObserver(self)
    }
    
    /// Cache mapping Section to Index.
    private var sectionIndexes: [Section: Int]?

    /// Returns the index of `section`.
    ///
    /// This uses a cache to avoid performing expensive `indexOf` operations multiple times.
    private func indexOfSection(section: Section) -> Int? {
        if let sectionIndex = sectionIndexes![section] {
            return sectionIndex
        } else if let sectionIndex = sections!.indexOf({ $0.name == section }) {
            sectionIndexes![section] = sectionIndex
            return sectionIndex
        } else {
            return nil
        }
    }
    
    /// Returns a new SectionInfo for `section` containing `object`.
    ///
    /// The closure passed during class initialization is used to make the actual `SectionInfo` object, and then it is populated with `object` and added to the `sections` list.
    private func makeSection(section section: Section, object: Entity) -> FetchedResultsSectionInfo<Section, Entity> {
        var sectionInfo = FetchedResultsSectionInfo<Section, Entity>(name: section)
        
        sectionInfo.numberOfObjects = 1
        sectionInfo.objects.append(object)
        
        sectionIndexes![section] = sections!.count
        sections!.append(sectionInfo)
        
        return sectionInfo
    }

    /// Cache of sort descriptor keys used for the previous fetch.
    private var sortKeys: Set<String>?
    
    /// Cache mapping ObjectIdentifier to Section it can be found within.
    private var objectSections: [ObjectIdentifier: Section]?
    
    /// Executes the fetch request.
    public func performFetch(notifyChanges notifyChanges: Bool = false) throws {
        fetchedObjects = try managedObjectContext.executeFetchRequest(fetchRequest) as? [Entity]
        
        // Save the old sections if we need to examine it for changes. This is done after performing the fetch since that can cause pending changes to be committed, and the notification observer called, and we don't want to double-notify changes.
        let oldSections = notifyChanges ? sections : nil

        sections = []
        objectSections = [:]
        sectionIndexes = [:]
        
        for object in fetchedObjects! {
            let section = sectionForObject(object)
            objectSections![ObjectIdentifier(object)] = section
            
            if let sectionIndex = indexOfSection(section) {
                sections![sectionIndex].numberOfObjects += 1
                sections![sectionIndex].objects.append(object)
            } else {
                makeSection(section: section, object: object)
            }
        }
        
        sortSections()
        
        if let oldSections = oldSections where notifyChanges {
            let changes = determineChanges(from: oldSections)
            handleChanges?(changes)
        }
        
        // Make a cache of the top-level keys that we used to sort the objects; when objects are updated later, we'll check the update keys against this set, and only re-sort when an update potentially changes the sort order.
        if let sortDescriptorKeys = fetchRequest.sortDescriptors?.flatMap({ $0.key?.componentsSeparatedByString(".").first }) {
            sortKeys = Set(sortDescriptorKeys)
        } else {
            sortKeys = nil
        }
    }
    
    /// Sorts the sections list.
    private func sortSections() {
        sections = sections!.sort({ $0.name < $1.name })
        sectionIndexes = [:]
    }
    
    private func determineChanges(from oldSections: [FetchedResultsSectionInfo<Section, Entity>]) -> [FetchedResultsChange<Section, Entity>] {
        var changes: [FetchedResultsChange<Section, Entity>] = []
        
        var indexes: [ObjectIdentifier: (oldSection: Section, oldSectionIndex: Int, oldIndex: Int)] = [:]
        var priorDeletes: [Section: [Int]] = [:]
        for (oldSectionIndex, oldSectionInfo) in oldSections.enumerate() {
            let oldSection = oldSectionInfo.name
            
            priorDeletes[oldSection] = []
            for (oldIndex, object) in oldSectionInfo.objects.enumerate() {
                if let section = objectSections![ObjectIdentifier(object)] {
                    indexes[ObjectIdentifier(object)] = (oldSection: oldSection, oldSectionIndex: oldSectionIndex, oldIndex: oldIndex)
                    if section != oldSection {
                        // Object has been moved from this section to another.
                        priorDeletes[oldSection]!.append((priorDeletes[oldSection]!.last ?? 0) + 1)
                    } else {
                        // Object remains in the same section.
                        priorDeletes[oldSection]!.append(priorDeletes[oldSection]!.last ?? 0)
                    }
                } else {
                    // Object has been deleted.
                    changes.append(.Delete(object: object, indexPath: NSIndexPath(forRow: oldIndex, inSection: oldSectionIndex)))
                    priorDeletes[oldSection]!.append((priorDeletes[oldSection]!.last ?? 0) + 1)
                }
            }

            // Delete the section if has no index in the new results.
            if indexOfSection(oldSection) == nil {
                // Strictly speaking this oldSectionInfo is wrong as it has members.
                changes.append(.DeleteSection(sectionInfo: oldSectionInfo, index: oldSectionIndex))
            }
        }
        
        for (sectionIndex, sectionInfo) in sections!.enumerate() {
            let section = sectionInfo.name
            
            // If the section didn't exist in the old results, insert it.
            let priorDeletes = priorDeletes[section]
            if priorDeletes == nil {
                // Strictly speaking this sectionInfo is wrong as it has members.
                changes.append(.InsertSection(sectionInfo: sectionInfo, newIndex: sectionIndex))
            }
            
            var priorInserts = 0
            for (index, object) in sectionInfo.objects.enumerate() {
                if let (oldSection, oldSectionIndex, oldIndex) = indexes[ObjectIdentifier(object)] {
                    if section != oldSection {
                        // Object has been moved into this section from another.
                        changes.append(.Move(object: object, indexPath: NSIndexPath(forRow: oldIndex, inSection: oldSectionIndex), newIndexPath: NSIndexPath(forRow: index, inSection: sectionIndex)))
                        priorInserts += 1
                    } else {
                        let adjustedIndex = oldIndex + (priorDeletes?[index] ?? 0)
                        let adjustedNewIndex = index - priorInserts
                        
                        if adjustedIndex != adjustedNewIndex {
                            // Moved.
                            changes.append(.Move(object: object, indexPath: NSIndexPath(forRow: oldIndex, inSection: oldSectionIndex), newIndexPath: NSIndexPath(forRow: index, inSection: sectionIndex)))
                        }
                    }
                } else {
                    // Object was inserted.
                    changes.append(.Insert(object: object, newIndexPath: NSIndexPath(forRow: index, inSection: sectionIndex)))
                    priorInserts += 1
                }
            }
        }
        
        return changes
    }

    @objc private func managedObjectContextObjectsDidChange(notification: NSNotification) {
        assert(notification.name == NSManagedObjectContextObjectsDidChangeNotification, "Notification method called for wrong notification.")
        assert(notification.object === managedObjectContext, "Notification called for incorrect managed object context.")
        
        var objects: Set<Entity> = Set()
        if let objectsSet = notification.userInfo?[NSInsertedObjectsKey] as? NSSet {
            objects.unionInPlace(objectsSet.allObjects as! [Entity])
        }
        if let objectsSet = notification.userInfo?[NSDeletedObjectsKey] as? NSSet {
            objects.unionInPlace(objectsSet.allObjects as! [Entity])
        }
        if let objectsSet = notification.userInfo?[NSUpdatedObjectsKey] as? NSSet {
            objects.unionInPlace(objectsSet.allObjects as! [Entity])
        }
        
        var changes: [FetchedResultsChange<Section, Entity>] = []
        sectionIndexes = [:]

        var insertedSections: [Section] = []
        var insertedObjects: [Section: [(object: Entity, sectionIndex: Int?, index: Int?)]] = [:]
        var deleteIndexes: [Section: NSMutableIndexSet] = [:]

        for object in objects {
            guard object.entity === fetchRequest.entity else { continue }
            
            if let section = objectSections![ObjectIdentifier(object)] {
                let sectionIndex = indexOfSection(section)!
                let index = sections![sectionIndex].objects.indexOf(object)!

                if object.deleted || !(fetchRequest.predicate?.evaluateWithObject(object) ?? true) {
                    // Object was deleted, or previously did, but no now longer does, match the predicate.
                    objectSections![ObjectIdentifier(object)] = nil
                    
                    if let _ = deleteIndexes[section] {
                        deleteIndexes[section]!.addIndex(index)
                    } else {
                        deleteIndexes[section] = NSMutableIndexSet(index: index)
                    }
                    
                    // Since we don't care about the indexes remaining stable, we can directly remove objects here.
                    fetchedObjects!.removeAtIndex(fetchedObjects!.indexOf(object)!)
                    
                    changes.append(.Delete(object: object, indexPath: NSIndexPath(forRow: index, inSection: sectionIndex)))
                    continue
                }
                
                // Determine the new section for the object, optimizing to avoid this where possible.
                let changedValues = object.changedValuesForCurrentEvent().keys
                let newSection: Section
                if let sectionKeys = sectionKeys {
                    if sectionKeys.isSubsetOf(changedValues) {
                        newSection = sectionForObject(object)
                    } else {
                        newSection = section
                    }
                } else {
                    newSection = sectionForObject(object)
                }
                
                let insertRecord = (object: object, sectionIndex: Int?.Some(sectionIndex), index: Int?.Some(index))
                if section != newSection {
                    // Object has changed section.
                    objectSections![ObjectIdentifier(object)] = newSection
                    
                    if let _ = deleteIndexes[section] {
                        deleteIndexes[section]!.addIndex(index)
                    } else {
                        deleteIndexes[section] = NSMutableIndexSet(index: index)
                    }
                    
                    if let newSectionIndex = indexOfSection(newSection) {
                        sections![newSectionIndex].numberOfObjects += 1
                        sections![newSectionIndex].objects.append(object)
                        
                        if let _ = insertedObjects[newSection] {
                            insertedObjects[newSection]!.append(insertRecord)
                        } else {
                            insertedObjects[newSection] = [insertRecord]
                        }
                    } else {
                        makeSection(section: newSection, object: object)
                        
                        insertedSections.append(newSection)
                        insertedObjects[newSection] = [insertRecord]
                    }
                } else if let sortKeys = sortKeys where sortKeys.isSubsetOf(changedValues) {
                    // Object may have moved within the sort order.
                    if let _ = insertedObjects[section] {
                        insertedObjects[section]!.append(insertRecord)
                    } else {
                        insertedObjects[section] = [insertRecord]
                    }
                } else {
                    // Object has changed in some other way.
                    changes.append(.Update(object: object, indexPath: NSIndexPath(forRow: index, inSection: sectionIndex)))
                }

            } else if fetchRequest.predicate?.evaluateWithObject(object) ?? true {
                // Object previous did not, but now does, match the predicate. This becomes an insert.
                let section = sectionForObject(object)
                objectSections![ObjectIdentifier(object)] = section
                
                let insertRecord = (object: object, sectionIndex: Int?.None, index: Int?.None)
                if let sectionIndex = indexOfSection(section) {
                    sections![sectionIndex].numberOfObjects += 1
                    sections![sectionIndex].objects.append(object)
                    
                    if let _ = insertedObjects[section] {
                        insertedObjects[section]!.append(insertRecord)
                    } else {
                        insertedObjects[section] = [insertRecord]
                    }
                } else {
                    makeSection(section: section, object: object)
                    
                    insertedSections.append(section)
                    insertedObjects[section] = [insertRecord]
                }
                
                fetchedObjects!.append(object)
            }
        }
        
        // Use the complete list of objects to be removed, including those moving out of a section to another, to do a single-pass of all removes.
        let deleteSectionIndexes = NSMutableIndexSet()
        
        for (section, indexes) in deleteIndexes {
            let sectionIndex = sectionIndexes![section]!
            
            sections![sectionIndex].numberOfObjects -= indexes.count
            for index in indexes.reverse() {
                sections![sectionIndex].objects.removeAtIndex(index)
            }
            
            if sections![sectionIndex].numberOfObjects == 0 {
                deleteSectionIndexes.addIndex(sectionIndex)

                changes.append(.DeleteSection(sectionInfo: sections![sectionIndex], index: sectionIndex))
            }
        }
        
        for sectionIndex in deleteSectionIndexes.reverse() {
            sections!.removeAtIndex(sectionIndex)
        }
        
        if deleteSectionIndexes.count > 0 {
            sectionIndexes = [:]
        }
        
        // Now that all of the objects are in the right section, we need to sort any sections into which we've inserted objects, including those into which objects have been moved. We also need to sort the sections list itself if any sections were inserted.
        if let sortDescriptors = fetchRequest.sortDescriptors {
            for section in insertedObjects.keys {
                let sectionIndex = sectionIndexes![section]!
                
                sections![sectionIndex].objects = (sections![sectionIndex].objects as NSArray).sortedArrayUsingDescriptors(sortDescriptors) as! [Entity]
            }
            
            if insertedObjects.count > 0 {
                fetchedObjects = (fetchedObjects! as NSArray).sortedArrayUsingDescriptors(sortDescriptors) as? [Entity]
            }

        }

        if insertedSections.count > 0 {
            sortSections()
        }
        
        // Finally we can iterate the set of inserted and moved objects to generate the changes with the correct new indexes for them.
        for newSection in insertedSections {
            let newSectionIndex = indexOfSection(newSection)!
            
            changes.append(.InsertSection(sectionInfo: sections![newSectionIndex], newIndex: newSectionIndex))
        }
        
        for (newSection, insertRecords) in insertedObjects {
            let newSectionIndex = indexOfSection(newSection)!
            for (object, sectionIndex, index) in insertRecords {
                let newIndex = sections![newSectionIndex].objects.indexOf(object)!
                
                if let index = index, sectionIndex = sectionIndex {
                    changes.append(.Move(object: object, indexPath: NSIndexPath(forRow: index, inSection: sectionIndex), newIndexPath: NSIndexPath(forRow: newIndex, inSection: newSectionIndex)))
                } else {
                    changes.append(.Insert(object: object, newIndexPath: NSIndexPath(forRow: newIndex, inSection: newSectionIndex)))
                }
            }
        }
        
        handleChanges?(changes)
    }
    
}

public enum FetchedResultsChange<Section, Entity : NSManagedObject> {
    case InsertSection(sectionInfo: FetchedResultsSectionInfo<Section, Entity>, newIndex: Int)
    case DeleteSection(sectionInfo: FetchedResultsSectionInfo<Section, Entity>, index: Int)

    case Insert(object: Entity, newIndexPath: NSIndexPath)
    case Delete(object: Entity, indexPath: NSIndexPath)
    case Move(object: Entity, indexPath: NSIndexPath, newIndexPath: NSIndexPath)
    case Update(object: Entity, indexPath: NSIndexPath)
}

public struct  FetchedResultsSectionInfo<Section, Entity : NSManagedObject> {
    public private(set) var name: Section
    public private(set) var numberOfObjects: Int
    public private(set) var objects: [Entity]
    
    private init(name: Section) {
        self.name = name
        self.numberOfObjects = 0
        self.objects = []
    }
    
}