Implement ASPageTable (#81)

* Implement ASPageTable
- It is a screen page table that can be used to quickly filter out objects in a certain rect without checking each and every one of them.
- ASCollectionLayoutState generates and keeps a table that maps page to layout attributes within that page.
- ASCollectionLayout (and later, ASCollectionGalleryLayoutDelegate) consults its layout state for `layoutAttributesForElementsInRect:`. This ensures the method can return as quickly as possible, especially on a large data set (I heard some people have galleries with thousands of photos!).

* Address comments

* Handle items that span multiple pages

* Make danger happy

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -389,6 +389,8 @@
 		DEFAD8131CC48914000527C4 /* ASVideoNode.mm in Sources */ = {isa = PBXBuildFile; fileRef = AEEC47E01C20C2DD00EC1693 /* ASVideoNode.mm */; };
 		E516FC7F1E9FE24200714FF4 /* ASHashing.h in Headers */ = {isa = PBXBuildFile; fileRef = E516FC7D1E9FE24200714FF4 /* ASHashing.h */; };
 		E516FC801E9FE24200714FF4 /* ASHashing.m in Sources */ = {isa = PBXBuildFile; fileRef = E516FC7E1E9FE24200714FF4 /* ASHashing.m */; };
+		E54E81FC1EB357BD00FFE8E1 /* ASPageTable.h in Headers */ = {isa = PBXBuildFile; fileRef = E54E81FA1EB357BD00FFE8E1 /* ASPageTable.h */; };
+		E54E81FD1EB357BD00FFE8E1 /* ASPageTable.m in Sources */ = {isa = PBXBuildFile; fileRef = E54E81FB1EB357BD00FFE8E1 /* ASPageTable.m */; };
 		E55D86331CA8A14000A0C26F /* ASLayoutElement.mm in Sources */ = {isa = PBXBuildFile; fileRef = E55D86311CA8A14000A0C26F /* ASLayoutElement.mm */; };
 		E5711A2C1C840C81009619D4 /* ASCollectionElement.h in Headers */ = {isa = PBXBuildFile; fileRef = E5711A2A1C840C81009619D4 /* ASCollectionElement.h */; settings = {ATTRIBUTES = (Private, ); }; };
 		E5711A301C840C96009619D4 /* ASCollectionElement.mm in Sources */ = {isa = PBXBuildFile; fileRef = E5711A2D1C840C96009619D4 /* ASCollectionElement.mm */; };
@@ -822,6 +824,8 @@
 		E516FC7E1E9FE24200714FF4 /* ASHashing.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASHashing.m; sourceTree = "<group>"; };
 		E52405B21C8FEF03004DC8E7 /* ASLayoutTransition.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASLayoutTransition.mm; sourceTree = "<group>"; };
 		E52405B41C8FEF16004DC8E7 /* ASLayoutTransition.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASLayoutTransition.h; sourceTree = "<group>"; };
+		E54E81FA1EB357BD00FFE8E1 /* ASPageTable.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASPageTable.h; sourceTree = "<group>"; };
+		E54E81FB1EB357BD00FFE8E1 /* ASPageTable.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASPageTable.m; sourceTree = "<group>"; };
 		E55D86311CA8A14000A0C26F /* ASLayoutElement.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASLayoutElement.mm; sourceTree = "<group>"; };
 		E5711A2A1C840C81009619D4 /* ASCollectionElement.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASCollectionElement.h; sourceTree = "<group>"; };
 		E5711A2D1C840C96009619D4 /* ASCollectionElement.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASCollectionElement.mm; sourceTree = "<group>"; };
@@ -1479,6 +1483,8 @@
 				E58E9E411E941D74004CFC59 /* ASCollectionLayoutDelegate.h */,
 				E58E9E3D1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.h */,
 				E58E9E3E1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.m */,
+				E54E81FA1EB357BD00FFE8E1 /* ASPageTable.h */,
+				E54E81FB1EB357BD00FFE8E1 /* ASPageTable.m */,
 			);
 			name = "Collection Layout";
 			sourceTree = "<group>";
@@ -1547,6 +1553,7 @@
 				B35061F71B010EFD0018CF92 /* ASCollectionViewProtocols.h in Headers */,
 				68FC85E31CE29B7E00EDD713 /* ASTabBarController.h in Headers */,
 				B35061FA1B010EFD0018CF92 /* ASControlNode+Subclasses.h in Headers */,
+				E54E81FC1EB357BD00FFE8E1 /* ASPageTable.h in Headers */,
 				B35061F81B010EFD0018CF92 /* ASControlNode.h in Headers */,
 				B35062171B010EFD0018CF92 /* ASDataController.h in Headers */,
 				34EFC75B1B701BAF00AD841F /* ASDimension.h in Headers */,
