Flatten NSDictionary and Associated References

This is a great exercise using Objective-C associated references to apply instance variables to a category. I first started by creating a utility class but also thought it would be cleaner as a NSDictionary category. For clarity on flattening the NSDictionary, lets first take a look at the utility class version before we get into associated references.

Lets create a test first so we know what to expect.

#import <senTestingKit/SenTestingKit.h>
#import "NSDictionaryUtils.h"
 
@interface NSDictionaryUtilsTests : SenTestCase {
 
}
 
- (void)testFlatten;
 
@end
 
@implementation NSDictionaryUtilsTests
 
- (void)setUp {
	[super setUp];
}
 
- (void)tearDown {
	[super tearDown];
}
 
- (void)testFlatten {
	NSDictionary *collection1 = @{ @"firstName" : @"Foo", @"lastName" : @"Bar", @"middleName" : @"middle" };
	NSArray *collection2 = @[@"134234", @"234323"];
	NSSet *collection3 = [[NSSet alloc] initWithObjects:@"fooset", @"barset", nil];
 
	NSDictionary *flatten:dictionaryToFlatten = [[NSDictionary alloc] initWithObjectsAndKeys:collection1, @"key1", collection2, @"key2", collection3, @"key3", nil];
 
        NSDictionaryUtils *dictUtils = [[NSDictionaryUtils alloc] init];
	NSDate *startTime = [NSDate date];
	NSMutableArray *testArray = [dictUtils flatten:dictionaryToFlatten];
	NSDate *endTime = [NSDate date];
 
	NSTimeInterval executionTime = [endTime timeIntervalSinceDate:startTime];
	NSLog(@"execution took approx: %2f seconds", executionTime );
 
	STAssertTrue( [testArray count] == 7, @"should be equal");
}
 
@end

 

#import <foundation/Foundation.h>
 
@interface NSDictionaryUtils : NSObject
 
- (NSMutableArray *)flatten:(NSDictionary *)dictionary;
 
@end
#import "NSDictionaryUtils.h"
 
@interface NSDictionaryUtils() {
	NSMutableArray *_tmpArray;
}
 
- (void)flattenWithRecursion:(id)object;
- (void)enumerate:(id)object depth:(int)depth parent:(id)parent;
 
@end
 
@implementation NSDictionaryUtils
 
- (NSMutableArray *)flatten:(NSDictionary *)dictionary {
	_tmpArray = [NSMutableArray array];
	[self flattenWithRecursion:dictionary];
	return _tmpArray;
}
 
- (void)flattenWithRecursion:(id)object {
	[self enumerate:object depth:0 parent:nil];
}
 
- (void)enumerate:(id)object depth:(int)depth parent:(id)parent {
 
	if( [object isKindOfClass:[NSDictionary class]] ) {
		for( NSString * key in [object allKeys] ) {
			id child = [object objectForKey:key];
			[self enumerate:child depth:depth+1 parent:object];
		}
 
	} else if( [object isKindOfClass:[NSArray class]] ) {
 
		for( id child in object ) {
			[self enumerate:child depth:depth+1 parent:object];
		}
 
	} else if( [object isKindOfClass:[NSSet class]] ) {
 
		for( id child in object ) {
			[self enumerate:child depth:depth+1 parent:object];
		}
 
	} else{
		//	not a collection/container it has ended
		//NSLog(@"Node: %@  depth: %d",[object description],depth);
		[_tmpArray addObject:object];
	}
}
 
@end

As you can see – (void)enumerate:depth:parent does all of the heavy lifting with recursion. It calls itself if the class is a collection type to enumerate through it again to get to the bottom most node. Also notice I have to use an instance variable here called _tmpArray to store all of the values while we still enumerate through other possible collections in the dictionary. Since we’re going to try to add this to a category of NSDictionary, we’ll have to use an associated reference variable as _tmpArray in our category since categories cannot have instance variables and we shouldn’t be extending NSDictionary.

Enter Associated References

Associated references are part of the Objective-c runtime that allow you to have property like functionality to a category. Lets take a look at the same bit of code but this time in a category for NSDictionary.

#import <foundation/Foundation.h>
 
@interface NSDictionary (Flatten)
 
- (NSMutableArray *)flattenedArray;
 
@end

 

#import "NSDictionary+Flatten.h"
#import <objc/runtime.h>
 
static char flattenedArrayKey;
 
@interface NSDictionary (FlattenPrivate)
 
- (void)flattenWithRecursion:(id)object;
- (void)enumerate:(id)object depth:(int)depth parent:(id)parent;
 
@end
 
 
@implementation NSDictionary (Flatten)
 
- (NSMutableArray *)flattenedArray {
	NSMutableArray *initArray = [[NSMutableArray alloc] init];
	objc_setAssociatedObject( self, &flattenedArrayKey, initArray, OBJC_ASSOCIATION_RETAIN );
	[self flattenWithRecursion:self];
	return (NSMutableArray *)objc_getAssociatedObject( self, &flattenedArrayKey );
}
 
- (void)flattenWithRecursion:(id)object {
	[self enumerate:object depth:0 parent:nil];
}
 
- (void)conquer:(id)object depth:(int)depth parent:(id)parent {
 
	if( [object isKindOfClass:[NSDictionary class]] ) {
 
		for( NSString * key in [object allKeys] ) {
			id child = [object objectForKey:key];
			[self enumerate:child depth:depth+1 parent:object];
		}
 
 
	} else if( [object isKindOfClass:[NSArray class]] ) {
 
		for( id child in object ) {
			[self enumerate:child depth:depth+1 parent:object];
		}
 
	} else if( [object isKindOfClass:[NSSet class]] ) {
 
		for( id child in object ) {
			[self enumerate:child depth:depth+1 parent:object];
		}
	}
	else{
		//	not a collection/container it has ended
		//NSLog(@"Node: %@  depth: %d",[object description],depth);
		NSMutableArray *assocObject = (NSMutableArray *)objc_getAssociatedObject(self, &flattenedArrayKey);
		[assocObject addObject:object];
		objc_setAssociatedObject( self, &flattenedArrayKey, assocObject, OBJC_ASSOCIATION_RETAIN );
	}
}
 
@end

Here are a few things that are worth noting:

  1. You have to import the Objective-C runtime
  2. We have to create a static var of type char as that’s what the runtime uses to keep a track of the associated reference.
  3. The Objective-C runtime is written in C, so you’ll see the C method call to objc_setAssociatedObject() and objc_getAssociatedObject()
  4. We have to use the OBJC_ASSOCIATION_RETAIN for the object behavior

It can get confusing

I’d like to show you a sample of using an associated reference that could look confusing if you’re not used to it.

NSMutableDictionary *myDictionary = [[NSMutableDictionary alloc] initWithObjectsAndKeys:@"foo", @"key", nil];
NSMutableArray *myArray = [[NSMutableArray alloc] init];
obj_setAssociatedObject( myDictionary, &myRef, myArray, OBJC_ASSOCIATION_RETAIN );

As you can see, the category implementation is a little cleaner because we use self but this can be done inline anywhere.

That’s it!. Full category source code is available on Github.