Immutable AVL Tree in Objective-C

AvlNode.h

//
//  AvlNode.h
//  AVL tree / Node
//
//  Created by Chris Cavanagh on 1/14/14.
//

#import <Foundation/Foundation.h>

typedef int (^EqualityComparer)( id v1, id v2 );

@interface AvlNode : NSObject

+ (AvlNode *) empty;

- (BOOL) isEmpty;

- (AvlNode *) left;

- (AvlNode *) right;

- (int) count;

- (int) balance;

- (int) depth;

- (AvlNode *) searchNode:(id)value
				comparer:(EqualityComparer)comparer;

- (AvlNode *) insertIntoNew:(int)index
					  value:(id)val;

- (AvlNode *) insertIntoNew:(id)val
				   comparer:(EqualityComparer)comparer;

- (AvlNode *) setItem:(int)index
				value:(id)val;

- (AvlNode *) getNodeAt:(int)index;

- (AvlNode *) removeFromNew:(int)index
					  found:(BOOL *)found;

- (AvlNode *) removeFromNew:(id)val
				   comparer:(EqualityComparer)comparer
					  found:(BOOL *)found;

- (NSEnumerator *)getEnumerator:(BOOL)reverse;

@property (nonatomic, strong) id value;

@end

AvlNode.m

//
//  AvlNode.m
//  AVL Tree / Node
//
//  Created by Chris Cavanagh on 1/14/14.
//

#import "AvlNode.h"

@class AvlNode;

@interface NullNode : AvlNode
@end

@interface MutableAvlNode : AvlNode
@end

@interface AvlNodeEnumerator : NSEnumerator

- (id) initWithNode:(AvlNode *)node
			reverse:(BOOL)reverse;

@end

@interface AvlNode ()
{
	AvlNode *_left;
	AvlNode *_right;
	int _count;
	int _depth;
}

@end

@implementation AvlNode

static AvlNode *_empty;

+ (void) initialize
{
	if ( [self class] == [AvlNode class] )
	{
		_empty = [NullNode new];
	}
}

+ (AvlNode *) empty
{
	return _empty;
}

- (BOOL) isEmpty
{
	return NO;
}

- (AvlNode *) left
{
	return _left;
}

- (AvlNode *) right
{
	return _right;
}

- (int) count
{
	return _count;
}

- (int) balance
{
	if ( [self isEmpty] ) return 0;
	return _left->_depth - _right->_depth;
}

- (int) depth
{
	return _depth;
}

- (id) init
{
	if ( ( self = [super init] ) )
	{
		_right = [AvlNode empty];
		_left = [AvlNode empty];
	}

	return self;
}

- (id) initWithValue:(id)value
{
	AvlNode *empty = [AvlNode empty];

	return [self initWithValue:value left:empty right:empty];
}

- (id) initWithValue:(id)value
				left:(AvlNode *)lt
			   right:(AvlNode *)gt
{
	if ( ( self = [super init] ) )
	{
		[self setValue:value left:lt right:gt];
	}

	return self;
}

- (void) setValue:(id)value
			 left:(AvlNode *)lt
			right:(AvlNode *)gt
{
	_value = value;
	_left = lt;
	_right = gt;
	_count = 1 + _left->_count + _right->_count;
	_depth = 1 + MAX( _left->_depth, _right->_depth );
}

- (AvlNode *) min
{
	AvlNode *empty = [AvlNode empty];

	if ( [self isEmpty] ) return empty;

	AvlNode *dict = self;
	AvlNode *next = dict->_left;

	while ( next != empty )
	{
		dict = next;
		next = dict->_left;
	}

	return dict;
}

- (AvlNode *) fixRootBalance
{
	int bal = [self balance];

	if ( abs( bal ) < 2 ) return self;

	if ( bal == 2 )
	{
		int leftBalance = [_left balance];

		if ( leftBalance == 1 || leftBalance == 0 )
		{
			//Easy case:
			return [self rotateToGT];
		}

		if ( leftBalance == -1 )
		{
			//Rotate LTDict:
			AvlNode *newlt = [[self toMutableIfNecessary:_left] rotateToLT];
			AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right];

			return [newroot rotateToGT];
		}

		[NSException raise:@"LTDict too unbalanced" format:@"LTDict too unbalanced: %d", leftBalance];
	}

	if ( bal == -2 )
	{
		int rightBalance = [_right balance];

		if ( rightBalance == -1 || rightBalance == 0 )
		{
			//Easy case:
			return [self rotateToLT];
		}

		if ( rightBalance == 1 )
		{
			//Rotate GTDict:
			AvlNode *newgt = [[self toMutableIfNecessary:_right] rotateToGT];
			AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt];

			return [newroot rotateToLT];
		}

		[NSException raise:@"RTDict too unbalanced" format:@"RTDict too unbalanced: %d", rightBalance];
	}

	//In this case we can show: |bal| > 2
	//if( Math.Abs(bal) > 2 ) {
	[NSException raise:@"Tree too out of balance" format:@"Tree too out of balance: %d", bal];

	return nil;
}

