support Cocoa collections as part of a cycle, which also required a c…

…hange to zero out the references within the candidate object rather than the references to the candidate object (which probably makes more sense anyway)

Circle/CircleIVarLayout.h

@@ -12,6 +12,4 @@
 unsigned *CalculateClassStrongLayout(Class c);
 unsigned *CalculateBlockStrongLayout(void *block);
 
-unsigned *GetStrongLayout(void *obj);
-
 void EnumerateStrongReferences(void *obj, void (^block)(void **reference, void *target));

Circle/CircleIVarLayout.m

@@ -36,9 +36,18 @@
 };
 
 
+enum Classification
+{
+    ENUMERABLE,
+    DICTIONARY,
+    BLOCK,
+    OTHER
+};
+
 static unsigned kNoStrongReferencesLayout[] = { 0 };
 
 static CFMutableDictionaryRef gLayoutCache;
+static CFMutableDictionaryRef gClassificationCache;
 
 
 @interface _CircleReleaseDetector : NSObject {
@@ -164,20 +173,60 @@ static BOOL IsBlock(void *obj)
     return layout;
 }
 
-unsigned *GetStrongLayout(void *obj)
+static enum Classification Classify(void *obj)
 {
-    return IsBlock(obj) ? GetBlockStrongLayout(obj) : GetClassStrongLayout(object_getClass((__bridge id)obj));
+    if(!gClassificationCache)
+        gClassificationCache = CFDictionaryCreateMutable(NULL, 0, NULL, NULL);
+
+    void *key = (__bridge void *)object_getClass((__bridge id)obj);
+
+    const void *value;
+    Boolean present = CFDictionaryGetValueIfPresent(gClassificationCache, key, &value);
+    if(present)
+        return (enum Classification)value;
+
+    enum Classification classification = OTHER;
+    if(IsBlock(obj))
+        classification = BLOCK;
+    else if([(__bridge id)obj isKindOfClass: [NSArray class]] || [(__bridge id)obj isKindOfClass: [NSSet class]])
+        classification = ENUMERABLE;
+    else if([(__bridge id)obj isKindOfClass: [NSDictionary class]])
+        classification = DICTIONARY;
+
+    CFDictionarySetValue(gClassificationCache, key, (const void *)classification);
+
+    return classification;
 }
 
 void EnumerateStrongReferences(void *obj, void (^block)(void **reference, void *target))
 {
-    void **objAsReferences = obj;
-    unsigned *layout = GetStrongLayout(obj);
-    for(int i = 0; layout[i]; i++)
+    enum Classification classification = Classify(obj);
+    if(classification == ENUMERABLE)
+    {
+        for(id target in (__bridge id)obj)
+            block(NULL, (__bridge void *)target);
+    }
+    else if(classification == DICTIONARY)
     {
-        void **reference = &objAsReferences[layout[i]];
-        void *target = reference ? *reference : NULL;
-        block(reference, target);
+        [(__bridge NSDictionary *)obj enumerateKeysAndObjectsUsingBlock: ^(id key, id obj, BOOL *stop) {
+            block(NULL, (__bridge void *)key);
+            block(NULL, (__bridge void *)obj);
+        }];
+    }
+    else
+    {
+        unsigned *layout;
+        if(classification == BLOCK)
+            layout = GetBlockStrongLayout(obj);
+        else
+            layout = GetClassStrongLayout(object_getClass((__bridge id)obj));
+
+        void **objAsReferences = obj;
+        for(int i = 0; layout[i]; i++)
+        {
+            void **reference = &objAsReferences[layout[i]];
+            void *target = reference ? *reference : NULL;
+            block(reference, target);
+        }
     }
 }
-

Circle/CircleSimpleCycleFinder.h

@@ -12,7 +12,7 @@
 struct CircleSearchResults
 {
     BOOL isUnclaimedCycle;
-    CFSetRef incomingReferences;
+    CFSetRef referencesToZero;
     CFDictionaryRef infos;
 };
 

Circle/CircleSimpleCycleFinder.m

@@ -109,7 +109,7 @@ struct CircleSearchResults CircleSimpleSearchCycle(id obj, BOOL gatherAll)
             // short circuit: if there's no info object for obj, then it's not part of any sort of cycle
             struct CircleSearchResults results;
             results.isUnclaimedCycle = NO;
-            results.incomingReferences = CFSetCreate(NULL, NULL, 0, NULL);
+            results.referencesToZero = CFSetCreate(NULL, NULL, 0, NULL);
             results.infos = infos;
             return results;
         }
