@@ -34,10 +34,11 @@
 #import "XMPPStanza.h"
 #import "XMPPJID.h"
 #import "XMPPIQ.h"
 #import "XMPPMessage.h"
 #import "XMPPPresence.h"
+#import "XMPPRoster.h"
 #import "XMPPRosterItem.h"
 #import "XMPPExceptions.h"
 
 #define NS_BIND @"urn:ietf:params:xml:ns:xmpp-bind"
 #define NS_CLIENT @"jabber:client"
@@ -79,11 +80,11 @@
 		useTLS = YES;
 
 		parser.delegate = self;
 		elementBuilder.delegate = self;
 
-		roster = [[OFMutableDictionary alloc] init];
+		roster = [[XMPPRoster alloc] initWithConnection: self];
 	} @catch (id e) {
 		[self release];
 		@throw e;
 	}
 
@@ -487,11 +488,11 @@
 		[self sendStanza: [OFXMLElement elementWithName: @"starttls"
 						      namespace: NS_STARTTLS]];
 		return;
 	}
 
-	if ([mechs count] > 0) {
+	if (mechs.count > 0) {
 		for (OFXMLElement *mech in [mechs.firstObject children])
 			[mechanisms addObject:
 			    [mech.children.firstObject stringValue]];
 
 		if ([mechanisms containsObject: @"SCRAM-SHA-1"]) {
@@ -652,45 +653,25 @@
 		    ![elem.ns isEqual: NS_ROSTER])
 			continue;
 
 		rosterItem = [XMPPRosterItem rosterItem];
 		rosterItem.JID = [XMPPJID JIDWithString:
-		    [rosterElem attributeForName: @"jid"].stringValue];
-		rosterItem.name =
-		    [rosterElem attributeForName: @"name"].stringValue;
+		    [elem attributeForName: @"jid"].stringValue];
+		rosterItem.name = [elem attributeForName: @"name"].stringValue;
 		rosterItem.subscription =
-		    [rosterElem attributeForName: @"subscription"].stringValue;
+		    [elem attributeForName: @"subscription"].stringValue;
 
 		for (OFXMLElement *groupElem in
 		     [elem elementsForName: @"group"
-				 namespace: NS_ROSTER]) {
-			OFMutableArray *rosterGroup =
-			     [roster objectForKey: rosterElem.stringValue];
-
-			if (rosterGroup == nil) {
-				rosterGroup = [OFMutableArray array];
-				[roster setObject: rosterGroup
-					   forKey: rosterElem.stringValue];
-			}
-
-			[rosterGroup addObject: rosterItem];
-		}
-
-		rosterItem.groups = groups;
-
-		if (groups.count == 0) {
-			OFMutableArray *rosterGroup =
-			    [roster objectForKey: @""];
-
-			if (rosterGroup == nil) {
-				rosterGroup = [OFMutableArray array];
-				[roster setObject: rosterGroup
-					   forKey: @""];
-			}
-
-			[rosterGroup addObject: rosterItem];
-		}
+				 namespace: NS_ROSTER])
+			[groups addObject:
+			    [groupElem.children.firstObject stringValue]];
+
+		if (groups.count > 0)
+			rosterItem.groups = groups;
+
+		[roster XMPP_addRosterItem: rosterItem];
 	}
 
 	if ([delegate respondsToSelector:
 	    @selector(connectionDidReceiveRoster:)])
 		[delegate connectionDidReceiveRoster: self];