- (AvlNode *) searchNode:(id)value
				comparer:(EqualityComparer)comparer
{
	AvlNode *dict = self;
	AvlNode *empty = [AvlNode empty];

	while ( dict != empty )
	{
		int comp = comparer( dict->_value, value );

		if ( comp < 0 )
		{
			dict = dict->_right;
		}
		else if ( comp > 0 )
		{
			dict = dict->_left;
		}
		else
		{
			//Awesome:
			return dict;
		}
	}

	return empty;
}

/// <summary>
/// Return a new tree with the key-value pair inserted
/// If the key is already present, it replaces the value
/// This operation is O(Log N) where N is the number of keys
/// </summary>

- (AvlNode *) insertIntoNew:(int)index
					  value:(id)val
{
	if ( [self isEmpty] ) return [[AvlNode alloc] initWithValue:val];

	AvlNode *newlt = _left;
	AvlNode *newgt = _right;

	if ( index <= _left->_count )
	{
		newlt = [[self toMutableIfNecessary:_left] insertIntoNew:index value:val];
	}
	else
	{
		newgt = [[self toMutableIfNecessary:_right] insertIntoNew:index - _left->_count - 1 value:val];
	}

	AvlNode *newroot = [self newOrMutate:_value left:newlt right:newgt];

	return [newroot fixRootBalance];
}

/// <summary>
/// Return a new tree with the key-value pair inserted
/// If the key is already present, it replaces the value
/// This operation is O(Log N) where N is the number of keys
/// </summary>
- (AvlNode *) insertIntoNew:(id)val
				   comparer:(EqualityComparer)comparer
{
	if ( [self isEmpty] ) return [[AvlNode alloc] initWithValue:val];

	AvlNode *newlt = _left;
	AvlNode *newgt = _right;

	int comp = comparer( _value, val );
	id newv = _value;

	if ( comp < 0 )
	{
		//Let the GTDict put it in:
		newgt = [[self toMutableIfNecessary:_right] insertIntoNew:val comparer:comparer];
	}
	else if ( comp > 0 )
	{
		//Let the LTDict put it in:
		newlt = [[self toMutableIfNecessary:_left] insertIntoNew:val comparer:comparer];
	}
	else
	{
		//Replace the current value:
		newv = val;
	}

	AvlNode *newroot = [self newOrMutate:newv left:newlt right:newgt];

	return [newroot fixRootBalance];
}

- (AvlNode *) setItem:(int)index
				value:(id)val
{
	AvlNode *newlt = _left;
	AvlNode *newgt = _right;

	if ( index < _left->_count)
	{
		newlt = [[self toMutableIfNecessary:_left] setItem:index value:val];
	}
	else if ( index > _left->_count )
	{
		newgt = [[self toMutableIfNecessary:_right] setItem:index - _left->_count - 1 value:val];
	}
	else
	{
		return [self newOrMutate:val left:newlt right:newgt];
	}

	return [self newOrMutate:_value left:newlt right:newgt];
}

- (AvlNode *) getNodeAt:(int)index
{
	if ( index < _left->_count ) return [_left getNodeAt:index];
	if ( index > _left->_count) return [_right getNodeAt:index - _left->_count - 1];

	return self;
}

/// <summary>
/// Try to remove the key, and return the resulting Dict
/// if the key is not found, old_node is Empty, else old_node is the Dict
/// with matching Key
/// </summary>
- (AvlNode *) removeFromNew:(int)index
					  found:(BOOL *)found
{
	if ( [self isEmpty] )
	{
		*found = NO;
		return [AvlNode empty];
	}

	if ( index < _left->_count )
	{
		AvlNode *newlt = [[self toMutableIfNecessary:_left] removeFromNew:index found:found];

		if ( !*found )
		{
			//Not found, so nothing changed
			return self;
		}

		AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right];

		return [newroot fixRootBalance];
	}

	if ( index > _left->_count )
	{
		AvlNode *newgt = [[self toMutableIfNecessary:_right] removeFromNew:index - _left->_count - 1 found:found];

		if ( !*found )
		{
			//Not found, so nothing changed
			return self;
		}

		AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt];

		return [newroot fixRootBalance];
	}

	//found it
	*found = YES;

	return [self removeRoot];
}

