Represent a set of index paths

SSMutableIndexPathSet.h

//
//  SSMutableIndexPathSet.h
//  MutableIndexPathSet Test
//
//  Created by Colin Barrett on 11/19/10.
//  Copyright 2010 __MyCompanyName__. All rights reserved.
//

#import <Foundation/Foundation.h>

/*
 If an NSIndexPath "represents the path to a specific node in a tree of nested array collections", then an index path set is a class that represents a collection of these
 tree paths and allows you to efficiently -- it uses a prefix tree, or trie, internally -- ask questions about the paths it contains. This is especially useful when
 dealing with a class like UITableView, which wants to know things like "how many rows in a particular section?" and which uses index paths extensively in its API to
 represent specific cells.
*/

@interface SSMutableIndexPathSet : NSObject {
    CFMutableDictionaryRef trie;
}

- (id)init;
- (void)addIndexPath:(NSIndexPath *)indexPath;
- (void)removeIndexPath:(NSIndexPath *)indexPath;

- (NSIndexSet *)childrenIndexSetAtIndexPath:(NSIndexPath *)indexPath; // ex: if you wanted all the subindexes of top level index 2, pass {2}.
- (NSIndexSet *)childrenIndexSet; // e.g. the set of all indexes at position 0 in the represented index paths. Equivalent to above with an empty index path.

- (SSMutableIndexPathSet *)childIndexPathSetAtIndexPath:(NSIndexPath *)indexPath; // prune the tree.

@end

@interface NSIndexPath (SSMIPS)
+ (id)ss_indexPath; // an empty index path
@end

SSMutableIndexPathSet.m

//
//  SSMutableIndexPathSet.m
//  MutableIndexPathSet Test
//
//  Created by Colin Barrett on 11/19/10.
//  Copyright 2010 __MyCompanyName__. All rights reserved.
//

#import "SSMutableIndexPathSet.h"

CFStringRef indexCopyDescription(const void *value)
{
    NSUInteger idx = (NSUInteger)value;
    return (CFStringRef)[[NSString alloc] initWithFormat:@"%u", idx];
}

const CFDictionaryKeyCallBacks kIndexDictionaryKeyCallbacks = {
    0, // version
    NULL, // retain
    NULL, // release
    indexCopyDescription, // descrption 
    NULL, // hash
    NULL // equal
};

@implementation NSIndexPath (SSMIPS)

+ (id)ss_indexPath
{
    return [[[self alloc] initWithIndexes:NULL length:0] autorelease];
}

@end


@interface NSIndexPath (SSMIPSInternal)
- (NSIndexPath *)ss_indexPathByRemovingFirstIndex;
@end

@implementation NSIndexPath (SSMIPSInternal)

- (NSIndexPath *)ss_indexPathByRemovingFirstIndex
{
    //NSLog(@"removing from: %@", self);
    NSUInteger length = [self length];
    NSParameterAssert(length > 1);
    NSUInteger *indexes = malloc(sizeof(NSUInteger) * length);
    [self getIndexes:indexes];
    NSIndexPath *childIndexPath = [[self class] indexPathWithIndexes:indexes+1 length:length-1];
    free(indexes);
    
    //NSLog(@"new: %@", childIndexPath);
    return childIndexPath;
}

@end

@implementation SSMutableIndexPathSet

- (id)init
{
    if ((self = [super init])) {
        trie = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, &kIndexDictionaryKeyCallbacks, &kCFTypeDictionaryValueCallBacks);
    }
    return self;
}

- (id)_initWithTrie:(CFMutableDictionaryRef)inTrie
{
    if ((self = [super init])) {
        CFRetain(inTrie);
        trie = inTrie;
    }
    return self;
}

- (void)dealloc
{
    CFRelease(trie);
    
    [super dealloc];
}

