Replace NSMutableSet with NSHashTable when Appropriate #trivial

Source/ASCollectionView.mm

@@ -90,8 +90,8 @@ @interface ASCollectionView () <ASRangeControllerDataSource, ASRangeControllerDe
   ASCollectionViewLayoutController *_layoutController;
   id<ASCollectionViewLayoutInspecting> _defaultLayoutInspector;
   __weak id<ASCollectionViewLayoutInspecting> _layoutInspector;
-  NSMutableSet *_cellsForVisibilityUpdates;
-  NSMutableSet *_cellsForLayoutUpdates;
+  NSHashTable<_ASCollectionViewCell *> *_cellsForVisibilityUpdates;
+  NSHashTable<ASCellNode *> *_cellsForLayoutUpdates;
   id<ASCollectionViewLayoutFacilitatorProtocol> _layoutFacilitator;
   CGFloat _leadingScreensForBatching;
   BOOL _inverted;
@@ -299,8 +299,8 @@ - (instancetype)_initWithFrame:(CGRect)frame collectionViewLayout:(UICollectionV
   _registeredSupplementaryKinds = [NSMutableSet set];
   _visibleElements = [[NSCountedSet alloc] init];
 
-  _cellsForVisibilityUpdates = [NSMutableSet set];
-  _cellsForLayoutUpdates = [NSMutableSet set];
+  _cellsForVisibilityUpdates = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
+  _cellsForLayoutUpdates = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
   self.backgroundColor = [UIColor whiteColor];
 
   [self registerClass:[_ASCollectionViewCell class] forCellWithReuseIdentifier:kReuseIdentifier];
@@ -1827,9 +1827,9 @@ - (ASRangeController *)rangeController
   return result;
 }
 
-- (NSArray<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController
+- (NSHashTable<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController
 {
-  return _visibleElements.allObjects;
+  return ASPointerTableByFlatMapping(_visibleElements, id element, element);
 }
 
 - (ASElementMap *)elementMapForRangeController:(ASRangeController *)rangeController

Source/ASDisplayNode.mm

@@ -854,10 +854,6 @@ - (void)setDebugName:(NSString *)debugName
 
 #pragma mark - Layout
 
-#if DEBUG
-  #define AS_DEDUPE_LAYOUT_SPEC_TREE 1
-#endif
-
 // At most a layoutSpecBlock or one of the three layout methods is overridden
 #define __ASDisplayNodeCheckForLayoutMethodOverrides \
     ASDisplayNodeAssert(_layoutSpecBlock != NULL || \
@@ -1002,7 +998,7 @@ - (ASLayout *)calculateLayoutThatFits:(ASSizeRange)constrainedSize
     ASLayoutSpec *layoutSpec = (ASLayoutSpec *)layoutElement;
 
 #if AS_DEDUPE_LAYOUT_SPEC_TREE
-    NSSet *duplicateElements = [layoutSpec findDuplicatedElementsInSubtree];
+    NSHashTable *duplicateElements = [layoutSpec findDuplicatedElementsInSubtree];
     if (duplicateElements.count > 0) {
       ASDisplayNodeFailAssert(@"Node %@ returned a layout spec that contains the same elements in multiple positions. Elements: %@", self, duplicateElements);
       // Use an empty layout spec to avoid crashes

Source/ASTableView.mm

@@ -191,13 +191,13 @@ @interface ASTableView () <ASRangeControllerDataSource, ASRangeControllerDelegat
   CGFloat _nodesConstrainedWidth;
   BOOL _queuedNodeHeightUpdate;
   BOOL _isDeallocating;
-  NSMutableSet *_cellsForVisibilityUpdates;
+  NSHashTable<_ASTableViewCell *> *_cellsForVisibilityUpdates;
 
   // CountedSet because UIKit may display the same element in multiple cells e.g. during animations.
   NSCountedSet<ASCollectionElement *> *_visibleElements;
 
   BOOL _remeasuringCellNodes;
-  NSMutableSet *_cellsForLayoutUpdates;
+  NSHashTable<ASCellNode *> *_cellsForLayoutUpdates;
 
   // See documentation on same property in ASCollectionView
   BOOL _hasEverCheckedForBatchFetchingDueToUpdate;
@@ -309,8 +309,8 @@ - (instancetype)_initWithFrame:(CGRect)frame style:(UITableViewStyle)style dataC
   if (!(self = [super initWithFrame:frame style:style])) {
     return nil;
   }
-  _cellsForVisibilityUpdates = [NSMutableSet set];
-  _cellsForLayoutUpdates = [NSMutableSet set];
+  _cellsForVisibilityUpdates = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
+  _cellsForLayoutUpdates = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
   if (!dataControllerClass) {
     dataControllerClass = [[self class] dataControllerClass];
   }
@@ -686,7 +686,7 @@ - (nullable NSIndexPath *)indexPathForNode:(ASCellNode *)cellNode waitingIfNeede
 
 - (NSArray<ASCellNode *> *)visibleNodes
 {
-  NSArray<ASCollectionElement *> *elements = [self visibleElementsForRangeController:_rangeController];
+  auto elements = [self visibleElementsForRangeController:_rangeController];
   return ASArrayByFlatMapping(elements, ASCollectionElement *e, e.node);
 }
 
@@ -1457,9 +1457,9 @@ - (ASRangeController *)rangeController
   return _rangeController;
 }
 
-- (NSArray<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController
+- (NSHashTable<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController
 {
-  return _visibleElements.allObjects;
+  return ASPointerTableByFlatMapping(_visibleElements, id element, element);
 }
 
 - (ASScrollDirection)scrollDirectionForRangeController:(ASRangeController *)rangeController

Source/Base/ASBaseDefines.h

@@ -235,6 +235,20 @@
   s; \
 })
 
+/**
+ * Create a new ObjectPointerPersonality NSHashTable by mapping `collection` over `work`, ignoring nil.
+ */
+#define ASPointerTableByFlatMapping(collection, decl, work) ({ \
+  NSHashTable *t = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality]; \
+  for (decl in collection) {\
+    id result = work; \
+    if (result != nil) { \
+      [t addObject:result]; \
+    } \
+  } \
+  t; \
+})
+
 /**
  * Create a new array by mapping `collection` over `work`, ignoring nil.
  */