/// <summary>
/// Try to remove the key, and return the resulting Dict
/// if the key is not found, old_node is Empty, else old_node is the Dict
/// with matching Key
/// </summary>
- (AvlNode *) removeFromNew:(id)val
				   comparer:(EqualityComparer)comparer
					  found:(BOOL *)found
{
	if ( [self isEmpty] )
	{
		*found = NO;
		return [AvlNode empty];
	}

	int comp = comparer( _value, val );

	if ( comp < 0 )
	{
		AvlNode *newgt = [[self toMutableIfNecessary:_right] removeFromNew:val comparer:comparer found:found];

		if ( !*found )
		{
			//Not found, so nothing changed
			return self;
		}

		AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt];

		return [newroot fixRootBalance];
	}
	if ( comp > 0 )
	{
		AvlNode *newlt = [[self toMutableIfNecessary:_left] removeFromNew:val comparer:comparer found:found];

		if ( !*found )
		{
			//Not found, so nothing changed
			return self;
		}

		AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right];

		return [newroot fixRootBalance];
	}

	//found it
	*found = YES;

	return [self removeRoot];
}

- (AvlNode *) removeMax:(AvlNode **)max
{
	if ( [self isEmpty] )
	{
		AvlNode *empty = [AvlNode empty];
		*max = empty;

		return empty;
	}

	if ( [_right isEmpty] )
	{
		//We are the max:
		*max = self;

		return _left;
	}
	else
	{
		//Go down:
		AvlNode *newgt = [[self toMutableIfNecessary:_right] removeMax:max];
		AvlNode *newroot = [self newOrMutate:_value left:_left right:newgt];

		return [newroot fixRootBalance];
	}
}

- (AvlNode *) removeMin:(AvlNode **)min
{
	if ( [self isEmpty] )
	{
		AvlNode *empty = [AvlNode empty];
		*min = empty;

		return empty;
	}

	if ( [_left isEmpty] )
	{
		//We are the minimum:
		*min = self;

		return _right;
	}
	else
	{
		//Go down:
		AvlNode *newlt = [[self toMutableIfNecessary:_left] removeMin:min];
		AvlNode *newroot = [self newOrMutate:_value left:newlt right:_right];

		return [newroot fixRootBalance];
	}
}

/// <summary>
/// Return a new dict with the root key-value pair removed
/// </summary>
- (AvlNode *) removeRoot
{
	if ( [self isEmpty] )
	{
		return self;
	}

	if ( [_left isEmpty] )
	{
		return _right;
	}

	if ( [_right isEmpty] )
	{
		return _left;
	}

	//Neither are empty:
	if ( _left->_count < _right->_count)
	{
		//LTDict has fewer, so promote from GTDict to minimize depth
		AvlNode *min;
		AvlNode *newgt = [[self toMutableIfNecessary:_right] removeMin:&min];
		AvlNode *newroot = [self newOrMutate:min.value left:_left right:newgt];

		return [newroot fixRootBalance];
	}
	else
	{
		AvlNode *max;
		AvlNode *newlt = [[self toMutableIfNecessary:_left] removeMax:&max];
		AvlNode *newroot = [self newOrMutate:max.value left:newlt right:_right];

		return [newroot fixRootBalance];
	}
}

/// <summary>
/// Move the Root into the GTDict and promote LTDict node up
/// If LTDict is empty, this operation returns this
/// </summary>
- (AvlNode *) rotateToGT
{
	if ( [_left isEmpty] || [self isEmpty] )
	{
		return self;
	}

	AvlNode *newLeft = [self toMutableIfNecessary:_left];
	AvlNode *lL = newLeft.left;
	AvlNode *lR = newLeft.right;
	AvlNode *newRight = [self newOrMutate:_value left:lR right:_right];

	return [newLeft newOrMutate:newLeft.value left:lL right:newRight];
}

