diff options
Diffstat (limited to 'Firestore/third_party/Immutable/FSTTreeSortedDictionary.m')
| -rw-r--r-- | Firestore/third_party/Immutable/FSTTreeSortedDictionary.m | 382 |
1 files changed, 382 insertions, 0 deletions
diff --git a/Firestore/third_party/Immutable/FSTTreeSortedDictionary.m b/Firestore/third_party/Immutable/FSTTreeSortedDictionary.m new file mode 100644 index 0000000..4059951 --- /dev/null +++ b/Firestore/third_party/Immutable/FSTTreeSortedDictionary.m @@ -0,0 +1,382 @@ +#import "FSTTreeSortedDictionary.h" + +#import "FSTLLRBEmptyNode.h" +#import "FSTLLRBValueNode.h" +#import "FSTTreeSortedDictionaryEnumerator.h" + +NS_ASSUME_NONNULL_BEGIN + +@interface FSTTreeSortedDictionary () + +- (FSTTreeSortedDictionary *)dictionaryBySettingObject:(id)aValue forKey:(id)aKey; + +@property(nonatomic, strong) id<FSTLLRBNode> root; +@property(nonatomic, copy, readwrite) NSComparator comparator; +@end + +@implementation FSTTreeSortedDictionary + ++ (FSTTreeSortedDictionary *)dictionaryWithDictionary:(NSDictionary *)dictionary + comparator:(NSComparator)comparator { + __block FSTTreeSortedDictionary *dict = + [[FSTTreeSortedDictionary alloc] initWithComparator:comparator]; + [dictionary enumerateKeysAndObjectsUsingBlock:^(id key, id obj, BOOL *stop) { + dict = [dict dictionaryBySettingObject:obj forKey:key]; + }]; + return dict; +} + +- (id)initWithComparator:(NSComparator)aComparator { + return [self initWithComparator:aComparator withRoot:[FSTLLRBEmptyNode emptyNode]]; +} + +// Designated initializer. +- (id)initWithComparator:(NSComparator)aComparator withRoot:(id<FSTLLRBNode>)aRoot { + self = [super init]; + if (self) { + self.root = aRoot; + self.comparator = aComparator; + } + return self; +} + +/** + * Returns a copy of the map, with the specified key/value added or replaced. + */ +- (FSTTreeSortedDictionary *)dictionaryBySettingObject:(id)aValue forKey:(id)aKey { + return [[FSTTreeSortedDictionary alloc] + initWithComparator:self.comparator + withRoot:[[self.root insertKey:aKey forValue:aValue withComparator:self.comparator] + copyWith:nil + withValue:nil + withColor:FSTLLRBColorBlack + withLeft:nil + withRight:nil]]; +} + +- (FSTTreeSortedDictionary *)dictionaryByRemovingObjectForKey:(id)aKey { + // Remove is somewhat expensive even if the key doesn't exist (the tree does rebalancing and + // stuff). So avoid it. + if (![self containsKey:aKey]) { + return self; + } else { + return [[FSTTreeSortedDictionary alloc] + initWithComparator:self.comparator + withRoot:[[self.root remove:aKey withComparator:self.comparator] + copyWith:nil + withValue:nil + withColor:FSTLLRBColorBlack + withLeft:nil + withRight:nil]]; + } +} + +- (nullable id)objectForKey:(id)key { + NSComparisonResult cmp; + id<FSTLLRBNode> node = self.root; + while (![node isEmpty]) { + cmp = self.comparator(key, node.key); + if (cmp == NSOrderedSame) { + return node.value; + } else if (cmp == NSOrderedAscending) { + node = node.left; + } else { + node = node.right; + } + } + return nil; +} + +- (nullable id)predecessorKey:(id)key { + NSComparisonResult cmp; + id<FSTLLRBNode> node = self.root; + id<FSTLLRBNode> rightParent = nil; + while (![node isEmpty]) { + cmp = self.comparator(key, node.key); + if (cmp == NSOrderedSame) { + if (![node.left isEmpty]) { + node = node.left; + while (![node.right isEmpty]) { + node = node.right; + } + return node.key; + } else if (rightParent != nil) { + return rightParent.key; + } else { + return nil; + } + } else if (cmp == NSOrderedAscending) { + node = node.left; + } else if (cmp == NSOrderedDescending) { + rightParent = node; + node = node.right; + } + } + @throw [NSException exceptionWithName:@"NonexistentKey" + reason:@"getPredecessorKey called with nonexistent key." + userInfo:@{@"key" : [key description]}]; +} + +- (NSUInteger)indexOfKey:(id)key { + NSUInteger prunedNodes = 0; + id<FSTLLRBNode> node = self.root; + while (![node isEmpty]) { + NSComparisonResult cmp = self.comparator(key, node.key); + if (cmp == NSOrderedSame) { + return prunedNodes + node.left.count; + } else if (cmp == NSOrderedAscending) { + node = node.left; + } else if (cmp == NSOrderedDescending) { + prunedNodes += node.left.count + 1; + node = node.right; + } + } + return NSNotFound; +} + +- (BOOL)isEmpty { + return [self.root isEmpty]; +} + +- (NSUInteger)count { + return [self.root count]; +} + +- (id)minKey { + return [self.root minKey]; +} + +- (id)maxKey { + return [self.root maxKey]; +} + +- (void)enumerateKeysAndObjectsUsingBlock:(void (^)(id, id, BOOL *))block { + [self enumerateKeysAndObjectsReverse:NO usingBlock:block]; +} + +- (void)enumerateKeysAndObjectsReverse:(BOOL)reverse usingBlock:(void (^)(id, id, BOOL *))block { + if (reverse) { + __block BOOL stop = NO; + [self.root reverseTraversal:^BOOL(id key, id value) { + block(key, value, &stop); + return stop; + }]; + } else { + __block BOOL stop = NO; + [self.root inorderTraversal:^BOOL(id key, id value) { + block(key, value, &stop); + return stop; + }]; + } +} + +- (BOOL)containsKey:(id)key { + return ([self objectForKey:key] != nil); +} + +- (NSEnumerator *)keyEnumerator { + return [[FSTTreeSortedDictionaryEnumerator alloc] initWithImmutableSortedDictionary:self + startKey:nil + endKey:nil + isReverse:NO]; +} + +- (NSEnumerator *)keyEnumeratorFrom:(id)startKey { + return [[FSTTreeSortedDictionaryEnumerator alloc] initWithImmutableSortedDictionary:self + startKey:startKey + endKey:nil + isReverse:NO]; +} + +- (NSEnumerator *)keyEnumeratorFrom:(id)startKey to:(nullable id)endKey { + return [[FSTTreeSortedDictionaryEnumerator alloc] initWithImmutableSortedDictionary:self + startKey:startKey + endKey:endKey + isReverse:NO]; +} + +- (NSEnumerator *)reverseKeyEnumerator { + return [[FSTTreeSortedDictionaryEnumerator alloc] initWithImmutableSortedDictionary:self + startKey:nil + endKey:nil + isReverse:YES]; +} + +- (NSEnumerator *)reverseKeyEnumeratorFrom:(id)startKey { + return [[FSTTreeSortedDictionaryEnumerator alloc] initWithImmutableSortedDictionary:self + startKey:startKey + endKey:nil + isReverse:YES]; +} + +#pragma mark - +#pragma mark Tree Builder + +// Code to efficiently build a red black tree. + +typedef struct { + unsigned int bits; + unsigned short count; + unsigned short current; +} Base12List; + +unsigned int LogBase2(unsigned int num) { + return (unsigned int)(log(num) / log(2)); +} + +/** + * Works like an iterator, so it moves to the next bit. Do not call more than list->count times. + * @return whether or not the next bit is a 1 in base {1,2}. + */ +BOOL Base12ListNext(Base12List *list) { + BOOL result = !(list->bits & (0x1 << list->current)); + list->current--; + return result; +} + +static inline unsigned BitMask(int x) { + return (x >= sizeof(unsigned) * CHAR_BIT) ? (unsigned)-1 : (1U << x) - 1; +} + +/** + * We represent the base{1,2} number as the combination of a binary number and a number of bits that + * we care about. We iterate backwards, from most significant bit to least, to build up the llrb + * nodes. 0 base 2 => 1 base {1,2}, 1 base 2 => 2 base {1,2} + */ +Base12List *NewBase12List(unsigned int length) { + size_t sz = sizeof(Base12List); + Base12List *list = calloc(1, sz); + // Calculate the number of bits that we care about + list->count = (unsigned short)LogBase2(length + 1); + unsigned int mask = BitMask(list->count); + list->bits = (length + 1) & mask; + list->current = list->count - 1; + return list; +} + +void FreeBase12List(Base12List *list) { + free(list); +} + ++ (id<FSTLLRBNode>)buildBalancedTree:(NSArray *)keys + dictionary:(NSDictionary *)dictionary + subArrayStartIndex:(NSUInteger)startIndex + length:(NSUInteger)length { + length = MIN(keys.count - startIndex, length); // Bound length by the actual length of the array + if (length == 0) { + return nil; + } else if (length == 1) { + id key = keys[startIndex]; + return [[FSTLLRBValueNode alloc] initWithKey:key + withValue:dictionary[key] + withColor:FSTLLRBColorBlack + withLeft:nil + withRight:nil]; + } else { + NSUInteger middle = length / 2; + id<FSTLLRBNode> left = [FSTTreeSortedDictionary buildBalancedTree:keys + dictionary:dictionary + subArrayStartIndex:startIndex + length:middle]; + id<FSTLLRBNode> right = [FSTTreeSortedDictionary buildBalancedTree:keys + dictionary:dictionary + subArrayStartIndex:(startIndex + middle + 1) + length:middle]; + id key = keys[startIndex + middle]; + return [[FSTLLRBValueNode alloc] initWithKey:key + withValue:dictionary[key] + withColor:FSTLLRBColorBlack + withLeft:left + withRight:right]; + } +} + ++ (nullable id<FSTLLRBNode>)rootFrom12List:(Base12List *)base12List + keyList:(NSArray *)keyList + dictionary:(NSDictionary *)dictionary { + __block FSTLLRBValueNode *root = nil; + __block FSTLLRBValueNode *node = nil; + __block NSUInteger index = keyList.count; + + void (^buildPennant)(FSTLLRBColor, NSUInteger) = ^(FSTLLRBColor color, NSUInteger chunkSize) { + NSUInteger startIndex = index - chunkSize + 1; + index -= chunkSize; + id key = keyList[index]; + FSTLLRBValueNode *childTree = [self buildBalancedTree:keyList + dictionary:dictionary + subArrayStartIndex:startIndex + length:(chunkSize - 1)]; + FSTLLRBValueNode *pennant = [[FSTLLRBValueNode alloc] initWithKey:key + withValue:dictionary[key] + withColor:color + withLeft:nil + withRight:childTree]; + if (node) { + // This is the only place this property is set. + node.left = pennant; + node = pennant; + } else { + root = pennant; + node = pennant; + } + }; + + for (int i = 0; i < base12List->count; ++i) { + BOOL isOne = Base12ListNext(base12List); + NSUInteger chunkSize = (NSUInteger)pow(2.0, base12List->count - (i + 1)); + if (isOne) { + buildPennant(FSTLLRBColorBlack, chunkSize); + } else { + buildPennant(FSTLLRBColorBlack, chunkSize); + buildPennant(FSTLLRBColorRed, chunkSize); + } + } + return root; +} + +/** + * Uses the algorithm linked here: + * http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.46.1458 + */ ++ (FSTImmutableSortedDictionary *)fromDictionary:(NSDictionary *)dictionary + withComparator:(NSComparator)comparator { + // Steps: + // 0. Sort the array + // 1. Calculate the 1-2 number + // 2. Build From 1-2 number + // 0. for each digit in 1-2 number + // 0. calculate chunk size + // 1. build 1 or 2 pennants of that size + // 2. attach pennants and update node pointer + // 1. return root + NSMutableArray *sortedKeyList = [NSMutableArray arrayWithCapacity:dictionary.count]; + [dictionary enumerateKeysAndObjectsUsingBlock:^(id key, id obj, BOOL *stop) { + [sortedKeyList addObject:key]; + }]; + [sortedKeyList sortUsingComparator:comparator]; + + [sortedKeyList enumerateObjectsUsingBlock:^(id obj, NSUInteger idx, BOOL *stop) { + if (idx > 0) { + if (comparator(sortedKeyList[idx - 1], obj) != NSOrderedAscending) { + [NSException raise:NSInvalidArgumentException + format: + @"Can't create FSTImmutableSortedDictionary " + @"with keys with same ordering!"]; + } + } + }]; + + Base12List *list = NewBase12List((unsigned int)sortedKeyList.count); + id<FSTLLRBNode> root = [self rootFrom12List:list keyList:sortedKeyList dictionary:dictionary]; + FreeBase12List(list); + + if (root != nil) { + return [[FSTTreeSortedDictionary alloc] initWithComparator:comparator withRoot:root]; + } else { + return [[FSTTreeSortedDictionary alloc] initWithComparator:comparator]; + } +} + +@end + +NS_ASSUME_NONNULL_END |