ReferenceCache.swift

/// The result of a binary search.
struct BinarySearchResult<Index> {
	/// The index for the searched key.
	///
	/// Use this to either insert new values while maintaining the sort order of the collection, or to lookup the value if `isPresent` is true.
	var index: Index
	
	/// Whether the key was found in the collection.
	///
	/// When performing a binary search, even if the value is not found in the collection, an index is returned which can be used to insert a new item at the correct location.
	var isPresent: Bool
}

extension Collection where Index: BinaryInteger {
	/// Search for an index for a given key.
	///
	/// Searches through a sorted collection to find the index where the item should be inserted, or already exists.
	///
	/// This method assumes that the collection is already sorted by the given key.
	func binarySearch<Key: Comparable>(for key: Key, using: (Element) -> Key) -> BinarySearchResult<Index> {
		if self.endIndex == self.startIndex {
			return BinarySearchResult(index: startIndex, isPresent: false)
		}
		
		let midIndex = startIndex + (endIndex - startIndex) / Index(exactly: 2)!
		let midKey = using(self[midIndex])
		
		// Is the search key in the left half?
		if midKey > key {
			return self[startIndex..<midIndex].binarySearch(for: key, using: using)
			// Is the search key in the right half?
		} else if midKey < key {
			return self[(midIndex + 1)..<endIndex].binarySearch(for: key, using: using)
			// If we get here, then we've found the search key!
		} else {
			return BinarySearchResult(index: midIndex, isPresent: true)
		}
	}
}

extension Collection where Index: BinaryInteger, Element: Comparable {
	/// Search for an index for a given element.
	///
	/// Searches through a sorted collection to find the index where the item should be inserted, or already exists.
	///
	/// This method assumes that the collection is already sorted.
	func binarySearch(for key: Element) -> BinarySearchResult<Index> {
		return self.binarySearch(for: key, using: { $0 })
	}
}


/// A cache that removes elements once they are no longer used (retained).
class ReferenceCache<Key: Hashable, Value: AnyObject> {
	private struct Content {
		var hash: Int
		
		struct Entry {
			var key: Key
			var value: Value
		}
		var entries: [Entry]
	}
	
	private var content: [Content] = []
	
	subscript(key: Key) -> Value? {
		get {
			let result = content.binarySearch(for: key.hashValue, using: { $0.hash })
			guard result.isPresent else { return nil }
			return content[result.index].entries.first(where: { $0.key == key })?.value
		}
		set(newValue) {
			let result = content.binarySearch(for: key.hashValue, using: { $0.hash })
			let index = result.index
			
			if result.isPresent {
				if let newValue = newValue {
					if let entryIndex = content[index].entries.firstIndex(where: { $0.key == key }) {
						content[index].entries[entryIndex].value = newValue
					} else {
						content[index].entries.append(Content.Entry(key: key, value: newValue))
					}
				} else {
					content[index].entries.removeAll(where: { $0.key == key })
					if content[index].entries.isEmpty {
						content.remove(at: index)
					}
				}
			} else {
				guard let newValue = newValue else { return }
				content.insert(Content(hash: key.hashValue, entries: [Content.Entry(key: key, value: newValue)]), at: index)
			}
		}
	}
	
	func clean() {
		for index in content.indices.reversed() {
			for entryIndex in content[index].entries.indices.reversed() {
				if isKnownUniquelyReferenced(&self.content[index].entries[entryIndex].value) {
					self.content[index].entries.remove(at: entryIndex)
					
					if self.content[index].entries.isEmpty {
						self.content[index].entries.remove(at: entryIndex)
					}
				}
			}
		}
	}
}