- (void)addIndexPath:(NSIndexPath *)indexPath
{
    CFMutableDictionaryRef current = trie;
    
    for (NSUInteger i = 0; i < [indexPath length]; i++) {
        NSUInteger idx = [indexPath indexAtPosition:i];        
        CFMutableDictionaryRef child = (CFMutableDictionaryRef)CFDictionaryGetValue(current, (const void *)idx);
        if (!child) {
            child = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, &kIndexDictionaryKeyCallbacks, &kCFTypeDictionaryValueCallBacks);
            CFDictionarySetValue(current, (const void *)idx, child);
            CFRelease(child);
        }
        current = child;
    }

    //NSLog(@"add: %@", [(NSString *)CFCopyDescription(trie) autorelease]);
}

- (NSIndexPath *)_innerRemoveIndexPath:(NSIndexPath *)indexPath fromTrie:(CFMutableDictionaryRef)inTrie
{
    NSIndexPath *current = indexPath;
    
    // Depth first down the tree to the leaf
    if ([current length] > 1) {
        CFMutableDictionaryRef childTrie = (CFMutableDictionaryRef)CFDictionaryGetValue(inTrie, (const void *)[current indexAtPosition:0]);
        NSParameterAssert(childTrie);
        current = [self _innerRemoveIndexPath:[current ss_indexPathByRemovingFirstIndex] fromTrie:childTrie];
    }
    
    // As we come back up, prune any empty dictionaries we find -- note that leaves have empty dictionaries.
    NSParameterAssert([current length] == 1);
    NSUInteger idx = [current indexAtPosition:0];
    CFMutableDictionaryRef child = (CFMutableDictionaryRef)CFDictionaryGetValue(inTrie, (const void *)idx);
    if (child && CFDictionaryGetCount(child) == 0) { 
        CFDictionaryRemoveValue(inTrie, (const void *)idx);
    }
    
    return current;
}

- (void)removeIndexPath:(NSIndexPath *)indexPath
{
    if (indexPath) {
        [self _innerRemoveIndexPath:indexPath fromTrie:trie];
    }
    
    //NSLog(@"remove: %@", [(NSString *)CFCopyDescription(trie) autorelease]);
}

- (NSIndexSet *)childrenIndexSetAtIndexPath:(NSIndexPath *)indexPath
{
    CFMutableDictionaryRef child = trie;
    
    for (NSUInteger i = 0; i < [indexPath length]; i++) {
        NSUInteger idx = [indexPath indexAtPosition:i];
        child = (CFMutableDictionaryRef)CFDictionaryGetValue(child, (const void *)idx);
        if (!child) {
            return nil;
        }
    }
    
    NSMutableIndexSet *indexSet = [NSMutableIndexSet indexSet];
    
    NSUInteger numberOfIndexes = CFDictionaryGetCount(child);
    NSUInteger *indexes = malloc(sizeof(NSUInteger) * numberOfIndexes);
    
    CFDictionaryGetKeysAndValues(child, (const void **)indexes, NULL);
    for (NSUInteger i = 0; i < numberOfIndexes; i++) {
        [indexSet addIndex:indexes[i]];
    }
    
    free(indexes);
    
    return [[indexSet copy] autorelease];    
}

- (NSIndexSet *)childrenIndexSet
{
    return [self childrenIndexSetAtIndexPath:[NSIndexPath ss_indexPath]];
}

- (SSMutableIndexPathSet *)childIndexPathSetAtIndexPath:(NSIndexPath *)indexPath
{
    CFMutableDictionaryRef child = trie;
    for (NSUInteger i = 0; i < [indexPath length]; i++) {
        NSUInteger idx = [indexPath indexAtPosition:i];
        child = (CFMutableDictionaryRef)CFDictionaryGetValue(child, (const void *)idx);
        if (!child) {
            return nil;
        }
    }
    
    return [[[SSMutableIndexPathSet alloc] _initWithTrie:child] autorelease];
}

- (NSString *)description
{
    return [(NSString *)CFCopyDescription(trie) autorelease]; // XXX temp
}

@end