Objective C Trie Implementation

Trie.h

#import <Foundation/Foundation.h>


typedef void (^TrieBlock)(id obj, BOOL *stop);

@interface Trie : NSObject

@property (readonly) unichar value;
@property (getter = isLeaf) BOOL isLeaf;
@property (readonly, getter=isHead) BOOL head;
@property (readonly, weak) Trie *parent;

- (id) init: (unichar) value parent: (id) parent isLeaf: (BOOL) isLeaf;
- (id) initHead;
- (NSString *) toString;
- (NSString *) valueString;

- (Trie *) getString: (NSString *) string;
- (Trie *) addString: (NSString *) string;
- (void) removeString: (NSString *) string;

- (Trie *) getUnichar: (unichar) value;

- (BOOL) hasString: (NSString *) string;
- (BOOL) hasChildren;

- (void) enumerateStringsUsingBlock: (TrieBlock) block;
- (void) enumerateNodesUsingBlock: (TrieBlock) block;

- (NSEnumerator<Trie *> *) children;

@end

Trie.m

#import "Trie.h"


NSString *stringFromUnichar(unichar value){
    return [[NSString alloc] initWithCharacters: &value length: 1];
}

@implementation Trie {
    NSMutableDictionary *_children;
}

@synthesize isLeaf = _isLeaf;
@synthesize value = _value;
@synthesize parent = _parent;


- (id) init: (unichar) value parent: (id) parent isLeaf: (BOOL) isLeaf {
    self = [super init];

    if(self){
        _isLeaf = isLeaf;
        _value = value;
        _parent = parent;
        _children = [NSMutableDictionary new];
    }

    return self;
}

- (id) initHead {
    return [self init: '\0' parent: nil isLeaf: NO];
}

- (BOOL) isHead {
    return [self parent] == nil;
}

- (NSString *) valueString {
    return stringFromUnichar([self value]);
}

- (Trie *) addUnichar: (unichar) value isLeaf: (bool) isLeaf {
    NSString *key = stringFromUnichar(value);

    Trie *child = _children[key];
    if(!child) {
        child = [[Trie alloc] init: value parent: self isLeaf: isLeaf];
        _children[key] = child;
    }
    [child setIsLeaf: child.isLeaf || isLeaf];
    return child;
}

- (Trie *) addString: (NSString *) string {
    unichar str[[string length] + 1];
    [string getCharacters: str range: NSMakeRange(0, [string length])];

    Trie *active = self;
    for(size_t i = 0; i < [string length]; i++){
        active = [active addUnichar: str[i] isLeaf: NO];
    }

    if(active){
        [active setIsLeaf: YES];
    }

    return active;
}

- (void) removeUnichar: (unichar) value {
    Trie *node = [self getUnichar: value];
    if(node){
        [node setIsLeaf: NO];

        if(![node hasChildren]){
            [_children removeObjectForKey: stringFromUnichar(value)];
        }
    }
}

- (void) removeString: (NSString *) string {
    Trie *active = [self getString: string];
    [active setIsLeaf: NO];

    while(active && ![active isHead] && ![active isLeaf]){
        Trie *parent = [active parent];
        [parent removeUnichar: [active value]];

        if([parent hasChildren]){
            return;
        }
        active = parent;
    }
}

- (NSString *) toString {
    NSMutableArray *arr = [[NSMutableArray alloc] initWithCapacity: 32];

    Trie *active = self;
    while(active && ![active isHead]){
        [arr addObject: active];
        active = [active parent];
    }

    unichar str[[arr count] + 1];
    size_t i = 0;
    for(Trie *c in [arr reverseObjectEnumerator]){
        str[i++] = [c value];

        if(i > sizeof(str) - 1)
            break;
    }

    NSString *result = [[NSString alloc] initWithCharacters: str 
                                                     length: [arr count]]; 

    return result;
}

- (Trie *) getUnichar: (unichar) value {
    return _children[stringFromUnichar(value)];
}

- (Trie *) getString: (NSString *) string {
    Trie *active = self;
    for(size_t i = 0; i < [string length] && active != nil; i++){
        active = [active getUnichar: [string characterAtIndex: i]];
    }
    return active;
}

- (BOOL) hasString: (NSString *) string {
    Trie *node = [self getString: string];
    return node && [node isLeaf];
}

- (BOOL) hasChildren {
    return [_children count] > 0;
}

- (NSEnumerator<Trie *> *) children {
    return [_children objectEnumerator];
}

- (void) __enumerateNodesUsingBlock: (TrieBlock) block stop: (BOOL *) stop {
    for(Trie *child in [self children]){
        if(*stop) return;

        block(child, stop);
        [child __enumerateNodesUsingBlock: block stop: stop];
    }
}

- (void) enumerateNodesUsingBlock: (TrieBlock) block { 
    BOOL stop = NO; 
    [self __enumerateNodesUsingBlock: block stop: &stop];
}

- (void) enumerateStringsUsingBlock: (TrieBlock) block {
    [self enumerateNodesUsingBlock: ^(id obj, BOOL *stop){
        if([obj isLeaf]){
            block(obj, stop);
        }
    }];
}

@end