@@ -2056,6 +2063,7 @@
 				509E68601B3AED8E009B9150 /* ASScrollDirection.m in Sources */,
 				B35062091B010EFD0018CF92 /* ASScrollNode.mm in Sources */,
 				8BDA5FC81CDBDF95007D13B2 /* ASVideoPlayerNode.mm in Sources */,
+				E54E81FD1EB357BD00FFE8E1 /* ASPageTable.m in Sources */,
 				34EFC7721B701D0300AD841F /* ASStackLayoutSpec.mm in Sources */,
 				7AB338661C55B3420055FDE8 /* ASRelativeLayoutSpec.mm in Sources */,
 				696F01EE1DD2AF450049FBD5 /* ASEventLog.mm in Sources */,

CHANGELOG.md

@@ -10,3 +10,4 @@
 - Simplified & optimized hashing code. [Adlai Holler](https://github.com/Adlai-Holler) [#86](https://github.com/TextureGroup/Texture/pull/86)
 - Improve the performance & safety of ASDisplayNode subnodes. [Adlai Holler](https://github.com/Adlai-Holler) [#223](https://github.com/TextureGroup/Texture/pull/223)
 - Remove finalLayoutElement [Michael Schneider] (https://github.com/maicki)[#96](https://github.com/TextureGroup/Texture/pull/96)
+- Add ASPageTable - A map table for fast retrieval of objects within a certain page [Huy Nguyen](https://github.com/nguyenhuy)

Source/Details/ASCollectionFlowLayoutDelegate.h

@@ -22,9 +22,15 @@ NS_ASSUME_NONNULL_BEGIN
 
 AS_SUBCLASSING_RESTRICTED
 
+/**
+ * A thread-safe, high performant layout delegate that arranges items into a flow layout. 
+ * It uses a concurrent and multi-line ASStackLayoutSpec under the hood. Thus, per-child flex properties (i.e alignSelf, 
+ * flexShrink, flexGrow, etc - see @ASStackLayoutElement) can be set directly on cell nodes to be used
+ * to calculate the final collection layout.
+ */
 @interface ASCollectionFlowLayoutDelegate : NSObject <ASCollectionLayoutDelegate>
 
-- (instancetype)initWithScrollableDirections:(ASScrollDirection)scrollableDirections;
+- (instancetype)initWithScrollableDirections:(ASScrollDirection)scrollableDirections NS_DESIGNATED_INITIALIZER;
 
 @end
 

Source/Details/ASCollectionFlowLayoutDelegate.m

@@ -31,16 +31,12 @@
 
 - (instancetype)init
 {
-  self = [super init];
+  return [self initWithScrollableDirections:ASScrollDirectionVerticalDirections];
-  if (self) {
-    _scrollableDirections = ASScrollDirectionVerticalDirections;
-  }
-  return self;
 }
 
 - (instancetype)initWithScrollableDirections:(ASScrollDirection)scrollableDirections
 {
-  self = [self init];
+  self = [super init];
   if (self) {
     _scrollableDirections = scrollableDirections;
   }
@@ -71,9 +67,9 @@
   ASElementMap *elements = context.elements;
   NSMutableArray<ASCellNode *> *children = ASArrayByFlatMapping(elements.itemElements, ASCollectionElement *element, element.node);
   if (children.count == 0) {
-    return [[ASCollectionLayoutState alloc] initWithElements:elements
+    return [[ASCollectionLayoutState alloc] initWithContext:context
                                                 contentSize:CGSizeZero
-                                elementToLayoutArrtibutesMap:[NSMapTable weakToStrongObjectsMapTable]];
+                             elementToLayoutAttributesTable:[NSMapTable elementToLayoutAttributesTable]];
   }
   
   ASStackLayoutSpec *stackSpec = [ASStackLayoutSpec stackLayoutSpecWithDirection:ASStackLayoutDirectionHorizontal
@@ -85,7 +81,7 @@
                                                                         children:children];
   stackSpec.concurrent = YES;
   ASLayout *layout = [stackSpec layoutThatFits:[self sizeRangeThatFits:context.viewportSize]];
-  return [[ASCollectionLayoutState alloc] initWithElements:elements layout:layout];
+  return [[ASCollectionLayoutState alloc] initWithContext:context layout:layout];
 }
 
 @end

Source/Details/ASCollectionLayoutDelegate.h

@@ -27,6 +27,8 @@ NS_ASSUME_NONNULL_BEGIN
 /**
  * @abstract Returns any additional information needed for a coming layout pass with the given elements.
  *
+ * @param elements The elements to be laid out later.
+ *
  * @discussion The returned object must support equality and hashing (i.e `-isEqual:` and `-hash` must be properly implemented).
  *
  * @discussion This method will be called on main thread.

Source/Details/ASCollectionLayoutState.h

@@ -19,45 +19,83 @@
 #import <UIKit/UIKit.h>
 #import <AsyncDisplayKit/ASBaseDefines.h>
 
-@class ASElementMap, ASCollectionElement, ASLayout;
+@class ASCollectionLayoutContext, ASLayout, ASCollectionElement;
 
 NS_ASSUME_NONNULL_BEGIN
 
+@interface NSMapTable (ASCollectionLayoutConvenience)
+
++ (NSMapTable<ASCollectionElement *, UICollectionViewLayoutAttributes *> *)elementToLayoutAttributesTable;
+
+@end
+
 AS_SUBCLASSING_RESTRICTED
 @interface ASCollectionLayoutState : NSObject
 
-/// The elements used to calculate this object
+/// The context used to calculate this object
-@property (nonatomic, strong, readonly) ASElementMap *elements;
+@property (nonatomic, strong, readonly) ASCollectionLayoutContext *context;
 
+/// The final content size of the collection's layout
 @property (nonatomic, assign, readonly) CGSize contentSize;
 
-/// Element to layout attributes map. Should use weak pointers for elements.
-@property (nonatomic, strong, readonly) NSMapTable<ASCollectionElement *, UICollectionViewLayoutAttributes *> *elementToLayoutArrtibutesMap;
-
 - (instancetype)init __unavailable;
 
 /**
  * Designated initializer.
  *
- * @param elements The elements used to calculate this object
+ * @param context The context used to calculate this object
  *
  * @param contentSize The content size of the collection's layout
  *
- * @param elementToLayoutArrtibutesMap Map between elements to their layout attributes. The map may contain all elements, or a subset of them and will be updated later. 
+ * @param table A map between elements to their layout attributes. It may contain all elements, or a subset of them that will be updated later.
- * Also, it should have NSMapTableObjectPointerPersonality and NSMapTableWeakMemory as key options.
+ * It should be initialized using +[NSMapTable elementToLayoutAttributesTable] convenience initializer.
  */
-- (instancetype)initWithElements:(ASElementMap *)elements contentSize:(CGSize)contentSize elementToLayoutArrtibutesMap:(NSMapTable<ASCollectionElement *, UICollectionViewLayoutAttributes *> *)elementToLayoutArrtibutesMap NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithContext:(ASCollectionLayoutContext *)context contentSize:(CGSize)contentSize elementToLayoutAttributesTable:(NSMapTable<ASCollectionElement *, UICollectionViewLayoutAttributes *> *)table NS_DESIGNATED_INITIALIZER;
 
 /**
  * Convenience initializer.
  *
- * @param elements The elements used to calculate this object
+ * @param context The context used to calculate this object
  *
- * @param layout The layout describes size and position of all elements, or a subset of them and will be updated later.
+ * @param layout The layout describes size and position of all elements, or a subset of them and will be updated over time.
  *
- * @discussion The sublayouts that describe position of elements must be direct children of the root layout object parameter.
  */
-- (instancetype)initWithElements:(ASElementMap *)elements layout:(ASLayout *)layout;
+- (instancetype)initWithContext:(ASCollectionLayoutContext *)context layout:(ASLayout *)layout;
+
+/**
+ * Returns all layout attributes present in this object.
+ */
+- (NSArray<UICollectionViewLayoutAttributes *> *)allLayoutAttributes;
+
+/**
+ * Returns layout attributes of elements in the specified rect.
+ *
+ * @param rect The rect containing the target elements.
+ */
+- (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect;
+
+/**
+ * Returns layout attributes of the element at the specified index path.
+ *
+ * @param indexPath The index path of the item.
+ */
+- (nullable UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath;
+
+/**
+ * Returns layout attributes of the specified supplementary element.
+ *
+ * @param kind A string that identifies the type of the supplementary element.
+ *
+ * @param indexPath The index path of the element.
+ */
+- (nullable UICollectionViewLayoutAttributes *)layoutAttributesForSupplementaryElementOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath;
+
+/**
+ * Returns layout attributes of the specified element.
+ *
+ * @element The element.
+ */
+- (nullable UICollectionViewLayoutAttributes *)layoutAttributesForElement:(ASCollectionElement *)element;
 
 @end
 

Source/Details/ASCollectionLayoutState.m

@@ -20,16 +20,37 @@
 #import <AsyncDisplayKit/ASAssert.h>
 #import <AsyncDisplayKit/ASCellNode+Internal.h>
 #import <AsyncDisplayKit/ASCollectionElement.h>
+#import <AsyncDisplayKit/ASCollectionLayoutContext.h>
 #import <AsyncDisplayKit/ASElementMap.h>
 #import <AsyncDisplayKit/ASLayout.h>
+#import <AsyncDisplayKit/ASPageTable.h>
 
-@implementation ASCollectionLayoutState
+@implementation NSMapTable (ASCollectionLayoutConvenience)
 
-- (instancetype)initWithElements:(ASElementMap *)elements layout:(ASLayout *)layout
++ (NSMapTable<ASCollectionElement *,UICollectionViewLayoutAttributes *> *)elementToLayoutAttributesTable
 {
-  NSMapTable<ASCollectionElement *, UICollectionViewLayoutAttributes *> *attrsMap = [NSMapTable mapTableWithKeyOptions:(NSMapTableObjectPointerPersonality | NSMapTableWeakMemory) valueOptions:NSMapTableStrongMemory];
+  return [NSMapTable mapTableWithKeyOptions:(NSMapTableWeakMemory | NSMapTableObjectPointerPersonality) valueOptions:NSMapTableStrongMemory];
+}
+
+@end
+
+@implementation ASCollectionLayoutState {
+  NSMapTable<ASCollectionElement *,UICollectionViewLayoutAttributes *> *_elementToLayoutAttributesTable;
+  ASPageTable<id, NSMutableArray<UICollectionViewLayoutAttributes *> *> *_pageToLayoutAttributesTable;
+}
+
+- (instancetype)initWithContext:(ASCollectionLayoutContext *)context layout:(ASLayout *)layout
+{
+  ASElementMap *elements = context.elements;
+  NSMapTable *table = [NSMapTable elementToLayoutAttributesTable];
+  
   for (ASLayout *sublayout in layout.sublayouts) {
     ASCollectionElement *element = ((ASCellNode *)sublayout.layoutElement).collectionElement;
+    if (element == nil) {
+      ASDisplayNodeFailAssert(@"Element not found!");
+      continue;
+    }
+    
     NSIndexPath *indexPath = [elements indexPathForElement:element];
     NSString *supplementaryElementKind = element.supplementaryElementKind;
     
@@ -41,21 +62,74 @@
     }
     
     attrs.frame = sublayout.frame;
-    [attrsMap setObject:attrs forKey:element];
+    [table setObject:attrs forKey:element];
   }
 
-  return [self initWithElements:elements contentSize:layout.size elementToLayoutArrtibutesMap:attrsMap];
+  return [self initWithContext:context contentSize:layout.size elementToLayoutAttributesTable:table];
 }
 
-- (instancetype)initWithElements:(ASElementMap *)elements contentSize:(CGSize)contentSize elementToLayoutArrtibutesMap:(NSMapTable<ASCollectionElement *,UICollectionViewLayoutAttributes *> *)attrsMap
+- (instancetype)initWithContext:(ASCollectionLayoutContext *)context contentSize:(CGSize)contentSize elementToLayoutAttributesTable:(NSMapTable<ASCollectionElement *,UICollectionViewLayoutAttributes *> *)table
 {
   self = [super init];
   if (self) {
-    _elements = elements;
+    _context = context;
     _contentSize = contentSize;
-    _elementToLayoutArrtibutesMap = attrsMap;
+    _elementToLayoutAttributesTable = table;
+    _pageToLayoutAttributesTable = [ASPageTable pageTableWithLayoutAttributes:_elementToLayoutAttributesTable.objectEnumerator contentSize:contentSize pageSize:context.viewportSize];
   }
   return self;
 }
 
+- (NSArray<UICollectionViewLayoutAttributes *> *)allLayoutAttributes
+{
+  return [_elementToLayoutAttributesTable.objectEnumerator allObjects];
+}
+
+- (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect
+{
+  CGSize pageSize = _context.viewportSize;
+  NSPointerArray *pages = ASPageCoordinatesForPagesThatIntersectRect(rect, _contentSize, pageSize);
+  if (pages.count == 0) {
+    return @[];
+  }
+  
+  // Use a mutable set here because some items may span multiple pages
+  NSMutableSet<UICollectionViewLayoutAttributes *> *result = [NSMutableSet set];
+  for (id pagePtr in pages) {
+    ASPageCoordinate page = (ASPageCoordinate)pagePtr;
+    NSArray<UICollectionViewLayoutAttributes *> *allAttrs = [_pageToLayoutAttributesTable objectForPage:page];
+    if (allAttrs.count > 0) {
+      CGRect pageRect = ASPageCoordinateGetPageRect(page, pageSize);
+      
+      if (CGRectContainsRect(rect, pageRect)) {
+        [result addObjectsFromArray:allAttrs];
+      } else {
+        for (UICollectionViewLayoutAttributes *attrs in allAttrs) {
+          if (CGRectIntersectsRect(rect, attrs.frame)) {
+            [result addObject:attrs];
+          }
+        }
+      }
+    }
+  }
+  return [result allObjects];
+}
+
+- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath
+{
+  ASCollectionElement *element = [_context.elements elementForItemAtIndexPath:indexPath];
+  return [_elementToLayoutAttributesTable objectForKey:element];
+}
+
+- (UICollectionViewLayoutAttributes *)layoutAttributesForSupplementaryElementOfKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
+{
+  ASCollectionElement *element = [_context.elements supplementaryElementOfKind:elementKind atIndexPath:indexPath];
+  return [_elementToLayoutAttributesTable objectForKey:element];
+}
+
+- (UICollectionViewLayoutAttributes *)layoutAttributesForElement:(ASCollectionElement *)element
+{
+  return [_elementToLayoutAttributesTable objectForKey:element];
+}
+
 @end

Source/Details/ASPageTable.h

@@ -0,0 +1,123 @@
+//
+//  ASPageTable.h
+//  Texture
+//
+//  Copyright (c) 2017-present, Pinterest, Inc.  All rights reserved.
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+
+#import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+#import <AsyncDisplayKit/ASAssert.h>
+#import <AsyncDisplayKit/ASBaseDefines.h>
+
+@class ASCollectionElement;
+
+NS_ASSUME_NONNULL_BEGIN
+
+ASDISPLAYNODE_EXTERN_C_BEGIN
+
+/**
+ * Represents x and y coordinates of a page.
+ */
+typedef uintptr_t ASPageCoordinate;
+
+/**
+ * Returns a page coordinate with the given x and y values. Both of them must be less than 65,535.
+ */
+extern ASPageCoordinate ASPageCoordinateMake(uint16_t x, uint16_t y) AS_WARN_UNUSED_RESULT;
+
+/**
+ * Returns coordinate of the page that contains the specified point.
+ * Similar to CGRectContainsPoint, a point is considered inside a page if its  lie inside the page or on the minimum X or minimum Y edge.
+ *
+ * @param point The point that the page at the returned  should contain. Any negative  of the point will be corrected to 0.0
+ *
+ * @param pageSize The size of each page.
+ */
+extern ASPageCoordinate ASPageCoordinateForPageThatContainsPoint(CGPoint point, CGSize pageSize) AS_WARN_UNUSED_RESULT;
+
+extern uint16_t ASPageCoordinateGetX(ASPageCoordinate pageCoordinate) AS_WARN_UNUSED_RESULT;
+
+extern uint16_t ASPageCoordinateGetY(ASPageCoordinate pageCoordinate) AS_WARN_UNUSED_RESULT;
+
+extern CGRect ASPageCoordinateGetPageRect(ASPageCoordinate pageCoordinate, CGSize pageSize) AS_WARN_UNUSED_RESULT;
+
+/**
+ * Returns coordinate pointers for pages that intersect the specified rect. For each pointer, use ASPageCoordinateFromPointer() to get the original coordinate.
+ * The specified rect is restricted to the bounds of a content rect that has an origin of {0, 0} and a size of the given contentSize.
+ *
+ * @param rect The rect intersecting the target pages.
+ *
+ * @param contentSize The combined size of all pages.
+ *
+ * @param pageSize The size of each page.
+ */
+extern NSPointerArray * _Nullable ASPageCoordinatesForPagesThatIntersectRect(CGRect rect, CGSize contentSize, CGSize pageSize) AS_WARN_UNUSED_RESULT;
+
+ASDISPLAYNODE_EXTERN_C_END
+
+/**
+ * An alias for an NSMapTable created to store objects using ASPageCoordinates as keys.
+ *
+ * You should not call -objectForKey:, -setObject:forKey:, or -removeObjectForKey:
+ * on these objects.
+ */
+typedef NSMapTable ASPageTable;
+
+/**
+ * A category for creating & using map tables meant for storing objects using ASPage as keys.
+ */
+@interface NSMapTable<id, ObjectType> (ASPageTableMethods)
+
+/**
+ * Creates a new page table with (NSMapTableStrongMemory | NSMapTableObjectPointerPersonality) for values.
+ */
++ (ASPageTable *)pageTableForStrongObjectPointers;
+
+/**
+ * Creates a new page table with (NSMapTableWeakMemory | NSMapTableObjectPointerPersonality) for values.
+ */
++ (ASPageTable *)pageTableForWeakObjectPointers;
+
+/**
+ * Builds a new page to layout attributes from the given layout attributes.
+ *
+ * @param layoutAttributesEnumerator The layout attributes to build from
+ *
+ * @param contentSize The combined size of all pages.
+ *
+ * @param pageSize The size of each page.
+ */
++ (ASPageTable<id, NSMutableArray<UICollectionViewLayoutAttributes *> *> *)pageTableWithLayoutAttributes:(id<NSFastEnumeration>)layoutAttributesEnumerator contentSize:(CGSize)contentSize pageSize:(CGSize)pageSize;
+
+/**
+ * Retrieves the object for a given page, or nil if the page is not found.
+ *
+ * @param page A page to lookup the object for.
+ */
+- (nullable ObjectType)objectForPage:(ASPageCoordinate)page;
+
+/**
+ * Sets the given object for the associated page.
+ *
+ * @param object The object to store as value.
+ *
+ * @param page The page to use for the rect.
+ */
+- (void)setObject:(ObjectType)object forPage:(ASPageCoordinate)page;
+
+/**
+ * Removes the object for the given page, if one exists.
+ *
+ * @param page The page to remove.
+ */
+- (void)removeObjectForPage:(ASPageCoordinate)page;
+
+@end
+
+NS_ASSUME_NONNULL_END

Source/Details/ASPageTable.m

@@ -0,0 +1,151 @@
+//
+//  ASPageTable.m
+//  Texture
+//
+//  Copyright (c) 2017-present, Pinterest, Inc.  All rights reserved.
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+
+#import <AsyncDisplayKit/ASPageTable.h>
+
+extern ASPageCoordinate ASPageCoordinateMake(uint16_t x, uint16_t y)
+{
+  // Add 1 to the end result because 0 is not accepted by NSArray and NSMapTable.
+  // To avoid overflow after adding, x and y can't be UINT16_MAX (0xFFFF) **at the same time**.
+  // But for API simplification, we enforce the same restriction to both values.
+  ASDisplayNodeCAssert(x < UINT16_MAX, @"x coordinate must be less than 65,535");
+  ASDisplayNodeCAssert(y < UINT16_MAX, @"y coordinate must be less than 65,535");
+  return (x << 16) + y + 1;
+}
+
+extern ASPageCoordinate ASPageCoordinateForPageThatContainsPoint(CGPoint point, CGSize pageSize)
+{
+  return ASPageCoordinateMake((MAX(0.0, point.x) / pageSize.width), (MAX(0.0, point.y) / pageSize.height));
+}
+
+extern uint16_t ASPageCoordinateGetX(ASPageCoordinate pageCoordinate)
+{
+  return (pageCoordinate - 1) >> 16;
+}
+
+extern uint16_t ASPageCoordinateGetY(ASPageCoordinate pageCoordinate)
+{
+  return (pageCoordinate - 1) & ~(0xFFFF<<16);
+}
+
+extern CGRect ASPageCoordinateGetPageRect(ASPageCoordinate pageCoordinate, CGSize pageSize)
+{
+  CGFloat pageWidth = pageSize.width;
+  CGFloat pageHeight = pageSize.height;
+  return CGRectMake(ASPageCoordinateGetX(pageCoordinate) * pageWidth, ASPageCoordinateGetY(pageCoordinate) * pageHeight, pageWidth, pageHeight);
+}
+
+extern NSPointerArray *ASPageCoordinatesForPagesThatIntersectRect(CGRect rect, CGSize contentSize, CGSize pageSize)
+{
+  CGRect contentRect = CGRectMake(0.0, 0.0, contentSize.width, contentSize.height);
+  // Make sure the specified rect is within contentRect
+  rect = CGRectIntersection(rect, contentRect);
+  if (CGRectIsNull(rect) || CGRectIsEmpty(rect)) {
+    return nil;
+  }
+  
+  NSPointerArray *result = [NSPointerArray pointerArrayWithOptions:(NSPointerFunctionsIntegerPersonality | NSPointerFunctionsOpaqueMemory)];
+  
+  ASPageCoordinate minPage = ASPageCoordinateForPageThatContainsPoint(CGPointMake(CGRectGetMinX(rect), CGRectGetMinY(rect)), pageSize);
+  ASPageCoordinate maxPage = ASPageCoordinateForPageThatContainsPoint(CGPointMake(CGRectGetMaxX(rect), CGRectGetMaxY(rect)), pageSize);
+  if (minPage == maxPage) {
+    [result addPointer:(void *)minPage];
+    return result;
+  }
+  
+  NSUInteger minX = ASPageCoordinateGetX(minPage);
+  NSUInteger minY = ASPageCoordinateGetY(minPage);
+  NSUInteger maxX = ASPageCoordinateGetX(maxPage);
+  NSUInteger maxY = ASPageCoordinateGetY(maxPage);
+  
+  for (NSUInteger x = minX; x <= maxX; x++) {
+    for (NSUInteger y = minY; y <= maxY; y++) {
+      ASPageCoordinate page = ASPageCoordinateMake(x, y);
+      [result addPointer:(void *)page];
+    }
+  }
+  
+  return result;
+}
+
+@implementation NSMapTable (ASPageTableMethods)
+
++ (instancetype)pageTableWithValuePointerFunctions:(NSPointerFunctions *)valueFuncs
+{
+  static NSPointerFunctions *pageCoordinatesFuncs;
+  static dispatch_once_t onceToken;
+  dispatch_once(&onceToken, ^{
+    pageCoordinatesFuncs = [NSPointerFunctions pointerFunctionsWithOptions:NSPointerFunctionsIntegerPersonality | NSPointerFunctionsOpaqueMemory];
+  });
+  
+  return [[NSMapTable alloc] initWithKeyPointerFunctions:pageCoordinatesFuncs valuePointerFunctions:valueFuncs capacity:0];
+}
+
++ (ASPageTable *)pageTableForStrongObjectPointers
+{
+  static NSPointerFunctions *strongObjectPointerFuncs;
+  static dispatch_once_t onceToken;
+  dispatch_once(&onceToken, ^{
+    strongObjectPointerFuncs = [NSPointerFunctions pointerFunctionsWithOptions:NSPointerFunctionsStrongMemory];
+  });
+  return [self pageTableWithValuePointerFunctions:strongObjectPointerFuncs];
+}
+
++ (ASPageTable *)pageTableForWeakObjectPointers
+{
+  static NSPointerFunctions *weakObjectPointerFuncs;
+  static dispatch_once_t onceToken;
+  dispatch_once(&onceToken, ^{
+    weakObjectPointerFuncs = [NSPointerFunctions pointerFunctionsWithOptions:NSPointerFunctionsWeakMemory];
+  });
+  return [self pageTableWithValuePointerFunctions:weakObjectPointerFuncs];
+}
+
++ (ASPageTable<id, NSMutableArray<UICollectionViewLayoutAttributes *> *> *)pageTableWithLayoutAttributes:(id<NSFastEnumeration>)layoutAttributesEnumerator contentSize:(CGSize)contentSize pageSize:(CGSize)pageSize
+{
+  ASPageTable *result = [ASPageTable pageTableForStrongObjectPointers];
+  for (UICollectionViewLayoutAttributes *attrs in layoutAttributesEnumerator) {
+    // This attrs may span multiple pages. Make sure it's registered to all of them
+    NSPointerArray *pages = ASPageCoordinatesForPagesThatIntersectRect(attrs.frame, contentSize, pageSize);
+    
+    for (id pagePtr in pages) {
+      ASPageCoordinate page = (ASPageCoordinate)pagePtr;
+      NSMutableArray<UICollectionViewLayoutAttributes *> *attrsInPage = [result objectForPage:page];
+      if (attrsInPage == nil) {
+        attrsInPage = [NSMutableArray array];
+        [result setObject:attrsInPage forPage:page];
+      }
+      [attrsInPage addObject:attrs];
+    }
+  }  
+  return result;
+}
+
+- (id)objectForPage:(ASPageCoordinate)page
+{
+  __unsafe_unretained id key = (__bridge id)(void *)page;
+  return [self objectForKey:key];
+}
+
+- (void)setObject:(id)object forPage:(ASPageCoordinate)page
+{
+  __unsafe_unretained id key = (__bridge id)(void *)page;
+  [self setObject:object forKey:key];
+}
+
+- (void)removeObjectForPage:(ASPageCoordinate)page
+{
+  __unsafe_unretained id key = (__bridge id)(void *)page;
+  [self removeObjectForKey:key];
+}
+
+@end

Source/Private/ASCollectionLayout.mm

@@ -33,12 +33,10 @@
   ASDN::Mutex __instanceLock__; // Non-recursive mutex, ftw!
   
   // Main thread only.
-  ASCollectionLayoutState *_state;
+  ASCollectionLayoutState *_layout;
   
   // The pending state calculated ahead of time, if any.
-  ASCollectionLayoutState *_pendingState;
+  ASCollectionLayoutState *_pendingLayout;
-  // The context used to calculate _pendingState
-  ASCollectionLayoutContext *_layoutContextForPendingState;
   
   BOOL _layoutDelegateImplementsAdditionalInfoForLayoutWithElements;
 }
@@ -63,20 +61,20 @@
 - (id)layoutContextWithElements:(ASElementMap *)elements
 {
   ASDisplayNodeAssertMainThread();
+  CGSize viewportSize = [self viewportSize];
   id additionalInfo = nil;
   if (_layoutDelegateImplementsAdditionalInfoForLayoutWithElements) {
     additionalInfo = [_layoutDelegate additionalInfoForLayoutWithElements:elements];
   }
-  return [[ASCollectionLayoutContext alloc] initWithViewportSize:[self viewportSize] elements:elements additionalInfo:additionalInfo];
+  return [[ASCollectionLayoutContext alloc] initWithViewportSize:viewportSize elements:elements additionalInfo:additionalInfo];
 }
 
 - (void)prepareLayoutWithContext:(id)context
 {
-  ASCollectionLayoutState *state = [_layoutDelegate calculateLayoutWithContext:context];
+  ASCollectionLayoutState *layout = [_layoutDelegate calculateLayoutWithContext:context];
   
   ASDN::MutexLocker l(__instanceLock__);
-  _pendingState = state;
+  _pendingLayout = layout;
-  _layoutContextForPendingState = context;
 }
 
 #pragma mark - UICollectionViewLayout overrides
@@ -87,70 +85,72 @@
   [super prepareLayout];
   ASCollectionLayoutContext *context = [self layoutContextWithElements:_collectionNode.visibleElements];
   
-  ASCollectionLayoutState *state = nil;
+  ASCollectionLayoutState *layout = nil;
   {
     ASDN::MutexLocker l(__instanceLock__);
-    if (_pendingState != nil && ASObjectIsEqual(_layoutContextForPendingState, context)) {
+    if (_pendingLayout != nil && ASObjectIsEqual(_pendingLayout.context, context)) {
-      // Looks like we can use the pending state. Great!
+      // Looks like we can use the pending layout. Great!
-      state = _pendingState;
+      layout = _pendingLayout;
-      _pendingState = nil;
+      _pendingLayout = nil;
-      _layoutContextForPendingState = nil;
     }
   }
   
-  if (state == nil) {
+  if (layout == nil) {
-    state = [_layoutDelegate calculateLayoutWithContext:context];
+    layout = [_layoutDelegate calculateLayoutWithContext:context];
   }
   
-  _state = state;
+  _layout = layout;
 }
 
 - (void)invalidateLayout
 {
   ASDisplayNodeAssertMainThread();
   [super invalidateLayout];
-  _state = nil;
+  _layout = nil;
 }
 
 - (CGSize)collectionViewContentSize
 {
   ASDisplayNodeAssertMainThread();
-  ASDisplayNodeAssertNotNil(_state, @"Collection layout state should not be nil at this point");
+  ASDisplayNodeAssertNotNil(_layout, @"Collection layout state should not be nil at this point");
-  return _state.contentSize;
+  return _layout.contentSize;
 }
 
 - (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect
 {
-  NSMutableArray *attributesInRect = [NSMutableArray array];
+  ASDisplayNodeAssertMainThread();
-  NSMapTable *attrsMap = _state.elementToLayoutArrtibutesMap;
+  NSArray<UICollectionViewLayoutAttributes *> *result = [_layout layoutAttributesForElementsInRect:rect];
-  for (ASCollectionElement *element in attrsMap) {
+  
-    UICollectionViewLayoutAttributes *attrs = [attrsMap objectForKey:element];
+  ASElementMap *elements = _layout.context.elements;
-    if (CGRectIntersectsRect(rect, attrs.frame)) {
+  for (UICollectionViewLayoutAttributes *attrs in result) {
+    ASCollectionElement *element = [elements elementForLayoutAttributes:attrs];
     [ASCollectionLayout setSize:attrs.frame.size toElement:element];
-      [attributesInRect addObject:attrs];
   }
-  }
+  
-  return attributesInRect;
+  return result;
 }
 
 - (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath
 {
-  ASCollectionLayoutState *state = _state;
+  ASCollectionElement *element = [_layout.context.elements elementForItemAtIndexPath:indexPath];
-  ASCollectionElement *element = [state.elements elementForItemAtIndexPath:indexPath];
+  UICollectionViewLayoutAttributes *attrs = [_layout layoutAttributesForElement:element];
-  UICollectionViewLayoutAttributes *attrs = [state.elementToLayoutArrtibutesMap objectForKey:element];
   [ASCollectionLayout setSize:attrs.frame.size toElement:element];
   return attrs;
 }
 
 - (UICollectionViewLayoutAttributes *)layoutAttributesForSupplementaryViewOfKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
 {
-  ASCollectionLayoutState *state = _state;
+  ASCollectionElement *element = [_layout.context.elements supplementaryElementOfKind:elementKind atIndexPath:indexPath];
-  ASCollectionElement *element = [state.elements supplementaryElementOfKind:elementKind atIndexPath:indexPath];
+  UICollectionViewLayoutAttributes *attrs = [_layout layoutAttributesForElement:element];
-  UICollectionViewLayoutAttributes *attrs = [state.elementToLayoutArrtibutesMap objectForKey:element];
   [ASCollectionLayout setSize:attrs.frame.size toElement:element];
   return attrs;
 }
 
+- (BOOL)shouldInvalidateLayoutForBoundsChange:(CGRect)newBounds
+{
+  return (! CGSizeEqualToSize([self viewportSize], newBounds.size));
+}
+
 #pragma mark - Private methods
 
 + (void)setSize:(CGSize)size toElement:(ASCollectionElement *)element