Source/Details/ASAbstractLayoutController.h

@@ -36,9 +36,9 @@ ASDISPLAYNODE_EXTERN_C_END
 
 @interface ASAbstractLayoutController (Unavailable)
 
-- (NSSet *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType __unavailable;
+- (NSHashTable *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType __unavailable;
 
-- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet * _Nullable * _Nullable)displaySet preloadSet:(NSSet * _Nullable * _Nullable)preloadSet __unavailable;
+- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSHashTable * _Nullable * _Nullable)displaySet preloadSet:(NSHashTable * _Nullable * _Nullable)preloadSet __unavailable;
 
 @end
 

Source/Details/ASAbstractLayoutController.mm

@@ -174,13 +174,13 @@ - (void)setTuningParameters:(ASRangeTuningParameters)tuningParameters forRangeMo
 
 #pragma mark - Abstract Index Path Range Support
 
-- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
+- (NSHashTable<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
 {
   ASDisplayNodeAssertNotSupported();
   return nil;
 }
 
-- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSHashTable<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSHashTable<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
 {
   ASDisplayNodeAssertNotSupported();
 }

Source/Details/ASCollectionViewLayoutController.m

@@ -54,14 +54,14 @@ - (instancetype)initWithCollectionView:(ASCollectionView *)collectionView
   return self;
 }
 
-- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
+- (NSHashTable<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
 {
   ASRangeTuningParameters tuningParameters = [self tuningParametersForRangeMode:rangeMode rangeType:rangeType];
   CGRect rangeBounds = [self rangeBoundsWithScrollDirection:scrollDirection rangeTuningParameters:tuningParameters];
   return [self elementsWithinRangeBounds:rangeBounds map:map];
 }
 
-- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSHashTable<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSHashTable<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
 {
   if (displaySet == NULL || preloadSet == NULL) {
     return;
@@ -74,9 +74,10 @@ - (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(AS
 
   CGRect unionBounds = CGRectUnion(displayBounds, preloadBounds);
   NSArray *layoutAttributes = [_collectionViewLayout layoutAttributesForElementsInRect:unionBounds];
+  NSInteger count = layoutAttributes.count;
 
-  NSMutableSet<ASCollectionElement *> *display = [NSMutableSet setWithCapacity:layoutAttributes.count];
-  NSMutableSet<ASCollectionElement *> *preload = [NSMutableSet setWithCapacity:layoutAttributes.count];
+  __auto_type display = [[NSHashTable<ASCollectionElement *> alloc] initWithOptions:NSHashTableObjectPointerPersonality capacity:count];
+  __auto_type preload = [[NSHashTable<ASCollectionElement *> alloc] initWithOptions:NSHashTableObjectPointerPersonality capacity:count];
 
   for (UICollectionViewLayoutAttributes *la in layoutAttributes) {
     // Manually filter out elements that don't intersect the range bounds.
@@ -105,10 +106,10 @@ - (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(AS
   return;
 }
 
-- (NSSet<ASCollectionElement *> *)elementsWithinRangeBounds:(CGRect)rangeBounds map:(ASElementMap *)map
+- (NSHashTable<ASCollectionElement *> *)elementsWithinRangeBounds:(CGRect)rangeBounds map:(ASElementMap *)map
 {
   NSArray *layoutAttributes = [_collectionViewLayout layoutAttributesForElementsInRect:rangeBounds];
-  NSMutableSet<ASCollectionElement *> *elementSet = [NSMutableSet setWithCapacity:layoutAttributes.count];
+  NSHashTable<ASCollectionElement *> *elementSet = [[NSHashTable alloc] initWithOptions:NSHashTableObjectPointerPersonality capacity:layoutAttributes.count];
 
   for (UICollectionViewLayoutAttributes *la in layoutAttributes) {
     // Manually filter out elements that don't intersect the range bounds.

Source/Details/ASLayoutController.h

@@ -41,9 +41,9 @@ ASDISPLAYNODE_EXTERN_C_END
 
 - (ASRangeTuningParameters)tuningParametersForRangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType;
 
-- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map;
+- (NSHashTable<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map;
 
-- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> * _Nullable * _Nullable)displaySet preloadSet:(NSSet<ASCollectionElement *> * _Nullable * _Nullable)preloadSet map:(ASElementMap *)map;
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSHashTable<ASCollectionElement *> * _Nullable * _Nullable)displaySet preloadSet:(NSHashTable<ASCollectionElement *> * _Nullable * _Nullable)preloadSet map:(ASElementMap *)map;
 
 @optional
 

Source/Details/ASRangeController.h

@@ -114,9 +114,9 @@ AS_SUBCLASSING_RESTRICTED
 /**
  * @param rangeController Sender.
  *
- * @return an array of elements corresponding to the data currently visible onscreen (i.e., the visible range).
+ * @return an table of elements corresponding to the data currently visible onscreen (i.e., the visible range).
  */
-- (NSArray<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController;
+- (nullable NSHashTable<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController;
 
 /**
  * @param rangeController Sender.

Source/Details/ASRangeController.mm

@@ -216,7 +216,7 @@ - (void)_updateVisibleNodeIndexPaths
 
   // TODO: Consider if we need to use this codepath, or can rely on something more similar to the data & display ranges
   // Example: ... = [_layoutController indexPathsForScrolling:scrollDirection rangeType:ASLayoutRangeTypeVisible];
-  NSSet<ASCollectionElement *> *visibleElements = [NSSet setWithArray:[_dataSource visibleElementsForRangeController:self]];
+  auto visibleElements = [_dataSource visibleElementsForRangeController:self];
   NSHashTable *newVisibleNodes = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
 
   if (visibleElements.count == 0) { // if we don't have any visibleNodes currently (scrolled before or after content)...
@@ -255,14 +255,14 @@ - (void)_updateVisibleNodeIndexPaths
   // Check if both Display and Preload are unique. If they are, we load them with a single fetch from the layout controller for performance.
   BOOL optimizedLoadingOfBothRanges = (equalDisplayPreload == NO && equalDisplayVisible == NO && emptyDisplayRange == NO);
 
-  NSSet<ASCollectionElement *> *displayElements = nil;
-  NSSet<ASCollectionElement *> *preloadElements = nil;
+  NSHashTable<ASCollectionElement *> *displayElements = nil;
+  NSHashTable<ASCollectionElement *> *preloadElements = nil;
 
   if (optimizedLoadingOfBothRanges) {
     [_layoutController allElementsForScrolling:scrollDirection rangeMode:rangeMode displaySet:&displayElements preloadSet:&preloadElements map:map];
   } else {
     if (emptyDisplayRange == YES) {
-      displayElements = [NSSet set];
+      displayElements = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
     } if (equalDisplayVisible == YES) {
       displayElements = visibleElements;
     } else {

Source/Details/ASTableLayoutController.m

@@ -38,17 +38,17 @@ - (instancetype)initWithTableView:(UITableView *)tableView
 
 #pragma mark - ASLayoutController
 
-- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
+- (NSHashTable<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
 {
   CGRect bounds = _tableView.bounds;
 
   ASRangeTuningParameters tuningParameters = [self tuningParametersForRangeMode:rangeMode rangeType:rangeType];
   CGRect rangeBounds = CGRectExpandToRangeWithScrollableDirections(bounds, tuningParameters, ASScrollDirectionVerticalDirections, scrollDirection);
   NSArray *array = [_tableView indexPathsForRowsInRect:rangeBounds];
-  return ASSetByFlatMapping(array, NSIndexPath *indexPath, [map elementForItemAtIndexPath:indexPath]);
+  return ASPointerTableByFlatMapping(array, NSIndexPath *indexPath, [map elementForItemAtIndexPath:indexPath]);
 }
 
-- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSHashTable<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSHashTable<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
 {
   if (displaySet == NULL || preloadSet == NULL) {
     return;

Source/Layout/ASLayoutSpec.mm

@@ -170,11 +170,12 @@ - (ASTraitCollection *)asyncTraitCollection
 
 #pragma mark - Framework Private
 
-- (nullable NSSet<id<ASLayoutElement>> *)findDuplicatedElementsInSubtree
+#if AS_DEDUPE_LAYOUT_SPEC_TREE
+- (nullable NSHashTable<id<ASLayoutElement>> *)findDuplicatedElementsInSubtree
 {
-  NSMutableSet *result = nil;
+  NSHashTable *result = nil;
   NSUInteger count = 0;
-  [self _findDuplicatedElementsInSubtreeWithWorkingSet:[[NSMutableSet alloc] init] workingCount:&count result:&result];
+  [self _findDuplicatedElementsInSubtreeWithWorkingSet:[NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality] workingCount:&count result:&result];
   return result;
 }
 
@@ -185,7 +186,7 @@ - (ASTraitCollection *)asyncTraitCollection
  * @param workingCount The current count of the set for use in the recursion.
  * @param result The set into which to put the result. This initially points to @c nil to save time if no duplicates exist.
  */
-- (void)_findDuplicatedElementsInSubtreeWithWorkingSet:(NSMutableSet<id<ASLayoutElement>> *)workingSet workingCount:(NSUInteger *)workingCount result:(NSMutableSet<id<ASLayoutElement>>  * _Nullable *)result
+- (void)_findDuplicatedElementsInSubtreeWithWorkingSet:(NSHashTable<id<ASLayoutElement>> *)workingSet workingCount:(NSUInteger *)workingCount result:(NSHashTable<id<ASLayoutElement>>  * _Nullable *)result
 {
   Class layoutSpecClass = [ASLayoutSpec class];
 
@@ -200,7 +201,7 @@ - (void)_findDuplicatedElementsInSubtreeWithWorkingSet:(NSMutableSet<id<ASLayout
     BOOL objectAlreadyExisted = (newCount != oldCount + 1);
     if (objectAlreadyExisted) {
       if (*result == nil) {
-        *result = [[NSMutableSet alloc] init];
+        *result = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
       }
       [*result addObject:child];
     } else {
@@ -212,6 +213,7 @@ - (void)_findDuplicatedElementsInSubtreeWithWorkingSet:(NSMutableSet<id<ASLayout
     }
   }
 }
+#endif
 
 #pragma mark - Debugging
 

Source/Private/Layout/ASLayoutSpecPrivate.h

@@ -18,6 +18,12 @@
 #import <AsyncDisplayKit/ASInternalHelpers.h>
 #import <AsyncDisplayKit/ASThread.h>
 
+#if DEBUG
+  #define AS_DEDUPE_LAYOUT_SPEC_TREE 1
+#else
+  #define AS_DEDUPE_LAYOUT_SPEC_TREE 0
+#endif
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface ASLayoutSpec() {
@@ -27,10 +33,12 @@ NS_ASSUME_NONNULL_BEGIN
   NSMutableArray *_childrenArray;
 }
 
+#if AS_DEDUPE_LAYOUT_SPEC_TREE
 /**
  * Recursively search the subtree for elements that occur more than once.
  */
-- (nullable NSSet<id<ASLayoutElement>> *)findDuplicatedElementsInSubtree;
+- (nullable NSHashTable<id<ASLayoutElement>> *)findDuplicatedElementsInSubtree;
+#endif
 
 @end