From e6643f8983302453a1a4026664d53876d11fe78d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ce=CC=81dric=20Luthi?= <cedric.luthi@gmail.com>
Date: Wed, 20 Aug 2014 09:41:47 +0200
Subject: [PATCH] Add support for non-lazy classes and categories

---
 Source/CDOCProtocol.m           | 10 ++++++++++
 Source/CDObjectiveC2Processor.m | 32 ++++++++++++++++++--------------
 Source/CDObjectiveCProcessor.m  | 12 ++++++------
 3 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/Source/CDOCProtocol.m b/Source/CDOCProtocol.m
index c25586cb..04337894 100644
--- a/Source/CDOCProtocol.m
+++ b/Source/CDOCProtocol.m
@@ -50,6 +50,16 @@ - (id)init;
     return self;
 }
 
+- (BOOL)isEqual:(id)object;
+{
+    return [object isKindOfClass:[CDOCProtocol class]] && [self.sortableName isEqualToString:((CDOCProtocol *)object).sortableName];
+}
+
+- (NSUInteger)hash;
+{
+    return [self.sortableName hash];
+}
+
 #pragma mark - Debugging
 
 - (NSString *)description;
diff --git a/Source/CDObjectiveC2Processor.m b/Source/CDObjectiveC2Processor.m
index 4bc07cf9..95f18aaf 100644
--- a/Source/CDObjectiveC2Processor.m
+++ b/Source/CDObjectiveC2Processor.m
@@ -36,26 +36,30 @@ - (void)loadProtocols;
 
 - (void)loadClasses;
 {
-    CDSection *section = [[self.machOFile segmentWithName:@"__DATA"] sectionWithName:@"__objc_classlist"];
-    
-    CDMachOFileDataCursor *cursor = [[CDMachOFileDataCursor alloc] initWithSection:section];
-    while ([cursor isAtEnd] == NO) {
-        uint64_t val = [cursor readPtr];
-        CDOCClass *aClass = [self loadClassAtAddress:val];
-        if (aClass != nil) {
-            [self addClass:aClass withAddress:val];
+    for (NSString *sectionName in @[@"__objc_classlist", @"__objc_nlclslist"]) {
+        CDSection *section = [[self.machOFile segmentWithName:@"__DATA"] sectionWithName:sectionName];
+        
+        CDMachOFileDataCursor *cursor = [[CDMachOFileDataCursor alloc] initWithSection:section];
+        while ([cursor isAtEnd] == NO) {
+            uint64_t val = [cursor readPtr];
+            CDOCClass *aClass = [self loadClassAtAddress:val];
+            if (aClass != nil) {
+                [self addClass:aClass withAddress:val];
+            }
         }
     }
 }
 
 - (void)loadCategories;
 {
-    CDSection *section = [[self.machOFile segmentWithName:@"__DATA"] sectionWithName:@"__objc_catlist"];
-    
-    CDMachOFileDataCursor *cursor = [[CDMachOFileDataCursor alloc] initWithSection:section];
-    while ([cursor isAtEnd] == NO) {
-        CDOCCategory *category = [self loadCategoryAtAddress:[cursor readPtr]];
-        [self addCategory:category];
+    for (NSString *sectionName in @[@"__objc_catlist", @"__objc_nlcatlist"]) {
+        CDSection *section = [[self.machOFile segmentWithName:@"__DATA"] sectionWithName:sectionName];
+        
+        CDMachOFileDataCursor *cursor = [[CDMachOFileDataCursor alloc] initWithSection:section];
+        while ([cursor isAtEnd] == NO) {
+            CDOCCategory *category = [self loadCategoryAtAddress:[cursor readPtr]];
+            [self addCategory:category];
+        }
     }
 }
 
diff --git a/Source/CDObjectiveCProcessor.m b/Source/CDObjectiveCProcessor.m
index ae3266af..e92d1ace 100644
--- a/Source/CDObjectiveCProcessor.m
+++ b/Source/CDObjectiveCProcessor.m
@@ -24,10 +24,10 @@ @implementation CDObjectiveCProcessor
 {
     CDMachOFile *_machOFile;
     
-    NSMutableArray *_classes;
+    NSMutableOrderedSet *_classes;
     NSMutableDictionary *_classesByAddress;
     
-    NSMutableArray *_categories;
+    NSMutableOrderedSet *_categories;
     
     CDProtocolUniquer *_protocolUniquer;
 }
@@ -36,9 +36,9 @@ - (id)initWithMachOFile:(CDMachOFile *)machOFile;
 {
     if ((self = [super init])) {
         _machOFile = machOFile;
-        _classes = [[NSMutableArray alloc] init];
+        _classes = [[NSMutableOrderedSet alloc] init];
         _classesByAddress = [[NSMutableDictionary alloc] init];
-        _categories = [[NSMutableArray alloc] init];
+        _categories = [[NSMutableOrderedSet alloc] init];
         
         _protocolUniquer = [[CDProtocolUniquer alloc] init];
     }
@@ -177,8 +177,8 @@ - (void)registerTypesWithObject:(CDTypeController *)typeController phase:(NSUInt
 - (void)recursivelyVisit:(CDVisitor *)visitor;
 {
     NSMutableArray *classesAndCategories = [[NSMutableArray alloc] init];
-    [classesAndCategories addObjectsFromArray:_classes];
-    [classesAndCategories addObjectsFromArray:_categories];
+    [classesAndCategories addObjectsFromArray:[_classes array]];
+    [classesAndCategories addObjectsFromArray:[_categories array]];
 
     [visitor willVisitObjectiveCProcessor:self];
     [visitor visitObjectiveCProcessor:self];