/// <summary>
/// Move the Root into the LTDict and promote GTDict node up
/// If GTDict is empty, this operation returns this
/// </summary>
- (AvlNode *) rotateToLT
{
	if ( [_right isEmpty] || [self isEmpty] )
	{
		return self;
	}

	AvlNode *newRight = [self toMutableIfNecessary:_right];
	AvlNode *rL = newRight.left;
	AvlNode* rR = newRight.right;
	AvlNode *newLeft = [self newOrMutate:_value left:_left right:rL];

	return [newRight newOrMutate:newRight.value left:newLeft right:rR];
}

- (NSEnumerator *)getEnumerator:(BOOL)reverse
{
	return [[AvlNodeEnumerator alloc] initWithNode:self reverse:reverse];
}

- (AvlNode *) toImmutable
{
	return self;
}

- (AvlNode *) newOrMutate:(id)newValue
					 left:(AvlNode *)newLeft
					right:(AvlNode *)newRight
{
	return [[AvlNode alloc] initWithValue:newValue left:newLeft right:newRight];
}

- (AvlNode *) toMutable
{
	return [[MutableAvlNode alloc] initWithValue:_value left:_left right:_right];
}

- (AvlNode *) toMutableIfNecessary:(AvlNode *)node
{
	return node;
}

- (BOOL) isMutable
{
	return NO;
}

@end

@implementation MutableAvlNode

- (id) initWithValue:(id)val
				left:(AvlNode *)lt
			   right:(AvlNode *)gt
{
	return [super initWithValue:val left:lt right:gt];
}

- (AvlNode *) toImmutable
{
	return [[AvlNode alloc] initWithValue:self.value left:[self.left toImmutable] right:[self.right toImmutable]];
}

- (AvlNode *) newOrMutate:(id)newValue
					 left:(AvlNode *)newLeft
					right:(AvlNode *)newRight
{
	[self setValue:newValue left:newLeft right:newRight];

	return self;
}

- (AvlNode *) toMutable
{
	return self;
}

- (AvlNode *) toMutableIfNecessary:(AvlNode *)node
{
	return [node toMutable];
}

- (BOOL) isMutable
{
	return YES;
}

@end

@implementation NullNode

- (BOOL)isEmpty
{
	return YES;
}

@end

@implementation AvlNodeEnumerator
{
	NSMutableArray *_toVisit;
	BOOL _reverse;
	BOOL _needsPop;
	AvlNode *_root;
	AvlNode *_this_d;
}

- (id) initWithNode:(AvlNode *)root
			reverse:(BOOL)reverse
{
	if ( ( self = [super init] ) )
	{
		_root = root;
		_toVisit = [NSMutableArray new];
		[_toVisit addObject:root];

		_reverse = reverse;
		_needsPop = YES;
	}

	return self;
}

- (id) nextObject
{
	while ( !_needsPop || [_toVisit count] > 0 )
	{
		if ( _needsPop )
		{
			_this_d = [_toVisit lastObject];
			[_toVisit removeLastObject];
		}
		else _needsPop = YES;

		if ( [_this_d isEmpty] )
		{
			continue;
		}

		if ( _reverse )
		{
			if ( [_this_d.right isEmpty] )
			{
				//This is the next biggest value in the Dict:
				id value = _this_d.value;
				_this_d = _this_d.left;
				_needsPop = NO;
				return value;
			}
			else
			{
				//Break it up
				[_toVisit addObject:_this_d.left];
				[_toVisit addObject:[[AvlNode alloc] initWithValue:_this_d.value]];
				_this_d = _this_d.right;
				_needsPop = NO;
			}
		}
		else
		{
			if ( [_this_d.left isEmpty] )
			{
				//This is the next biggest value in the Dict:
				id value = _this_d.value;
				_this_d = _this_d.right;
				_needsPop = NO;
				return value;
			}
			else
			{
				//Break it up
				if ( ![_this_d.right isEmpty] )
				{
					[_toVisit addObject:_this_d.right];
				}

				[_toVisit addObject:[[AvlNode alloc] initWithValue:_this_d.value]];
				_this_d = _this_d.left;
				_needsPop = NO;
			}
		}
	}

	return nil;
}

- (NSArray *)allObjects
{
	NSMutableArray *result = [NSMutableArray new];
	AvlNodeEnumerator *enumerator = [[AvlNodeEnumerator alloc] initWithNode:_root reverse:_reverse];

	id item = nil;

	while ( nil != ( item = [enumerator nextObject] ) )
	{
		[result addObject:item];
	}

	return result;
}

@end