@@ -175,9 +175,15 @@ struct CircleSearchResults CircleSimpleSearchCycle(id obj, BOOL gatherAll)
 
     LOG(@"foundExternallyRetained is %d", foundExternallyRetained);
 
+    CFMutableSetRef referencesToZero = CFSetCreateMutable(NULL, 0, NULL);
+    EnumerateStrongReferences((__bridge void *)obj, ^(void **reference, void *target) {
+        if(target)
+            CFSetAddValue(referencesToZero, reference);
+    });
+
     struct CircleSearchResults results;
     results.isUnclaimedCycle = !foundExternallyRetained;
-    results.incomingReferences = CFRetain(incomingReferences);
+    results.referencesToZero = referencesToZero;
     results.infos = infos;
 
     CFRelease(toSearchObjs);
@@ -188,11 +194,11 @@ struct CircleSearchResults CircleSimpleSearchCycle(id obj, BOOL gatherAll)
 
 void CircleZeroReferences(CFSetRef references)
 {
-    NSUInteger incomingReferencesCount = CFSetGetCount(references);
+    NSUInteger referencesCount = CFSetGetCount(references);
 
-    const void *locations[incomingReferencesCount];
+    const void *locations[referencesCount];
     CFSetGetValues(references, locations);
-    for(unsigned i = 0; i < incomingReferencesCount; i++)
+    for(unsigned i = 0; i < referencesCount; i++)
     {
         void **reference = (void **)locations[i];
         void *target = *reference;
@@ -243,7 +249,7 @@ - (void)_enumerateObjectsGatherAll: (BOOL) gatherAll resultsCallback: (void (^)(
         {
             struct CircleSearchResults results = CircleSimpleSearchCycle(obj, gatherAll);
             block(results);
-            CFRelease(results.incomingReferences);
+            CFRelease(results.referencesToZero);
             CFRelease(results.infos);
         }
         else
@@ -265,7 +271,7 @@ - (void)collect
 {
     [self _enumerateObjectsGatherAll: NO resultsCallback: ^(struct CircleSearchResults results) {
         if(results.isUnclaimedCycle)
-            CircleZeroReferences(results.incomingReferences);
+            CircleZeroReferences(results.referencesToZero);
     }];
 }
 

CircleTests/CircleTests.m

@@ -107,6 +107,8 @@ - (void)testComplexCycle
         weakObj = a;
 
         [collector addCandidate: a];
+        [collector addCandidate: b];
+        [collector addCandidate: c];
     }
 
     @autoreleasepool {
@@ -197,4 +199,67 @@ - (void)testManyCollections
         STAssertNil(weakObjs[i], @"Collector failed to collect a cycle");
 }
 
+- (void)testArrayCycle
+{
+    CircleSimpleCycleFinder *collector = [[CircleSimpleCycleFinder alloc] init];
+
+    __weak id weakObj;
+    @autoreleasepool {
+        Referrer *a = [[Referrer alloc] init];
+        NSArray *b = @[ a ];
+        [a setPtr1: b];
+        weakObj = a;
+
+        [collector addCandidate: a];
+    }
+
+    @autoreleasepool {
+        STAssertNotNil(weakObj, @"Weak pointer to cycle should not be nil before running the collector");
+    }
+    [collector collect];
+    STAssertNil(weakObj, @"Collector didn't collect a cycle");
+}
+
+- (void)testSetCycle
+{
+    CircleSimpleCycleFinder *collector = [[CircleSimpleCycleFinder alloc] init];
+
+    __weak id weakObj;
+    @autoreleasepool {
+        Referrer *a = [[Referrer alloc] init];
+        NSSet *b = [NSSet setWithObject: a];
+        [a setPtr1: b];
+        weakObj = a;
+
+        [collector addCandidate: a];
+    }
+
+    @autoreleasepool {
+        STAssertNotNil(weakObj, @"Weak pointer to cycle should not be nil before running the collector");
+    }
+    [collector collect];
+    STAssertNil(weakObj, @"Collector didn't collect a cycle");
+}
+
+- (void)testDictionaryCycle
+{
+    CircleSimpleCycleFinder *collector = [[CircleSimpleCycleFinder alloc] init];
+
+    __weak id weakObj;
+    @autoreleasepool {
+        Referrer *a = [[Referrer alloc] init];
+        NSDictionary *b = @{ @"a" : a };
+        [a setPtr1: b];
+        weakObj = a;
+
+        [collector addCandidate: a];
+    }
+
+    @autoreleasepool {
+        STAssertNotNil(weakObj, @"Weak pointer to cycle should not be nil before running the collector");
+    }
+    [collector collect];
+    STAssertNil(weakObj, @"Collector didn't collect a cycle");
+}
+
 @end