From 85f84fc806d8e255f139efcc1dc5d083e3c30544 Mon Sep 17 00:00:00 2001
From: zhongwuzw <zhongwuzw@qq.com>
Date: Thu, 1 Mar 2018 09:43:49 +0800
Subject: [PATCH] Fix issue when object dealloc in DFS

---
 .../Detector/FBRetainCycleDetector.mm         | 37 +++++++++----------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/FBRetainCycleDetector/Detector/FBRetainCycleDetector.mm b/FBRetainCycleDetector/Detector/FBRetainCycleDetector.mm
index b340943..8889176 100644
--- a/FBRetainCycleDetector/Detector/FBRetainCycleDetector.mm
+++ b/FBRetainCycleDetector/Detector/FBRetainCycleDetector.mm
@@ -24,7 +24,7 @@ @implementation FBRetainCycleDetector
 {
   NSMutableArray *_candidates;
   FBObjectGraphConfiguration *_configuration;
-  NSMutableSet *_objectSet;
+  NSHashTable *_objectSet;
 }
 
 - (instancetype)initWithConfiguration:(FBObjectGraphConfiguration *)configuration
@@ -32,7 +32,7 @@ - (instancetype)initWithConfiguration:(FBObjectGraphConfiguration *)configuratio
   if (self = [super init]) {
     _configuration = configuration;
     _candidates = [NSMutableArray new];
-    _objectSet = [NSMutableSet new];
+    _objectSet = [[NSHashTable alloc] initWithOptions:NSPointerFunctionsWeakMemory | NSPointerFunctionsObjectPointerPersonality capacity:kFBRetainCycleDetectorDefaultStackDepth];
   }
 
   return self;
@@ -99,8 +99,8 @@ - (void)addCandidate:(id)candidate
   NSMutableArray<FBNodeEnumerator *> *stack = [NSMutableArray new];
 
   // To make the search non-linear we will also keep
-  // a set of previously visited nodes.
-  NSMutableSet<FBNodeEnumerator *> *objectsOnPath = [NSMutableSet new];
+  // a set of previously visited objects that nodes monitor.
+  NSHashTable *objectsOnPath = [[NSHashTable alloc] initWithOptions:NSPointerFunctionsWeakMemory | NSPointerFunctionsObjectPointerPersonality capacity:kFBRetainCycleDetectorDefaultStackDepth];
 
   // Let's start with the root
   [stack addObject:wrappedObject];
@@ -112,31 +112,33 @@ - (void)addCandidate:(id)candidate
     @autoreleasepool {
       // Take topmost node in stack and mark it as visited
       FBNodeEnumerator *top = [stack lastObject];
+      // Take object that topmost node monitor
+      __weak id object = top.object.object;
 
       // We don't want to retraverse the same subtree
-      if (![objectsOnPath containsObject:top]) {
-        if ([_objectSet containsObject:@([top.object objectAddress])]) {
+      if (![objectsOnPath containsObject:object]) {
+        if ([_objectSet containsObject:object]) {
           [stack removeLastObject];
           continue;
         }
-        // Add the object address to the set as an NSNumber to avoid
-        // unnecessarily retaining the object
-        [_objectSet addObject:@([top.object objectAddress])];
+
+        [_objectSet addObject:object];
       }
 
-      [objectsOnPath addObject:top];
+      [objectsOnPath addObject:object];
 
       // Take next adjecent node to that child. Wrapper object can
       // persist iteration state. If we see that node again, it will
       // give us new adjacent node unless it runs out of them
       FBNodeEnumerator *firstAdjacent = [top nextObject];
-      if (firstAdjacent) {
+      __weak id firstAdjacentObject = firstAdjacent.object.object;
+      if (firstAdjacentObject) {
         // Current node still has some adjacent not-visited nodes
 
         BOOL shouldPushToStack = NO;
 
         // Check if child was already seen in that path
-        if ([objectsOnPath containsObject:firstAdjacent]) {
+        if ([objectsOnPath containsObject:firstAdjacentObject]) {
           // We have caught a retain cycle
 
           // Ignore the first element which is equal to firstAdjacent, use firstAdjacent
@@ -145,10 +147,7 @@ - (void)addCandidate:(id)candidate
           NSUInteger index = [stack indexOfObject:firstAdjacent];
           NSInteger length = [stack count] - index;
 
-          if (index == NSNotFound) {
-            // Object got deallocated between checking if it exists and grabbing its index
-            shouldPushToStack = YES;
-          } else {
+          if (firstAdjacentObject) {
             NSRange cycleRange = NSMakeRange(index, length);
             NSMutableArray<FBNodeEnumerator *> *cycle = [[stack subarrayWithRange:cycleRange] mutableCopy];
             [cycle replaceObjectAtIndex:0 withObject:firstAdjacent];
@@ -157,10 +156,10 @@ - (void)addCandidate:(id)candidate
             // 2. Shift to lowest address (if we omit that, and the cycle is created by same class,
             //    we might have duplicates)
             // 3. Shift by class (lexicographically)
-
+          
             [retainCycles addObject:[self _shiftToUnifiedCycle:[self _unwrapCycle:cycle]]];
           }
-        } else {
+        } else if (firstAdjacentObject) {
           // Node is clear to check, add it to stack and continue
           shouldPushToStack = YES;
         }
@@ -173,7 +172,7 @@ - (void)addCandidate:(id)candidate
       } else {
         // Node has no more adjacent nodes, it itself is done, move on
         [stack removeLastObject];
-        [objectsOnPath removeObject:top];
+        [objectsOnPath removeObject:object];
       }
     }
   }