Introduce ASCollectionLayout and friends (#3130)

* Introduce ASCollectionViewLayout
- `ASCollectionViewLayout` is an async `UICollectionViewLayout` that encapsulates its layout calculation logic into a separate thread-safe object which can be used ahead of time and/or on multiple threads.
- `ASDataController` now can prepare for a new layout resulted from a change set before `ASCollectionView` even knows about it. By the time the change set it ready to be consumed by `ASCollectionView`, its new layout is also ready.
- New `ASCollectionViewLayoutCalculating` protocol that is simple and generic enough that many types of calculators can be built on top. `ASCollectionViewLayoutSpecCalculator` conforms to `ASCollectionViewLayoutCalculating` protocol and can be backed by any layout spec (e.g `ASStackLayoutSpec`, `PIMasonryLayoutSpec`, etc). We can even build a `ASCollectionViewLayoutYogaCalculator` that uses Yoga internally.
- A built-in `ASCollectionViewFlowLayoutCalculator` that is a subclass of `ASCollectionViewLayoutSpecCalculator` and uses a multi-threaded multi-line `ASStackLayoutSpec` internally. The result is a performant and thread-safe flow layout calculator.
- Finally, `ASCollectionViewLayout` can be subclassed to handle a specific type of calculator with optimizations implemented based on the knowledge of such calculator. For example, `ASCollectionViewFlowLayout` can have a highly optimized implementation of `-layoutAttributesForElementsInRect:`.

Protocolize layout calculator providing and consuming

Add flex wrap documentation

Add a `multithreaded` flag to ASStackLayoutSpec that forces it to dispatch even if it's off main
- Update ASCollectionViewFlowLayoutSpecCalculator to use that flag.

Minor change in ASCollectionViewLayout

Implement Mosaic layout calculator

Minor change

Fix project file

Rename and fix project file

Skip fetching constrained size only if a layout calculator is available

Update examples/ASCollectionView

Remove unnecessary change in ASTableView

Address comments

Rename collection view calculator protocols

Minor changes after rebasing with master

Add ASLegacyCollectionLayoutCalculator for backward compatibility

Remove ASCollectionLayoutSpecCalculator

Remove ASLegacyCollectionLayoutCalculator

Introduce ASCollectionLayout
- A wrapper object that contains content size and an element to rect table.
- Collection layout calculators to return this new object instead of an ASLayout.

Before adding a content cache

Finishing hooking up ASCollectionLayoutDataSource to ASCollectionNode

Stash

Finish ASCollectionLayout

Rough impl of ASCollectionFlowLayout

Revert changes in CustomCollectionView example

Move ASRectTable back to Private

* Rename ASCollectionContentAttributes to ASCollectionLayoutState

* Address other comments

* Introduce ASCollectionLayoutDelegate and make ASCollectionLayout private

* Address comments

* API tweaks:
- Replace `-layoutContextWithElementMap:` in ASCollectionLayoutDelegate with `-additionalInfoForLayoutWithElements:`. The returned object is then stored in ASCollectionLayoutContext for later lookups.
- ASCollectionLayoutContext has no public initializer.
- ASDataControllerLayoutDelegate no longer requires a context of type ASCollectionLayoutContext but simply an `id`. This helps decouple ASDataController and ASCollectionLayout.
- Rename `elementMap` to `elements`.
- Rename `visibleMap` to `visibleElements`.
- Other minor changes.

* Rename ASCGSizeHash to ASHashFromCGSize

* Make sure to call super in -[ASCollectionLayout prepareLayout]

* Update example/ASCollectionView to use ASCollectionFlowLayoutDelegate

* Remove unnecessary change

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -375,6 +375,20 @@
 		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 */; };
+		E58E9E421E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.h in Headers */ = {isa = PBXBuildFile; fileRef = E58E9E3D1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.h */; settings = {ATTRIBUTES = (Public, ); }; };
+		E58E9E431E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.m in Sources */ = {isa = PBXBuildFile; fileRef = E58E9E3E1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.m */; };
+		E58E9E441E941D74004CFC59 /* ASCollectionLayoutContext.h in Headers */ = {isa = PBXBuildFile; fileRef = E58E9E3F1E941D74004CFC59 /* ASCollectionLayoutContext.h */; settings = {ATTRIBUTES = (Public, ); }; };
+		E58E9E451E941D74004CFC59 /* ASCollectionLayoutContext.mm in Sources */ = {isa = PBXBuildFile; fileRef = E58E9E401E941D74004CFC59 /* ASCollectionLayoutContext.mm */; };
+		E58E9E461E941D74004CFC59 /* ASCollectionLayoutDelegate.h in Headers */ = {isa = PBXBuildFile; fileRef = E58E9E411E941D74004CFC59 /* ASCollectionLayoutDelegate.h */; settings = {ATTRIBUTES = (Public, ); }; };
+		E58E9E491E941DA5004CFC59 /* ASCollectionLayout.h in Headers */ = {isa = PBXBuildFile; fileRef = E58E9E471E941DA5004CFC59 /* ASCollectionLayout.h */; settings = {ATTRIBUTES = (Private, ); }; };
+		E58E9E4A1E941DA5004CFC59 /* ASCollectionLayout.mm in Sources */ = {isa = PBXBuildFile; fileRef = E58E9E481E941DA5004CFC59 /* ASCollectionLayout.mm */; };
+		E5ABAC7B1E8564EE007AC15C /* ASRectTable.h in Headers */ = {isa = PBXBuildFile; fileRef = E5ABAC791E8564EE007AC15C /* ASRectTable.h */; };
+		E5ABAC7C1E8564EE007AC15C /* ASRectTable.m in Sources */ = {isa = PBXBuildFile; fileRef = E5ABAC7A1E8564EE007AC15C /* ASRectTable.m */; };
+		E5B077FF1E69F4EB00C24B5B /* ASElementMap.h in Headers */ = {isa = PBXBuildFile; fileRef = E5B077FD1E69F4EB00C24B5B /* ASElementMap.h */; settings = {ATTRIBUTES = (Public, ); }; };
+		E5B078001E69F4EB00C24B5B /* ASElementMap.m in Sources */ = {isa = PBXBuildFile; fileRef = E5B077FE1E69F4EB00C24B5B /* ASElementMap.m */; };
+		E5B5B9D11E9BAD9800A6B726 /* ASCollectionLayoutContext+Private.h in Headers */ = {isa = PBXBuildFile; fileRef = E5B5B9D01E9BAD9800A6B726 /* ASCollectionLayoutContext+Private.h */; };
+		E5E281741E71C833006B67C2 /* ASCollectionLayoutState.h in Headers */ = {isa = PBXBuildFile; fileRef = E5E281731E71C833006B67C2 /* ASCollectionLayoutState.h */; settings = {ATTRIBUTES = (Public, ); }; };
+		E5E281761E71C845006B67C2 /* ASCollectionLayoutState.m in Sources */ = {isa = PBXBuildFile; fileRef = E5E281751E71C845006B67C2 /* ASCollectionLayoutState.m */; };
 		F711994E1D20C21100568860 /* ASDisplayNodeExtrasTests.m in Sources */ = {isa = PBXBuildFile; fileRef = F711994D1D20C21100568860 /* ASDisplayNodeExtrasTests.m */; };
 /* End PBXBuildFile section */
 
@@ -440,7 +454,7 @@
 		058D09DF195D050800B7D73C /* ASTextNode.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASTextNode.h; sourceTree = "<group>"; };
 		058D09E0195D050800B7D73C /* ASTextNode.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; lineEnding = 0; path = ASTextNode.mm; sourceTree = "<group>"; };
 		058D09E2195D050800B7D73C /* _ASDisplayLayer.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = _ASDisplayLayer.h; sourceTree = "<group>"; };
-		058D09E3195D050800B7D73C /* _ASDisplayLayer.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; lineEnding = 0; path = _ASDisplayLayer.mm; sourceTree = "<group>"; xcLanguageSpecificationIdentifier = xcode.lang.objcpp; };
+		058D09E3195D050800B7D73C /* _ASDisplayLayer.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; lineEnding = 0; path = _ASDisplayLayer.mm; sourceTree = "<group>"; };
 		058D09E4195D050800B7D73C /* _ASDisplayView.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = _ASDisplayView.h; sourceTree = "<group>"; };
 		058D09E5195D050800B7D73C /* _ASDisplayView.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = _ASDisplayView.mm; sourceTree = "<group>"; };
 		058D09E6195D050800B7D73C /* ASHighlightOverlayLayer.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASHighlightOverlayLayer.h; sourceTree = "<group>"; };
@@ -450,7 +464,7 @@
 		058D09F5195D050800B7D73C /* NSMutableAttributedString+TextKitAdditions.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "NSMutableAttributedString+TextKitAdditions.h"; sourceTree = "<group>"; };
 		058D09F6195D050800B7D73C /* NSMutableAttributedString+TextKitAdditions.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "NSMutableAttributedString+TextKitAdditions.m"; sourceTree = "<group>"; };
 		058D09F8195D050800B7D73C /* _ASAsyncTransaction.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = _ASAsyncTransaction.h; sourceTree = "<group>"; };
-		058D09F9195D050800B7D73C /* _ASAsyncTransaction.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; lineEnding = 0; path = _ASAsyncTransaction.mm; sourceTree = "<group>"; xcLanguageSpecificationIdentifier = xcode.lang.objcpp; };
+		058D09F9195D050800B7D73C /* _ASAsyncTransaction.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; lineEnding = 0; path = _ASAsyncTransaction.mm; sourceTree = "<group>"; };
 		058D09FA195D050800B7D73C /* _ASAsyncTransactionContainer+Private.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "_ASAsyncTransactionContainer+Private.h"; sourceTree = "<group>"; };
 		058D09FB195D050800B7D73C /* _ASAsyncTransactionContainer.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = _ASAsyncTransactionContainer.h; sourceTree = "<group>"; };
 		058D09FC195D050800B7D73C /* _ASAsyncTransactionContainer.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; lineEnding = 0; path = _ASAsyncTransactionContainer.m; sourceTree = "<group>"; xcLanguageSpecificationIdentifier = xcode.lang.objc; };
@@ -777,6 +791,20 @@
 		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>"; };
+		E58E9E3D1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASCollectionFlowLayoutDelegate.h; sourceTree = "<group>"; };
+		E58E9E3E1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASCollectionFlowLayoutDelegate.m; sourceTree = "<group>"; };
+		E58E9E3F1E941D74004CFC59 /* ASCollectionLayoutContext.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASCollectionLayoutContext.h; sourceTree = "<group>"; };
+		E58E9E401E941D74004CFC59 /* ASCollectionLayoutContext.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASCollectionLayoutContext.mm; sourceTree = "<group>"; };
+		E58E9E411E941D74004CFC59 /* ASCollectionLayoutDelegate.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASCollectionLayoutDelegate.h; sourceTree = "<group>"; };
+		E58E9E471E941DA5004CFC59 /* ASCollectionLayout.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASCollectionLayout.h; sourceTree = "<group>"; };
+		E58E9E481E941DA5004CFC59 /* ASCollectionLayout.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASCollectionLayout.mm; sourceTree = "<group>"; };
+		E5ABAC791E8564EE007AC15C /* ASRectTable.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASRectTable.h; sourceTree = "<group>"; };
+		E5ABAC7A1E8564EE007AC15C /* ASRectTable.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASRectTable.m; sourceTree = "<group>"; };
+		E5B077FD1E69F4EB00C24B5B /* ASElementMap.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASElementMap.h; sourceTree = "<group>"; };
+		E5B077FE1E69F4EB00C24B5B /* ASElementMap.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASElementMap.m; sourceTree = "<group>"; };
+		E5B5B9D01E9BAD9800A6B726 /* ASCollectionLayoutContext+Private.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "ASCollectionLayoutContext+Private.h"; sourceTree = "<group>"; };
+		E5E281731E71C833006B67C2 /* ASCollectionLayoutState.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASCollectionLayoutState.h; sourceTree = "<group>"; };
+		E5E281751E71C845006B67C2 /* ASCollectionLayoutState.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASCollectionLayoutState.m; sourceTree = "<group>"; };
 		EFA731F0396842FF8AB635EE /* libPods-AsyncDisplayKitTests.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-AsyncDisplayKitTests.a"; sourceTree = BUILT_PRODUCTS_DIR; };
 		F711994D1D20C21100568860 /* ASDisplayNodeExtrasTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASDisplayNodeExtrasTests.m; sourceTree = "<group>"; };
 		FB07EABBCF28656C6297BC2D /* Pods-AsyncDisplayKitTests.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-AsyncDisplayKitTests.debug.xcconfig"; path = "Pods/Target Support Files/Pods-AsyncDisplayKitTests/Pods-AsyncDisplayKitTests.debug.xcconfig"; sourceTree = "<group>"; };
@@ -1117,6 +1145,7 @@
 				205F0E1F1B376416007741D0 /* CoreGraphics+ASConvenience.h */,
 				205F0E201B376416007741D0 /* CoreGraphics+ASConvenience.m */,
 				25B171EA1C12242700508A7A /* Data Controller */,
+				E5B077EB1E6843AF00C24B5B /* Collection Layout */,
 				DBC452D91C5BF64600B16017 /* NSArray+Diffing.h */,
 				DBC452DA1C5BF64600B16017 /* NSArray+Diffing.m */,
 				CC4981BA1D1C7F65004E13CC /* NSIndexSet+ASHelpers.h */,
@@ -1150,10 +1179,8 @@
 			children = (
 				CC2F65EC1E5FFB1600DA57C9 /* ASMutableElementMap.h */,
 				CC2F65ED1E5FFB1600DA57C9 /* ASMutableElementMap.m */,
-				CC0349FF1E5FAF9700626263 /* ASElementMap.h */,
+				E5ABAC791E8564EE007AC15C /* ASRectTable.h */,
-				CC034A001E5FAF9700626263 /* ASElementMap.m */,
+				E5ABAC7A1E8564EE007AC15C /* ASRectTable.m */,
-				CC034A0B1E60C3D500626263 /* ASRectTable.h */,
-				CC034A0C1E60C3D500626263 /* ASRectTable.m */,
 				CC55A70F1E52A0F200594372 /* ASResponderChainEnumerator.h */,
 				CC55A7101E52A0F200594372 /* ASResponderChainEnumerator.m */,
 				6947B0BB1E36B4E30007C478 /* Layout */,
@@ -1171,6 +1198,9 @@
 				044285051BAA63FE00D16268 /* ASBatchFetching.h */,
 				044285061BAA63FE00D16268 /* ASBatchFetching.m */,
 				CC87BB941DA8193C0090E380 /* ASCellNode+Internal.h */,
+				E58E9E471E941DA5004CFC59 /* ASCollectionLayout.h */,
+				E58E9E481E941DA5004CFC59 /* ASCollectionLayout.mm */,
+				E5B5B9D01E9BAD9800A6B726 /* ASCollectionLayoutContext+Private.h */,
 				CC2E317F1DAC353700EEE891 /* ASCollectionView+Undeprecated.h */,
 				CC0F885A1E42807F00576FED /* ASCollectionViewFlowLayoutInspector.h */,
 				CC0F88591E42807F00576FED /* ASCollectionViewFlowLayoutInspector.m */,
@@ -1267,6 +1297,8 @@
 				4640521A1A3F83C40061C0BA /* ASDataController.mm */,
 				E5711A2A1C840C81009619D4 /* ASCollectionElement.h */,
 				E5711A2D1C840C96009619D4 /* ASCollectionElement.mm */,
+				E5B077FD1E69F4EB00C24B5B /* ASElementMap.h */,
+				E5B077FE1E69F4EB00C24B5B /* ASElementMap.m */,
 				AC6145401D8AFAE8003D62A2 /* ASSection.h */,
 				AC6145421D8AFD4F003D62A2 /* ASSection.m */,
 			);
@@ -1382,6 +1414,20 @@
 			path = Debug;
 			sourceTree = "<group>";
 		};
+		E5B077EB1E6843AF00C24B5B /* Collection Layout */ = {
+			isa = PBXGroup;
+			children = (
+				E58E9E3F1E941D74004CFC59 /* ASCollectionLayoutContext.h */,
+				E58E9E401E941D74004CFC59 /* ASCollectionLayoutContext.mm */,
+				E5E281731E71C833006B67C2 /* ASCollectionLayoutState.h */,
+				E5E281751E71C845006B67C2 /* ASCollectionLayoutState.m */,
+				E58E9E411E941D74004CFC59 /* ASCollectionLayoutDelegate.h */,
+				E58E9E3D1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.h */,
+				E58E9E3E1E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.m */,
+			);
+			name = "Collection Layout";
+			sourceTree = "<group>";
+		};
 		FD40E2760492F0CAAEAD552D /* Pods */ = {
 			isa = PBXGroup;
 			children = (
@@ -1399,6 +1445,11 @@
 			isa = PBXHeadersBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				E58E9E461E941D74004CFC59 /* ASCollectionLayoutDelegate.h in Headers */,
+				E5E281741E71C833006B67C2 /* ASCollectionLayoutState.h in Headers */,
+				E5B077FF1E69F4EB00C24B5B /* ASElementMap.h in Headers */,
+				E58E9E441E941D74004CFC59 /* ASCollectionLayoutContext.h in Headers */,
+				E58E9E421E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.h in Headers */,
 				696F01EC1DD2AF450049FBD5 /* ASEventLog.h in Headers */,
 				690C35641E055C7B00069B91 /* ASDimensionInternal.h in Headers */,
 				690C356B1E05680300069B91 /* ASDimensionDeprecated.h in Headers */,
@@ -1428,7 +1479,6 @@
 				B35062591B010F070018CF92 /* ASBaseDefines.h in Headers */,
 				B35062131B010EFD0018CF92 /* ASBasicImageDownloader.h in Headers */,
 				B35062151B010EFD0018CF92 /* ASBatchContext.h in Headers */,
-				CC034A0D1E60C3D500626263 /* ASRectTable.h in Headers */,
 				B35061F31B010EFD0018CF92 /* ASCellNode.h in Headers */,
 				34EFC7631B701CBF00AD841F /* ASCenterLayoutSpec.h in Headers */,
 				CC55A7111E52A0F200594372 /* ASResponderChainEnumerator.h in Headers */,
@@ -1497,13 +1547,14 @@
 				CC57EAF71E3939350034C595 /* ASCollectionView+Undeprecated.h in Headers */,
 				B35062521B010EFD0018CF92 /* ASDisplayNodeInternal.h in Headers */,
 				AC7A2C181BDE11DF0093FE1A /* ASTableViewInternal.h in Headers */,
-				CC034A011E5FAF9700626263 /* ASElementMap.h in Headers */,
 				B35062531B010EFD0018CF92 /* ASImageNode+CGExtras.h in Headers */,
+				E58E9E491E941DA5004CFC59 /* ASCollectionLayout.h in Headers */,
 				254C6B7F1BF94DF4003EC431 /* ASTextKitTruncating.h in Headers */,
 				CC58AA4B1E398E1D002C8CB4 /* ASBlockTypes.h in Headers */,
 				6977965F1D8AC8D3007E93D7 /* ASLayoutSpec+Subclasses.h in Headers */,
 				692BE8D71E36B65B00C86D87 /* ASLayoutSpecPrivate.h in Headers */,
 				34EFC75D1B701BE900AD841F /* ASInternalHelpers.h in Headers */,
+				E5B5B9D11E9BAD9800A6B726 /* ASCollectionLayoutContext+Private.h in Headers */,
 				DEC146B71C37A16A004A0EE7 /* ASCollectionInternal.h in Headers */,
 				68B8A4E21CBDB958007E4543 /* ASWeakProxy.h in Headers */,
 				9F98C0271DBE29FC00476D92 /* ASControlTargetAction.h in Headers */,
@@ -1541,6 +1592,7 @@
 				B13CA0F81C519EBA00E031AB /* ASCollectionViewLayoutFacilitatorProtocol.h in Headers */,
 				B35062061B010EFD0018CF92 /* ASNetworkImageNode.h in Headers */,
 				34EFC76C1B701CED00AD841F /* ASOverlayLayoutSpec.h in Headers */,
+				E5ABAC7B1E8564EE007AC15C /* ASRectTable.h in Headers */,
 				B35062261B010EFD0018CF92 /* ASRangeController.h in Headers */,
 				34EFC76E1B701CF400AD841F /* ASRatioLayoutSpec.h in Headers */,
 				DB55C2671C641AE4004EDCF5 /* ASContextTransitioning.h in Headers */,
@@ -1852,6 +1904,7 @@
 				509E68621B3AEDA5009B9150 /* ASAbstractLayoutController.mm in Sources */,
 				254C6B861BF94F8A003EC431 /* ASTextKitContext.mm in Sources */,
 				DBDB83971C6E879900D0098C /* ASPagerFlowLayout.m in Sources */,
+				E5B078001E69F4EB00C24B5B /* ASElementMap.m in Sources */,
 				9C8898BC1C738BA800D6B02E /* ASTextKitFontSizeAdjuster.mm in Sources */,
 				690ED59B1E36D118000627C0 /* ASImageNode+tvOS.m in Sources */,
 				34EFC7621B701CA400AD841F /* ASBackgroundLayoutSpec.mm in Sources */,
@@ -1859,6 +1912,7 @@
 				B35062141B010EFD0018CF92 /* ASBasicImageDownloader.mm in Sources */,
 				B35062161B010EFD0018CF92 /* ASBatchContext.mm in Sources */,
 				AC47D9421B3B891B00AAEE9D /* ASCellNode.mm in Sources */,
+				E58E9E451E941D74004CFC59 /* ASCollectionLayoutContext.mm in Sources */,
 				34EFC7641B701CC600AD841F /* ASCenterLayoutSpec.mm in Sources */,
 				18C2ED831B9B7DE800F627B3 /* ASCollectionNode.mm in Sources */,
 				E55D86331CA8A14000A0C26F /* ASLayoutElement.mm in Sources */,
@@ -1867,7 +1921,6 @@
 				68B8A4E41CBDB958007E4543 /* ASWeakProxy.m in Sources */,
 				9C70F20A1CDBE949007D6C76 /* ASTableNode.mm in Sources */,
 				69CB62AE1CB8165900024920 /* _ASDisplayViewAccessiblity.mm in Sources */,
-				CC034A021E5FAF9700626263 /* ASElementMap.m in Sources */,
 				B35061F61B010EFD0018CF92 /* ASCollectionView.mm in Sources */,
 				509E68641B3AEDB7009B9150 /* ASCollectionViewLayoutController.m in Sources */,
 				B35061F91B010EFD0018CF92 /* ASControlNode.mm in Sources */,
@@ -1881,10 +1934,10 @@
 				B35062501B010EFD0018CF92 /* ASDisplayNode+DebugTiming.mm in Sources */,
 				DEC146B91C37A16A004A0EE7 /* ASCollectionInternal.m in Sources */,
 				254C6B891BF94F8A003EC431 /* ASTextKitRenderer+Positioning.mm in Sources */,
-				CC034A0E1E60C3D500626263 /* ASRectTable.m in Sources */,
 				68355B341CB579B9001D4E68 /* ASImageNode+AnimatedImage.mm in Sources */,
 				E5711A301C840C96009619D4 /* ASCollectionElement.mm in Sources */,
 				B35062511B010EFD0018CF92 /* ASDisplayNode+UIViewBridge.mm in Sources */,
+				E5E281761E71C845006B67C2 /* ASCollectionLayoutState.m in Sources */,
 				B35061FC1B010EFD0018CF92 /* ASDisplayNode.mm in Sources */,
 				B35061FF1B010EFD0018CF92 /* ASDisplayNodeExtras.mm in Sources */,
 				B35062011B010EFD0018CF92 /* ASEditableTextNode.mm in Sources */,
@@ -1896,6 +1949,7 @@
 				CC0F885F1E4280B800576FED /* _ASCollectionViewCell.m in Sources */,
 				CC2F65EF1E5FFB1600DA57C9 /* ASMutableElementMap.m in Sources */,
 				B35062541B010EFD0018CF92 /* ASImageNode+CGExtras.m in Sources */,
+				E58E9E4A1E941DA5004CFC59 /* ASCollectionLayout.mm in Sources */,
 				6947B0C01E36B4E30007C478 /* ASStackUnpositionedLayout.mm in Sources */,
 				68355B401CB57A69001D4E68 /* ASImageContainerProtocolCategories.m in Sources */,
 				B35062031B010EFD0018CF92 /* ASImageNode.mm in Sources */,
@@ -1923,6 +1977,7 @@
 				254C6B8B1BF94F8A003EC431 /* ASTextKitShadower.mm in Sources */,
 				254C6B851BF94F8A003EC431 /* ASTextKitAttributes.mm in Sources */,
 				90FC784F1E4BFE1B00383C5A /* ASDisplayNode+Yoga.mm in Sources */,
+				E5ABAC7C1E8564EE007AC15C /* ASRectTable.m in Sources */,
 				509E68601B3AED8E009B9150 /* ASScrollDirection.m in Sources */,
 				B35062091B010EFD0018CF92 /* ASScrollNode.mm in Sources */,
 				8BDA5FC81CDBDF95007D13B2 /* ASVideoPlayerNode.mm in Sources */,
@@ -1932,6 +1987,7 @@
 				9C70F2051CDA4F06007D6C76 /* ASTraitCollection.m in Sources */,
 				83A7D95B1D44547700BF333E /* ASWeakMap.m in Sources */,
 				CC034A0A1E60BEB400626263 /* ASDisplayNode+Convenience.m in Sources */,
+				E58E9E431E941D74004CFC59 /* ASCollectionFlowLayoutDelegate.m in Sources */,
 				DE84918E1C8FFF9F003D89E9 /* ASRunLoopQueue.mm in Sources */,
 				68FC85E51CE29B7E00EDD713 /* ASTabBarController.m in Sources */,
 				34EFC7741B701D0A00AD841F /* ASAbsoluteLayoutSpec.mm in Sources */,

Source/ASCollectionNode+Beta.h

@@ -10,7 +10,8 @@
 
 #import <AsyncDisplayKit/ASCollectionNode.h>
 
-@protocol ASCollectionViewLayoutFacilitatorProtocol;
+@protocol ASCollectionViewLayoutFacilitatorProtocol, ASCollectionLayoutDelegate;
+@class ASElementMap;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -23,8 +24,17 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property (strong, nonatomic, nullable) Class collectionViewClass;
 
+/**
+ * The elements that are currently displayed. The "UIKit index space". Must be accessed on main thread.
+ */
+@property (strong, nonatomic, readonly) ASElementMap *visibleElements;
+
+@property (strong, readonly, nullable) id<ASCollectionLayoutDelegate> layoutDelegate;
+
 - (instancetype)initWithFrame:(CGRect)frame collectionViewLayout:(UICollectionViewLayout *)layout layoutFacilitator:(nullable id<ASCollectionViewLayoutFacilitatorProtocol>)layoutFacilitator;
 
+- (instancetype)initWithLayoutDelegate:(id<ASCollectionLayoutDelegate>)layoutDelegate layoutFacilitator:(nullable id<ASCollectionViewLayoutFacilitatorProtocol>)layoutFacilitator;
+
 - (void)beginUpdates ASDISPLAYNODE_DEPRECATED_MSG("Use -performBatchUpdates:completion: instead.");
 
 - (void)endUpdatesAnimated:(BOOL)animated ASDISPLAYNODE_DEPRECATED_MSG("Use -performBatchUpdates:completion: instead.");

Source/ASCollectionNode.h

@@ -96,6 +96,12 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property (nonatomic, assign) BOOL allowsMultipleSelection;
 
+/**
+ * The layout used to organize the node's items.
+ *
+ * @discussion Assigning a new layout object to this property causes the new layout to be applied (without animations) to the node’s items.
+ */
+@property (nonatomic, strong) UICollectionViewLayout *collectionViewLayout;
 
 /**
  * Tuning parameters for a range type in full mode.

Source/ASCollectionNode.mm

@@ -11,10 +11,12 @@
 //
 
 #import <AsyncDisplayKit/ASCollectionNode.h>
+#import <AsyncDisplayKit/ASCollectionNode+Beta.h>
 
 #import <AsyncDisplayKit/ASCollectionElement.h>
 #import <AsyncDisplayKit/ASElementMap.h>
 #import <AsyncDisplayKit/ASCollectionInternal.h>
+#import <AsyncDisplayKit/ASCollectionLayout.h>
 #import <AsyncDisplayKit/ASCollectionViewLayoutFacilitatorProtocol.h>
 #import <AsyncDisplayKit/ASDisplayNode+Beta.h>
 #import <AsyncDisplayKit/ASDisplayNode+Subclasses.h>
@@ -33,6 +35,7 @@
 @interface _ASCollectionPendingState : NSObject
 @property (weak, nonatomic) id <ASCollectionDelegate>   delegate;
 @property (weak, nonatomic) id <ASCollectionDataSource> dataSource;
+@property (strong, nonatomic) UICollectionViewLayout *collectionViewLayout;
 @property (nonatomic, assign) ASLayoutRangeMode rangeMode;
 @property (nonatomic, assign) BOOL allowsSelection; // default is YES
 @property (nonatomic, assign) BOOL allowsMultipleSelection; // default is NO
@@ -133,13 +136,21 @@
   return [self initWithFrame:frame collectionViewLayout:layout layoutFacilitator:nil];
 }
 
+- (instancetype)initWithLayoutDelegate:(id<ASCollectionLayoutDelegate>)layoutDelegate layoutFacilitator:(id<ASCollectionViewLayoutFacilitatorProtocol>)layoutFacilitator
+{
+  return [self initWithFrame:CGRectZero collectionViewLayout:[[ASCollectionLayout alloc] initWithLayoutDelegate:layoutDelegate] layoutFacilitator:layoutFacilitator];
+}
+
 - (instancetype)initWithFrame:(CGRect)frame collectionViewLayout:(UICollectionViewLayout *)layout layoutFacilitator:(id<ASCollectionViewLayoutFacilitatorProtocol>)layoutFacilitator
 {
   if (self = [super init]) {
+    // Must call the setter here to make sure pendingState is created and the layout is configured.
+    [self setCollectionViewLayout:layout];
+    
     __weak __typeof__(self) weakSelf = self;
     [self setViewBlock:^{
       __typeof__(self) strongSelf = weakSelf;
-      return [[[strongSelf collectionViewClass] alloc] _initWithFrame:frame collectionViewLayout:layout layoutFacilitator:layoutFacilitator eventLog:ASDisplayNodeGetEventLog(strongSelf)];
+      return [[[strongSelf collectionViewClass] alloc] _initWithFrame:frame collectionViewLayout:strongSelf->_pendingState.collectionViewLayout layoutFacilitator:layoutFacilitator eventLog:ASDisplayNodeGetEventLog(strongSelf)];
     }];
   }
   return self;
@@ -166,6 +177,8 @@
     if (pendingState.rangeMode != ASLayoutRangeModeUnspecified) {
       [view.rangeController updateCurrentRangeWithMode:pendingState.rangeMode];
     }
+    
+    // Don't need to set collectionViewLayout to the view as the layout was already used to init the view in view block.
   }
 }
 
@@ -340,6 +353,42 @@
   }
 }
 
+- (void)setCollectionViewLayout:(UICollectionViewLayout *)layout
+{
+  if ([self pendingState]) {
+    [self _configureCollectionViewLayout:layout];
+    _pendingState.collectionViewLayout = layout;
+  } else {
+    [self _configureCollectionViewLayout:layout];
+    self.view.collectionViewLayout = layout;
+  }
+}
+
+- (UICollectionViewLayout *)collectionViewLayout
+{
+  if ([self pendingState]) {
+    return _pendingState.collectionViewLayout;
+  } else {
+    return self.view.collectionViewLayout;
+  }
+}
+
+- (ASElementMap *)visibleElements
+{
+  ASDisplayNodeAssertMainThread();
+  // TODO Own the data controller when view is not yet loaded
+  return self.dataController.visibleMap;
+}
+
+- (id<ASCollectionLayoutDelegate>)layoutDelegate
+{
+  UICollectionViewLayout *layout = self.collectionViewLayout;
+  if ([layout isKindOfClass:[ASCollectionLayout class]]) {
+    return ((ASCollectionLayout *)layout).layoutDelegate;
+  }
+  return nil;
+}
+
 #pragma mark - Range Tuning
 
 - (ASRangeTuningParameters)tuningParametersForRangeType:(ASLayoutRangeType)rangeType
@@ -662,4 +711,14 @@ ASLayoutElementCollectionTableSetTraitCollection(_environmentStateLock)
   return result;
 }
 
+#pragma mark - Private methods
+
+- (void)_configureCollectionViewLayout:(UICollectionViewLayout *)layout
+{
+  if ([layout isKindOfClass:[ASCollectionLayout class]]) {
+    ASCollectionLayout *collectionLayout = (ASCollectionLayout *)layout;
+    collectionLayout.collectionNode = self;
+  }
+}
+
 @end

Source/ASCollectionView.mm

@@ -14,7 +14,9 @@
 #import <AsyncDisplayKit/ASCellNode+Internal.h>
 #import <AsyncDisplayKit/ASCollectionElement.h>
 #import <AsyncDisplayKit/ASCollectionInternal.h>
+#import <AsyncDisplayKit/ASCollectionLayout.h>
 #import <AsyncDisplayKit/ASCollectionViewLayoutController.h>
+#import <AsyncDisplayKit/ASCollectionViewLayoutFacilitatorProtocol.h>
 #import <AsyncDisplayKit/ASCollectionViewFlowLayoutInspector.h>
 #import <AsyncDisplayKit/ASDataController.h>
 #import <AsyncDisplayKit/ASDisplayNodeExtras.h>
@@ -23,16 +25,15 @@
 #import <AsyncDisplayKit/ASInternalHelpers.h>
 #import <AsyncDisplayKit/UICollectionViewLayout+ASConvenience.h>
 #import <AsyncDisplayKit/ASRangeController.h>
-#import <AsyncDisplayKit/ASCollectionNode.h>
 #import <AsyncDisplayKit/_ASCollectionViewCell.h>
 #import <AsyncDisplayKit/_ASDisplayLayer.h>
-#import <AsyncDisplayKit/ASCollectionViewLayoutFacilitatorProtocol.h>
 #import <AsyncDisplayKit/ASPagerNode.h>
 #import <AsyncDisplayKit/ASSectionContext.h>
 #import <AsyncDisplayKit/ASCollectionView+Undeprecated.h>
 #import <AsyncDisplayKit/_ASHierarchyChangeSet.h>
 #import <AsyncDisplayKit/CoreGraphics+ASConvenience.h>
 #import <AsyncDisplayKit/ASLayout.h>
+#import <AsyncDisplayKit/ASThread.h>
 
 /**
  * A macro to get self.collectionNode and assign it to a local variable, or return
@@ -213,6 +214,8 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
     unsigned int didChangeCollectionViewDataSource:1;
     unsigned int didChangeCollectionViewDelegate:1;
   } _layoutInspectorFlags;
+  
+  BOOL _hasDataControllerLayoutDelegate;
 }
 
 @end
@@ -292,6 +295,8 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
     _retainedLayer = self.layer;
   }
   
+  [self _configureCollectionViewLayout:layout];
+  
   return self;
 }
 
@@ -533,10 +538,13 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
   }
 }
 
-- (void)setCollectionViewLayout:(UICollectionViewLayout *)collectionViewLayout
+- (void)setCollectionViewLayout:(nonnull UICollectionViewLayout *)collectionViewLayout
 {
+  ASDisplayNodeAssertMainThread();
   [super setCollectionViewLayout:collectionViewLayout];
   
+  [self _configureCollectionViewLayout:collectionViewLayout];
+  
   // Trigger recreation of layout inspector with new collection view layout
   if (_layoutInspector != nil) {
     _layoutInspector = nil;
@@ -747,6 +755,14 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
 
 #pragma mark Internal
 
+- (void)_configureCollectionViewLayout:(nonnull UICollectionViewLayout *)layout
+{
+  _hasDataControllerLayoutDelegate = [layout conformsToProtocol:@protocol(ASDataControllerLayoutDelegate)];
+  if (_hasDataControllerLayoutDelegate) {
+    _dataController.layoutDelegate = (id<ASDataControllerLayoutDelegate>)layout;
+  }
+}
+
 /**
  Performing nested batch updates with super (e.g. resizing a cell node & updating collection view during same frame)
  can cause super to throw data integrity exceptions because it checks the data source counts before
@@ -1515,7 +1531,6 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
   }
 }
 
-
 #pragma mark - ASDataControllerSource
 
 - (ASCellNodeBlock)dataController:(ASDataController *)dataController nodeBlockAtIndexPath:(NSIndexPath *)indexPath
@@ -1578,11 +1593,6 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
   return block;
 }
 
-- (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath
-{
-  return [self.layoutInspector collectionView:self constrainedSizeForNodeAtIndexPath:indexPath];
-}
-
 - (NSUInteger)dataController:(ASDataController *)dataController rowsInSection:(NSUInteger)section
 {
   if (_asyncDataSourceFlags.collectionNodeNumberOfItemsInSection) {
@@ -1677,6 +1687,11 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
   }
 }
 
+- (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath
+{
+  return [self.layoutInspector collectionView:self constrainedSizeForNodeAtIndexPath:indexPath];
+}
+
 - (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForSupplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
 {
   if (_layoutInspectorFlags.constrainedSizeForSupplementaryNodeOfKindAtIndexPath) {
@@ -1892,6 +1907,7 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
 }
 
 #pragma mark - ASCellNodeDelegate
+
 - (void)nodeSelectedStateDidChange:(ASCellNode *)node
 {
   NSIndexPath *indexPath = [self indexPathForNode:node];
@@ -2030,6 +2046,10 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
  */
 - (void)layer:(CALayer *)layer didChangeBoundsWithOldValue:(CGRect)oldBounds newValue:(CGRect)newBounds
 {
+  if (_hasDataControllerLayoutDelegate) {
+    // Let the layout delegate handle bounds changes if it's available.
+    return;
+  }
   if (self.collectionViewLayout == nil) {
     return;
   }

Source/ASDisplayNode.mm

@@ -38,19 +38,6 @@
 #import <AsyncDisplayKit/ASWeakProxy.h>
 #import <AsyncDisplayKit/ASResponderChainEnumerator.h>
 
-/**
- * Assert if the current thread owns a mutex.
- * This assertion is useful when you want to indicate and enforce the locking policy/expectation of methods.
- * To determine when and which methods acquired a (recursive) mutex (to debug deadlocks, for example), 
- * put breakpoints at some of these assertions. When the breakpoints hit, walk through stack trace frames 
- * and check ownership count of the mutex.
- */
-#if CHECK_LOCKING_SAFETY
-  #define ASDisplayNodeAssertLockUnownedByCurrentThread(lock) ASDisplayNodeAssertFalse(lock.ownedByCurrentThread())
-#else
-  #define ASDisplayNodeAssertLockUnownedByCurrentThread(lock)
-#endif
-
 #if ASDisplayNodeLoggingEnabled
   #define LOG(...) NSLog(__VA_ARGS__)
 #else

Source/ASPagerFlowLayout.m

@@ -20,6 +20,7 @@
 
 @end
 
+//TODO make this an ASCollectionViewLayout
 @implementation ASPagerFlowLayout
 
 - (ASCollectionView *)asCollectionView

Source/ASTableView.mm

@@ -1555,7 +1555,7 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
   [changeSet executeCompletionHandlerWithFinished:YES];
 }
 
-#pragma mark - ASDataControllerDelegate
+#pragma mark - ASDataControllerSource
 
 - (ASCellNodeBlock)dataController:(ASDataController *)dataController nodeBlockAtIndexPath:(NSIndexPath *)indexPath {
   ASCellNodeBlock block = nil;

Source/AsyncDisplayKit.h

@@ -38,6 +38,11 @@
 #import <AsyncDisplayKit/ASCellNode.h>
 #import <AsyncDisplayKit/ASSectionContext.h>
 
+#import <AsyncDisplayKit/ASElementMap.h>
+#import <AsyncDisplayKit/ASCollectionLayoutContext.h>
+#import <AsyncDisplayKit/ASCollectionLayoutState.h>
+#import <AsyncDisplayKit/ASCollectionFlowLayoutDelegate.h>
+
 #import <AsyncDisplayKit/ASSectionController.h>
 #import <AsyncDisplayKit/ASSupplementaryNodeSource.h>
 

Source/Base/ASAssert.h

@@ -59,6 +59,19 @@
 #define ASDisplayNodeCAssertPositiveReal(description, num) ASDisplayNodeCAssert(num >= 0 && num <= CGFLOAT_MAX, @"%@ must be a real positive integer.", description)
 #define ASDisplayNodeCAssertInfOrPositiveReal(description, num) ASDisplayNodeCAssert(isinf(num) || (num >= 0 && num <= CGFLOAT_MAX), @"%@ must be infinite or a real positive integer.", description)
 
+/**
+ * Assert if the current thread owns a mutex.
+ * This assertion is useful when you want to indicate and enforce the locking policy/expectation of methods.
+ * To determine when and which methods acquired a (recursive) mutex (to debug deadlocks, for example),
+ * put breakpoints at some of these assertions. When the breakpoints hit, walk through stack trace frames
+ * and check ownership count of the mutex.
+ */
+#if CHECK_LOCKING_SAFETY
+#define ASDisplayNodeAssertLockUnownedByCurrentThread(lock) ASDisplayNodeAssertFalse(lock.ownedByCurrentThread())
+#else
+#define ASDisplayNodeAssertLockUnownedByCurrentThread(lock)
+#endif
+
 #define ASDisplayNodeErrorDomain @"ASDisplayNodeErrorDomain"
 #define ASDisplayNodeNonFatalErrorCode 1
 

Source/Details/ASCollectionFlowLayoutDelegate.h

@@ -0,0 +1,22 @@
+//
+//  ASCollectionFlowLayoutDelegate.h
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 28/2/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <AsyncDisplayKit/ASCollectionLayoutDelegate.h>
+#import <AsyncDisplayKit/ASScrollDirection.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+AS_SUBCLASSING_RESTRICTED
+
+@interface ASCollectionFlowLayoutDelegate : NSObject <ASCollectionLayoutDelegate>
+
+- (instancetype)initWithScrollableDirections:(ASScrollDirection)scrollableDirections;
+
+@end
+
+NS_ASSUME_NONNULL_END

Source/Details/ASCollectionFlowLayoutDelegate.m

@@ -0,0 +1,81 @@
+//  ASCollectionFlowLayoutDelegate.m
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 28/2/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <AsyncDisplayKit/ASCollectionFlowLayoutDelegate.h>
+
+#import <AsyncDisplayKit/ASCellNode.h>
+#import <AsyncDisplayKit/ASCollectionLayoutState.h>
+#import <AsyncDisplayKit/ASCollectionElement.h>
+#import <AsyncDisplayKit/ASCollectionLayoutContext.h>
+#import <AsyncDisplayKit/ASElementMap.h>
+#import <AsyncDisplayKit/ASLayout.h>
+#import <AsyncDisplayKit/ASStackLayoutSpec.h>
+
+@implementation ASCollectionFlowLayoutDelegate {
+  ASScrollDirection _scrollableDirections;
+}
+
+- (instancetype)init
+{
+  self = [super init];
+  if (self) {
+    _scrollableDirections = ASScrollDirectionVerticalDirections;
+  }
+  return self;
+}
+
+- (instancetype)initWithScrollableDirections:(ASScrollDirection)scrollableDirections
+{
+  self = [self init];
+  if (self) {
+    _scrollableDirections = scrollableDirections;
+  }
+  return self;
+}
+
+- (ASSizeRange)sizeRangeThatFits:(CGSize)viewportSize
+{
+  ASSizeRange sizeRange = ASSizeRangeUnconstrained;
+  if (ASScrollDirectionContainsVerticalDirection(_scrollableDirections) == NO) {
+    sizeRange.min.height = viewportSize.height;
+    sizeRange.max.height = viewportSize.height;
+  }
+  if (ASScrollDirectionContainsHorizontalDirection(_scrollableDirections) == NO) {
+    sizeRange.min.width = viewportSize.width;
+    sizeRange.max.width = viewportSize.width;
+  }
+  return sizeRange;
+}
+
+- (id)additionalInfoForLayoutWithElements:(ASElementMap *)elements
+{
+  return nil;
+}
+
+- (ASCollectionLayoutState *)calculateLayoutWithContext:(ASCollectionLayoutContext *)context
+{
+  ASElementMap *elements = context.elements;
+  NSMutableArray<ASCellNode *> *children = ASArrayByFlatMapping(elements.itemElements, ASCollectionElement *element, element.node);
+  if (children.count == 0) {
+    return [[ASCollectionLayoutState alloc] initWithElements:elements
+                                                 contentSize:CGSizeZero
+                                elementToLayoutArrtibutesMap:[NSMapTable weakToStrongObjectsMapTable]];
+  }
+  
+  ASStackLayoutSpec *stackSpec = [ASStackLayoutSpec stackLayoutSpecWithDirection:ASStackLayoutDirectionHorizontal
+                                                                         spacing:0
+                                                                  justifyContent:ASStackLayoutJustifyContentStart
+                                                                      alignItems:ASStackLayoutAlignItemsStart
+                                                                        flexWrap:ASStackLayoutFlexWrapWrap
+                                                                    alignContent:ASStackLayoutAlignContentStart
+                                                                        children:children];
+  stackSpec.concurrent = YES;
+  ASLayout *layout = [stackSpec layoutThatFits:[self sizeRangeThatFits:context.viewportSize]];
+  return [[ASCollectionLayoutState alloc] initWithElements:elements layout:layout];
+}
+
+@end

Source/Details/ASCollectionLayoutContext.h

@@ -0,0 +1,29 @@
+//
+//  ASCollectionLayoutContext.h
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 21/3/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+#import <AsyncDisplayKit/ASBaseDefines.h>
+
+@class ASElementMap;
+
+NS_ASSUME_NONNULL_BEGIN
+
+AS_SUBCLASSING_RESTRICTED
+
+@interface ASCollectionLayoutContext : NSObject
+
+@property (nonatomic, assign, readonly) CGSize viewportSize;
+@property (nonatomic, strong, readonly) ASElementMap *elements;
+@property (nonatomic, strong, readonly, nullable) id additionalInfo;
+
+- (instancetype)init __unavailable;
+
+@end
+
+NS_ASSUME_NONNULL_END

Source/Details/ASCollectionLayoutContext.mm

@@ -0,0 +1,59 @@
+//
+//  ASCollectionLayoutContext.mm
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 21/3/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <AsyncDisplayKit/ASCollectionLayoutContext.h>
+#import <AsyncDisplayKit/ASCollectionLayoutContext+Private.h>
+
+#import <AsyncDisplayKit/ASAssert.h>
+#import <AsyncDisplayKit/ASElementMap.h>
+#import <AsyncDisplayKit/ASEqualityHelpers.h>
+#import <AsyncDisplayKit/ASEqualityHashHelpers.h>
+
+@implementation ASCollectionLayoutContext
+
+- (instancetype)initWithViewportSize:(CGSize)viewportSize elements:(ASElementMap *)elements additionalInfo:(id)additionalInfo
+{
+  self = [super init];
+  if (self) {
+    _viewportSize = viewportSize;
+    _elements = elements;
+    _additionalInfo = additionalInfo;
+  }
+  return self;
+}
+
+- (BOOL)isEqualToContext:(ASCollectionLayoutContext *)context
+{
+  if (context == nil) {
+    return NO;
+  }
+  return CGSizeEqualToSize(_viewportSize, context.viewportSize) && ASObjectIsEqual(_elements, context.elements) && ASObjectIsEqual(_additionalInfo, context.additionalInfo);
+}
+
+- (BOOL)isEqual:(id)other
+{
+  if (self == other) {
+    return YES;
+  }
+  if (! [other isKindOfClass:[ASCollectionLayoutContext class]]) {
+    return NO;
+  }
+  return [self isEqualToContext:other];
+}
+
+- (NSUInteger)hash
+{
+  NSUInteger subhashes[] = {
+    ASHashFromCGSize(_viewportSize),
+    [_elements hash],
+    [_additionalInfo hash]
+  };
+  return ASIntegerArrayHash(subhashes, sizeof(subhashes) / sizeof(subhashes[0]));
+}
+
+@end

Source/Details/ASCollectionLayoutDelegate.h

@@ -0,0 +1,45 @@
+//
+//  ASCollectionLayoutDelegate.h
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 21/3/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+
+@class ASElementMap, ASCollectionLayoutContext, ASCollectionLayoutState;
+
+NS_ASSUME_NONNULL_BEGIN
+
+@protocol ASCollectionLayoutDelegate <NSObject>
+
+/**
+ * @abstract Returns any additional information needed for a coming layout pass with the given elements.
+ *
+ * @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.
+ */
+- (nullable id)additionalInfoForLayoutWithElements:(ASElementMap *)elements;
+
+/**
+ * @abstract Prepares and returns a new layout for given context.
+ *
+ * @param context A context that contains all elements to be laid out and any additional information needed.
+ *
+ * @return The new layout calculated for the given context.
+ *
+ * @discussion This method is called ahead of time, i.e before the underlying collection/table view is aware of the provided elements.
+ * As a result, this method should rely solely on the given context and should not reach out to other objects for information not available in the context.
+ *
+ * @discussion This method will be called on background theads. It must be thread-safe and should not change any internal state of this object.
+ *
+ * @discussion This method must block its calling thread. It can dispatch to other theads to reduce blocking time.
+ */
+- (ASCollectionLayoutState *)calculateLayoutWithContext:(ASCollectionLayoutContext *)context;
+
+@end
+
+NS_ASSUME_NONNULL_END

Source/Details/ASCollectionLayoutState.h

@@ -0,0 +1,55 @@
+//
+//  ASCollectionLayoutState.h
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 9/3/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+#import <AsyncDisplayKit/ASBaseDefines.h>
+
+@class ASElementMap, ASCollectionElement, ASLayout;
+
+NS_ASSUME_NONNULL_BEGIN
+
+AS_SUBCLASSING_RESTRICTED
+@interface ASCollectionLayoutState : NSObject
+
+/// The elements used to calculate this object
+@property (nonatomic, strong, readonly) ASElementMap *elements;
+
+@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 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. 
+ * Also, it should have NSMapTableObjectPointerPersonality and NSMapTableWeakMemory as key options.
+ */
+- (instancetype)initWithElements:(ASElementMap *)elements contentSize:(CGSize)contentSize elementToLayoutArrtibutesMap:(NSMapTable<ASCollectionElement *, UICollectionViewLayoutAttributes *> *)attrsMap NS_DESIGNATED_INITIALIZER;
+
+/**
+ * Convenience initializer.
+ *
+ * @param elements The elements 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.
+ *
+ * @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;
+
+@end
+
+NS_ASSUME_NONNULL_END

Source/Details/ASCollectionLayoutState.m

@@ -0,0 +1,52 @@
+//
+//  ASCollectionLayoutState.m
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 9/3/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <AsyncDisplayKit/ASCollectionLayoutState.h>
+
+#import <AsyncDisplayKit/ASAssert.h>
+#import <AsyncDisplayKit/ASCellNode+Internal.h>
+#import <AsyncDisplayKit/ASCollectionElement.h>
+#import <AsyncDisplayKit/ASElementMap.h>
+#import <AsyncDisplayKit/ASLayout.h>
+
+@implementation ASCollectionLayoutState
+
+- (instancetype)initWithElements:(ASElementMap *)elements layout:(ASLayout *)layout
+{
+  NSMapTable<ASCollectionElement *, UICollectionViewLayoutAttributes *> *attrsMap = [NSMapTable mapTableWithKeyOptions:(NSMapTableObjectPointerPersonality | NSMapTableWeakMemory) valueOptions:NSMapTableStrongMemory];
+  for (ASLayout *sublayout in layout.sublayouts) {
+    ASCollectionElement *element = ((ASCellNode *)sublayout.layoutElement).collectionElement;
+    NSIndexPath *indexPath = [elements indexPathForElement:element];
+    NSString *supplementaryElementKind = element.supplementaryElementKind;
+    
+    UICollectionViewLayoutAttributes *attrs;
+    if (supplementaryElementKind == nil) {
+      attrs = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];
+    } else {
+      attrs = [UICollectionViewLayoutAttributes layoutAttributesForSupplementaryViewOfKind:supplementaryElementKind withIndexPath:indexPath];
+    }
+    
+    attrs.frame = sublayout.frame;
+    [attrsMap setObject:attrs forKey:element];
+  }
+
+  return [self initWithElements:elements contentSize:layout.size elementToLayoutArrtibutesMap:attrsMap];
+}
+
+- (instancetype)initWithElements:(ASElementMap *)elements contentSize:(CGSize)contentSize elementToLayoutArrtibutesMap:(NSMapTable<ASCollectionElement *,UICollectionViewLayoutAttributes *> *)attrsMap
+{
+  self = [super init];
+  if (self) {
+    _elements = elements;
+    _contentSize = contentSize;
+    _elementToLayoutArrtibutesMap = attrsMap;
+  }
+  return self;
+}
+
+@end

Source/Details/ASDataController.h

@@ -27,9 +27,10 @@ NS_ASSUME_NONNULL_BEGIN
 #endif
 
 @class ASCellNode;
+@class ASCollectionElement;
 @class ASDataController;
 @class ASElementMap;
-@class ASCollectionElement;
+@class ASLayout;
 @class _ASHierarchyChangeSet;
 @protocol ASTraitEnvironment;
 @protocol ASSectionContext;
@@ -51,11 +52,6 @@ extern NSString * const ASCollectionInvalidUpdateException;
  */
 - (ASCellNodeBlock)dataController:(ASDataController *)dataController nodeBlockAtIndexPath:(NSIndexPath *)indexPath;
 
-/**
- The constrained size range for layout.
- */
-- (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath;
-
 /**
  Fetch the number of rows in specific section.
  */
@@ -73,12 +69,20 @@ extern NSString * const ASCollectionInvalidUpdateException;
 
 @optional
 
+/**
+ The constrained size range for layout. Called only if collection layout delegate is not provided.
+ */
+- (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath;
+
 - (NSArray<NSString *> *)dataController:(ASDataController *)dataController supplementaryNodeKindsInSections:(NSIndexSet *)sections;
 
 - (NSUInteger)dataController:(ASDataController *)dataController supplementaryNodesOfKind:(NSString *)kind inSection:(NSUInteger)section;
 
 - (ASCellNodeBlock)dataController:(ASDataController *)dataController supplementaryNodeBlockOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath;
 
+/**
+ The constrained size range for layout. Called only if no data controller layout delegate is provided.
+ */
 - (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForSupplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath;
 
 - (nullable id<ASSectionContext>)dataController:(ASDataController *)dataController contextForSection:(NSInteger)section;
@@ -113,6 +117,33 @@ extern NSString * const ASCollectionInvalidUpdateException;
 
 @end
 
+@protocol ASDataControllerLayoutDelegate <NSObject>
+
+/**
+ * @abstract Returns a layout context needed for a coming layout pass with the given elements.
+ * The context should contain the elements and any additional information needed.
+ *
+ * @discussion This method will be called on main thread.
+ */
+- (id)layoutContextWithElements:(ASElementMap *)elements;
+
+/**
+ * @abstract Prepares in advance a new layout with the given context.
+ *
+ * @param context A context that was previously returned by `-layoutContextWithElements:`.
+ *
+ * @discussion This method is called ahead of time, i.e before the underlying collection/table view is aware of the provided elements.
+ * As a result, this method should rely solely on the given context and should not reach out to its collection/table view for information regarding items.
+ *
+ * @discussion This method will be called on background theads. It must be thread-safe and should not change any internal state of the conforming object.
+ * It's recommended to put the resulting layouts of this method into a thread-safe cache that can be looked up later on.
+ *
+ * @discussion This method must block its calling thread. It can dispatch to other theads to reduce blocking time.
+ */
+- (void)prepareLayoutWithContext:(id)context;
+
+@end
+
 /**
  * Controller to layout data in background, and managed data updating.
  *
@@ -154,6 +185,11 @@ extern NSString * const ASCollectionInvalidUpdateException;
  */
 @property (nonatomic, weak) id<ASDataControllerEnvironmentDelegate> environmentDelegate;
 
+/**
+ * Delegate for preparing layouts. Main thead only.
+ */
+@property (nonatomic, weak) id<ASDataControllerLayoutDelegate> layoutDelegate;
+
 #ifdef __cplusplus
 /**
  * Returns the most recently gathered item counts from the data source. If the counts
@@ -193,7 +229,7 @@ extern NSString * const ASCollectionInvalidUpdateException;
 - (void)relayoutAllNodes;
 
 /**
- * Re-measures given noades in the backing store.
+ * Re-measures given nodes in the backing store.
  *
  * @discussion Used to respond to setNeedsLayout calls in ASCellNode
  */

Source/Details/ASDataController.mm

@@ -13,15 +13,17 @@
 #import <AsyncDisplayKit/_ASHierarchyChangeSet.h>
 #import <AsyncDisplayKit/ASAssert.h>
 #import <AsyncDisplayKit/ASCellNode.h>
-#import <AsyncDisplayKit/ASLayout.h>
-#import <AsyncDisplayKit/ASMainSerialQueue.h>
-#import <AsyncDisplayKit/ASTwoDimensionalArrayUtils.h>
-#import <AsyncDisplayKit/ASSection.h>
-#import <AsyncDisplayKit/ASThread.h>
 #import <AsyncDisplayKit/ASCollectionElement.h>
+#import <AsyncDisplayKit/ASCollectionLayoutContext.h>
 #import <AsyncDisplayKit/ASDispatch.h>
 #import <AsyncDisplayKit/ASElementMap.h>
+#import <AsyncDisplayKit/ASLayout.h>
+#import <AsyncDisplayKit/ASMainSerialQueue.h>
 #import <AsyncDisplayKit/ASMutableElementMap.h>
+#import <AsyncDisplayKit/ASThread.h>
+#import <AsyncDisplayKit/ASTwoDimensionalArrayUtils.h>
+#import <AsyncDisplayKit/ASSection.h>
+
 #import <AsyncDisplayKit/ASInternalHelpers.h>
 #import <AsyncDisplayKit/ASCellNode+Internal.h>
 #import <AsyncDisplayKit/ASDisplayNode+Subclasses.h>
@@ -52,6 +54,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
 #endif
 
 @interface ASDataController () {
+  id<ASDataControllerLayoutDelegate> _layoutDelegate;
 
   NSInteger _nextSectionID;
   
@@ -69,6 +72,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
     unsigned int supplementaryNodeKindsInSections:1;
     unsigned int supplementaryNodesOfKindInSection:1;
     unsigned int supplementaryNodeBlockOfKindAtIndexPath:1;
+    unsigned int constrainedSizeForNodeAtIndexPath:1;
     unsigned int constrainedSizeForSupplementaryNodeOfKindAtIndexPath:1;
     unsigned int contextForSection:1;
   } _dataSourceFlags;
@@ -91,6 +95,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   _dataSourceFlags.supplementaryNodeKindsInSections = [_dataSource respondsToSelector:@selector(dataController:supplementaryNodeKindsInSections:)];
   _dataSourceFlags.supplementaryNodesOfKindInSection = [_dataSource respondsToSelector:@selector(dataController:supplementaryNodesOfKind:inSection:)];
   _dataSourceFlags.supplementaryNodeBlockOfKindAtIndexPath = [_dataSource respondsToSelector:@selector(dataController:supplementaryNodeBlockOfKind:atIndexPath:)];
+  _dataSourceFlags.constrainedSizeForNodeAtIndexPath = [_dataSource respondsToSelector:@selector(dataController:constrainedSizeForNodeAtIndexPath:)];
   _dataSourceFlags.constrainedSizeForSupplementaryNodeOfKindAtIndexPath = [_dataSource respondsToSelector:@selector(dataController:constrainedSizeForSupplementaryNodeOfKind:atIndexPath:)];
   _dataSourceFlags.contextForSection = [_dataSource respondsToSelector:@selector(dataController:contextForSection:)];
   
@@ -132,9 +137,23 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   return parallelProcessorCount;
 }
 
+- (id<ASDataControllerLayoutDelegate>)layoutDelegate
+{
+  ASDisplayNodeAssertMainThread();
+  return _layoutDelegate;
+}
+
+- (void)setLayoutDelegate:(id<ASDataControllerLayoutDelegate>)layoutDelegate
+{
+  ASDisplayNodeAssertMainThread();
+  if (layoutDelegate != _layoutDelegate) {
+    _layoutDelegate = layoutDelegate;
+  }
+}
+
 #pragma mark - Cell Layout
 
-- (void)batchLayoutNodesFromContexts:(NSArray<ASCollectionElement *> *)elements batchSize:(NSInteger)batchSize batchCompletion:(ASDataControllerCompletionBlock)batchCompletionHandler
+- (void)batchAllocateNodesFromElements:(NSArray<ASCollectionElement *> *)elements andLayout:(BOOL)shouldLayout batchSize:(NSInteger)batchSize batchCompletion:(ASDataControllerCompletionBlock)batchCompletionHandler
 {
   ASSERT_ON_EDITING_QUEUE;
 #if AS_MEASURE_AVOIDED_DATACONTROLLER_WORK
@@ -156,9 +175,9 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   // Processing in batches
   for (NSUInteger i = 0; i < count; i += batchSize) {
     NSRange batchedRange = NSMakeRange(i, MIN(count - i, batchSize));
-    NSArray<ASCollectionElement *> *batchedContexts = [elements subarrayWithRange:batchedRange];
+    NSArray<ASCollectionElement *> *batchedElements = [elements subarrayWithRange:batchedRange];
-    NSArray<ASCellNode *> *nodes = [self _layoutNodesFromContexts:batchedContexts];
+    NSArray<ASCellNode *> *nodes = [self _allocateNodesFromElements:batchedElements andLayout:shouldLayout];
-    batchCompletionHandler(batchedContexts, nodes);
+    batchCompletionHandler(batchedElements, nodes);
   }
   
   ASProfilingSignpostEnd(2, _dataSource);
@@ -176,7 +195,8 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   node.frame = frame;
 }
 
-- (NSArray<ASCellNode *> *)_layoutNodesFromContexts:(NSArray<ASCollectionElement *> *)elements
+// TODO Is returned array still needed? Can it be removed?
+- (NSArray<ASCellNode *> *)_allocateNodesFromElements:(NSArray<ASCollectionElement *> *)elements andLayout:(BOOL)shouldLayout
 {
   ASSERT_ON_EDITING_QUEUE;
   
@@ -199,6 +219,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
       node = [[ASCellNode alloc] init]; // Fallback to avoid crash for production apps.
     }
     
+    if (shouldLayout) {
       // Layout the node if the size range is valid.
       ASSizeRange sizeRange = context.constrainedSize;
       if (ASSizeRangeHasSignificantArea(sizeRange)) {
@@ -208,6 +229,8 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
 #if AS_MEASURE_AVOIDED_DATACONTROLLER_WORK
       [ASDataController _didLayoutNode];
 #endif
+    }
+
     allocatedNodeBuffer[i] = node;
   });
 
@@ -269,6 +292,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
  * @param owningNode The node that owns the new elements.
  * @param traitCollection The trait collection needed to initialize elements
  * @param indexPathsAreNew YES if index paths are "after the update," NO otherwise.
+ * @param shouldFetchSizeRanges Whether constrained sizes should be fetched from data source
  */
 - (void)_repopulateSupplementaryNodesIntoMap:(ASMutableElementMap *)map
              forSectionsContainingIndexPaths:(NSArray<NSIndexPath *> *)indexPaths
@@ -276,6 +300,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                                   owningNode:(ASDisplayNode *)owningNode
                              traitCollection:(ASPrimitiveTraitCollection)traitCollection
                             indexPathsAreNew:(BOOL)indexPathsAreNew
+                       shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
 {
   ASDisplayNodeAssertMainThread();
 
@@ -298,7 +323,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
 
   for (NSString *kind in [self supplementaryKindsInSections:newSections]) {
-    [self _insertElementsIntoMap:map kind:kind forSections:newSections owningNode:owningNode traitCollection:traitCollection];
+    [self _insertElementsIntoMap:map kind:kind forSections:newSections owningNode:owningNode traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
   }
 }
 
@@ -309,12 +334,14 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
  * @param sections The sections that should be populated by new elements
  * @param owningNode The node that owns the new elements.
  * @param traitCollection The trait collection needed to initialize elements
+ * @param shouldFetchSizeRanges Whether constrained sizes should be fetched from data source
  */
 - (void)_insertElementsIntoMap:(ASMutableElementMap *)map
                           kind:(NSString *)kind
                    forSections:(NSIndexSet *)sections
                     owningNode:(ASDisplayNode *)owningNode
                traitCollection:(ASPrimitiveTraitCollection)traitCollection
+         shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
 {
   ASDisplayNodeAssertMainThread();
   
@@ -323,7 +350,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
   
   NSArray<NSIndexPath *> *indexPaths = [self _allIndexPathsForItemsOfKind:kind inSections:sections];
-  [self _insertElementsIntoMap:map kind:kind atIndexPaths:indexPaths owningNode:owningNode traitCollection:traitCollection];
+  [self _insertElementsIntoMap:map kind:kind atIndexPaths:indexPaths owningNode:owningNode traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
 }
 
 /**
@@ -334,12 +361,14 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
  * @param indexPaths The index paths at which new elements should be populated
  * @param owningNode The node that owns the new elements.
  * @param traitCollection The trait collection needed to initialize elements
+ * @param shouldFetchSizeRanges Whether constrained sizes should be fetched from data source
  */
 - (void)_insertElementsIntoMap:(ASMutableElementMap *)map
                           kind:(NSString *)kind
                   atIndexPaths:(NSArray<NSIndexPath *> *)indexPaths
                     owningNode:(ASDisplayNode *)owningNode
                traitCollection:(ASPrimitiveTraitCollection)traitCollection
+         shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
 {
   ASDisplayNodeAssertMainThread();
   
@@ -362,7 +391,11 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
       nodeBlock = [_dataSource dataController:self supplementaryNodeBlockOfKind:kind atIndexPath:indexPath];
     }
     
-    ASSizeRange constrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
+    ASSizeRange constrainedSize;
+    if (shouldFetchSizeRanges) {
+      constrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
+    }
+    
     ASCollectionElement *element = [[ASCollectionElement alloc] initWithNodeBlock:nodeBlock
                                                          supplementaryElementKind:isRowKind ? nil : kind
                                                                   constrainedSize:constrainedSize
@@ -404,18 +437,34 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   return @[];
 }
 
+- (ASSizeRange)constrainedSizeForElement:(ASCollectionElement *)element inElementMap:(ASElementMap *)map
+{
+  ASDisplayNodeAssertMainThread();
+  NSString *kind = element.supplementaryElementKind ?: ASDataControllerRowNodeKind;
+  NSIndexPath *indexPath = [map indexPathForElement:element];
+  return [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
+}
+
+
 - (ASSizeRange)constrainedSizeForNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
 {
   ASDisplayNodeAssertMainThread();
+  
+  id<ASDataControllerSource> dataSource = _dataSource;
+  if (dataSource == nil) {
+    return ASSizeRangeZero;
+  }
+  
   if ([kind isEqualToString:ASDataControllerRowNodeKind]) {
-    return [_dataSource dataController:self constrainedSizeForNodeAtIndexPath:indexPath];
+    ASDisplayNodeAssert(_dataSourceFlags.constrainedSizeForNodeAtIndexPath, @"-dataController:constrainedSizeForNodeAtIndexPath: must also be implemented");
+    return [dataSource dataController:self constrainedSizeForNodeAtIndexPath:indexPath];
   }
   
   if (_dataSourceFlags.constrainedSizeForSupplementaryNodeOfKindAtIndexPath){
-    return [_dataSource dataController:self constrainedSizeForSupplementaryNodeOfKind:kind atIndexPath:indexPath];
+    return [dataSource dataController:self constrainedSizeForSupplementaryNodeOfKind:kind atIndexPath:indexPath];
   }
   
-  ASDisplayNodeAssert(NO, @"Unknown constrained size for node of kind %@ by data source %@", kind, _dataSource);
+  ASDisplayNodeAssert(NO, @"Unknown constrained size for node of kind %@ by data source %@", kind, dataSource);
   return ASSizeRangeZero;
 }
 
@@ -484,27 +533,51 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   // Mutable copy of current data.
   ASMutableElementMap *mutableMap = [_pendingMap mutableCopy];
   
-  // Step 1: update the mutable copies to match the data source's state
+  BOOL canDelegateLayout = (_layoutDelegate != nil);
+
+  // Step 1: Update the mutable copies to match the data source's state
   [self _updateSectionContextsInMap:mutableMap changeSet:changeSet];
   __weak id<ASTraitEnvironment> environment = [self.environmentDelegate dataControllerEnvironment];
   __weak ASDisplayNode *owningNode = (ASDisplayNode *)environment; // This is gross!
   ASPrimitiveTraitCollection existingTraitCollection = [environment primitiveTraitCollection];
-  [self _updateElementsInMap:mutableMap changeSet:changeSet owningNode:owningNode traitCollection:existingTraitCollection];
+  [self _updateElementsInMap:mutableMap changeSet:changeSet owningNode:owningNode traitCollection:existingTraitCollection shouldFetchSizeRanges:(! canDelegateLayout)];
   
   // Step 2: Clone the new data
   ASElementMap *newMap = [mutableMap copy];
-
   _pendingMap = newMap;
   
+  // Step 3: Ask layout delegate for contexts
+  id layoutContext = nil;
+  if (canDelegateLayout) {
+    layoutContext = [_layoutDelegate layoutContextWithElements:newMap];
+  }
+  
   dispatch_group_async(_editingTransactionGroup, _editingTransactionQueue, ^{
-    // Step 3: Layout **all** new elements without batching in background.
+    // Step 4: Allocate and layout elements if can't delegate
-    NSArray<ASCollectionElement *> *unmeasuredElements = [ASDataController unmeasuredElementsFromMap:newMap];
+    NSArray<ASCollectionElement *> *elementsToProcess;
-    [self batchLayoutNodesFromContexts:unmeasuredElements batchSize:unmeasuredElements.count batchCompletion:^(id, id) {
+    if (canDelegateLayout) {
+      // Allocate all nodes before handling them to the layout delegate.
+      // In the future, we may want to let the delegate drive allocation as well.
+      elementsToProcess = ASArrayByFlatMapping(newMap,
+                                               ASCollectionElement *element,
+                                               (element.nodeIfAllocated == nil ? element : nil));
+    } else {
+      elementsToProcess = ASArrayByFlatMapping(newMap,
+                                               ASCollectionElement *element,
+                                               (element.nodeIfAllocated.calculatedLayout == nil ? element : nil));
+    }
+    
+    [self batchAllocateNodesFromElements:elementsToProcess andLayout:(! canDelegateLayout) batchSize:elementsToProcess.count batchCompletion:^(NSArray<ASCollectionElement *> *elements, NSArray<ASCellNode *> *nodes) {
       ASSERT_ON_EDITING_QUEUE;
+
+      if (canDelegateLayout) {
+        [_layoutDelegate prepareLayoutWithContext:layoutContext];
+      }
+      
       [_mainSerialQueue performBlockOnMainThread:^{
         [_delegate dataController:self willUpdateWithChangeSet:changeSet];
 
-        // Step 4: Deploy the new data as "completed" and inform delegate
+        // Step 5: Deploy the new data as "completed" and inform delegate
         _visibleMap = newMap;
         
         [_delegate dataController:self didUpdateWithChangeSet:changeSet];
@@ -566,6 +639,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                    changeSet:(_ASHierarchyChangeSet *)changeSet
                   owningNode:(ASDisplayNode *)owningNode
              traitCollection:(ASPrimitiveTraitCollection)traitCollection
+       shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
 {
   ASDisplayNodeAssertMainThread();
 
@@ -575,7 +649,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
     NSUInteger sectionCount = [self itemCountsFromDataSource].size();
     if (sectionCount > 0) {
       NSIndexSet *sectionIndexes = [NSIndexSet indexSetWithIndexesInRange:NSMakeRange(0, sectionCount)];
-      [self _insertElementsIntoMap:map sections:sectionIndexes owningNode:owningNode traitCollection:traitCollection];
+      [self _insertElementsIntoMap:map sections:sectionIndexes owningNode:owningNode traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
     }
     // Return immediately because reloadData can't be used in conjuntion with other updates.
     return;
@@ -588,7 +662,8 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                                      changeSet:changeSet
                                     owningNode:owningNode
                                traitCollection:traitCollection
-                              indexPathsAreNew:NO];
+                              indexPathsAreNew:NO
+                         shouldFetchSizeRanges:shouldFetchSizeRanges];
   }
 
   for (_ASHierarchySectionChange *change in [changeSet sectionChangesOfType:_ASHierarchyChangeTypeDelete]) {
@@ -598,17 +673,18 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
   
   for (_ASHierarchySectionChange *change in [changeSet sectionChangesOfType:_ASHierarchyChangeTypeInsert]) {
-    [self _insertElementsIntoMap:map sections:change.indexSet owningNode:owningNode traitCollection:traitCollection];
+    [self _insertElementsIntoMap:map sections:change.indexSet owningNode:owningNode traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
   }
   
   for (_ASHierarchyItemChange *change in [changeSet itemChangesOfType:_ASHierarchyChangeTypeInsert]) {
-    [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind atIndexPaths:change.indexPaths owningNode:owningNode traitCollection:traitCollection];
+    [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind atIndexPaths:change.indexPaths owningNode:owningNode traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
     // Aggressively reload supplementary nodes (#1773 & #1629)
     [self _repopulateSupplementaryNodesIntoMap:map forSectionsContainingIndexPaths:change.indexPaths
                                      changeSet:changeSet
                                     owningNode:owningNode
                                traitCollection:traitCollection
-                              indexPathsAreNew:YES];
+                              indexPathsAreNew:YES
+                         shouldFetchSizeRanges:shouldFetchSizeRanges];
   }
 }
 
@@ -616,6 +692,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                       sections:(NSIndexSet *)sectionIndexes
                     owningNode:(ASDisplayNode *)owningNode
                traitCollection:(ASPrimitiveTraitCollection)traitCollection
+         shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
 {
   ASDisplayNodeAssertMainThread();
   
@@ -625,12 +702,12 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
 
   // Items
   [map insertEmptySectionsOfItemsAtIndexes:sectionIndexes];
-  [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind forSections:sectionIndexes owningNode:owningNode traitCollection:traitCollection];
+  [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind forSections:sectionIndexes owningNode:owningNode traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
 
   // Supplementaries
   for (NSString *kind in [self supplementaryKindsInSections:sectionIndexes]) {
     // Step 2: Populate new elements for all sections
-    [self _insertElementsIntoMap:map kind:kind forSections:sectionIndexes owningNode:owningNode traitCollection:traitCollection];
+    [self _insertElementsIntoMap:map kind:kind forSections:sectionIndexes owningNode:owningNode traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
   }
 }
 
@@ -646,9 +723,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
   
   for (ASCellNode *node in nodes) {
-    NSString *kind = node.collectionElement.supplementaryElementKind ?: ASDataControllerRowNodeKind;
+    ASSizeRange constrainedSize = [self constrainedSizeForElement:node.collectionElement inElementMap:_pendingMap];
-    NSIndexPath *indexPath = [_pendingMap indexPathForElement:node.collectionElement];
-    ASSizeRange constrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
     [self _layoutNode:node withConstrainedSize:constrainedSize];
     BOOL matchesSize = [_dataSource dataController:self presentedSizeForElement:node.collectionElement matchesSize:node.frame.size];
     if (! matchesSize) {
@@ -676,9 +751,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
 {
   ASDisplayNodeAssertMainThread();
   for (ASCollectionElement *element in _visibleMap) {
-    NSIndexPath *indexPath = [_visibleMap indexPathForElement:element];
+    ASSizeRange constrainedSize = [self constrainedSizeForElement:element inElementMap:_visibleMap];
-    NSString *kind = element.supplementaryElementKind ?: ASDataControllerRowNodeKind;
-    ASSizeRange constrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
     if (ASSizeRangeHasSignificantArea(constrainedSize)) {
       element.constrainedSize = constrainedSize;
 
@@ -721,17 +794,6 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   [_mainSerialQueue performBlockOnMainThread:block];
 }
 
-+ (NSArray<ASCollectionElement *> *)unmeasuredElementsFromMap:(ASElementMap *)map
-{
-  NSMutableArray<ASCollectionElement *> *unloadedContexts = [NSMutableArray array];
-  for (ASCollectionElement *element in map) {
-    if (element.nodeIfAllocated.calculatedLayout == nil) {
-      [unloadedContexts addObject:element];
-    }
-  }
-  return unloadedContexts;
-}
-
 @end
 
 #if AS_MEASURE_AVOIDED_DATACONTROLLER_WORK

Source/Details/ASElementMap.h

@@ -49,6 +49,11 @@ AS_SUBCLASSING_RESTRICTED
  */
 @property (copy, readonly) NSArray<NSIndexPath *> *itemIndexPaths;
 
+/**
+ * All the item elements in this map, in ascending order. O(N)
+ */
+@property (copy, readonly) NSArray<ASCollectionElement *> *itemElements;
+
 /**
  * Returns the index path that corresponds to the same element in @c map at the given @c indexPath. O(1)
  */

Source/Details/ASElementMap.m

@@ -69,6 +69,11 @@
   return ASIndexPathsForTwoDimensionalArray(_sectionsOfItems);
 }
 
+- (NSArray<ASCollectionElement *> *)itemElements
+{
+  return ASElementsInTwoDimensionalArray(_sectionsOfItems);
+}
+
 - (NSInteger)numberOfSections
 {
   return _sectionsOfItems.count;

Source/Layout/ASStackLayoutDefines.h

@@ -86,13 +86,13 @@ typedef NS_ENUM(NSUInteger, ASStackLayoutAlignSelf) {
   ASStackLayoutAlignSelfStretch,
 };
 
-// TODO documentation
+/** Whether children are stacked into a single or multiple lines. */
 typedef NS_ENUM(NSUInteger, ASStackLayoutFlexWrap) {
   ASStackLayoutFlexWrapNoWrap,
   ASStackLayoutFlexWrapWrap,
 };
 
-// TODO documentation
+/** Orientation of lines along cross axis if there are multiple lines. */
 typedef NS_ENUM(NSUInteger, ASStackLayoutAlignContent) {
   ASStackLayoutAlignContentStart,
   ASStackLayoutAlignContentCenter,

Source/Layout/ASStackLayoutSpec.h

@@ -63,7 +63,6 @@ NS_ASSUME_NONNULL_BEGIN
 @property (nonatomic, assign) ASStackLayoutFlexWrap flexWrap;
 /** Orientation of lines along cross axis if there are multiple lines. Defaults to ASStackLayoutAlignContentStart */
 @property (nonatomic, assign) ASStackLayoutAlignContent alignContent;
-
 /** Whether this stack can dispatch to other threads, regardless of which thread it's running on */
 @property (nonatomic, assign, getter=isConcurrent) BOOL concurrent;
 

Source/Private/ASCollectionLayout.h

@@ -0,0 +1,51 @@
+//
+//  ASCollectionLayout.h
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 28/2/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+#import <AsyncDisplayKit/ASBaseDefines.h>
+
+@protocol ASCollectionLayoutDelegate;
+@class ASElementMap, ASCollectionLayout, ASCollectionNode;
+
+NS_ASSUME_NONNULL_BEGIN
+
+AS_SUBCLASSING_RESTRICTED
+
+@interface ASCollectionLayout : UICollectionViewLayout
+
+/**
+ * The collection node object currently using this layout object.
+ *
+ * @discussion The collection node object sets the value of this property when a new layout object is assigned to it.
+ *
+ * @discussion To get the truth on the current state of the collection, call methods on the collection node or the data source rather than the collection view because:
+ *  1. The view might not yet be allocated.
+ *  2. The collection node and data source are thread-safe.
+ */
+@property (nonatomic, weak) ASCollectionNode *collectionNode;
+
+@property (nonatomic, strong, readonly) id<ASCollectionLayoutDelegate> layoutDelegate;
+
+/**
+ * Initializes with a layout delegate.
+ *
+ * @discussion For developers' convenience, the delegate is retained by this layout object, similar to UICollectionView retains its UICollectionViewLayout object.
+ *
+ * @discussion For simplicity, the delegate is read-only. If a new layout delegate is needed, construct a new layout object with that delegate and notify ASCollectionView about it.
+ * This ensures the underlying UICollectionView purges its cache and properly loads the new layout.
+ */
+- (instancetype)initWithLayoutDelegate:(id<ASCollectionLayoutDelegate>)layoutDelegate NS_DESIGNATED_INITIALIZER;
+
+- (instancetype)init __unavailable;
+
+- (instancetype)initWithCoder:(NSCoder *)aDecoder __unavailable;
+
+@end
+
+NS_ASSUME_NONNULL_END

Source/Private/ASCollectionLayout.mm

@@ -0,0 +1,153 @@
+//
+//  ASCollectionLayout.mm
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 28/2/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <AsyncDisplayKit/ASCollectionLayout.h>
+
+#import <AsyncDisplayKit/ASAssert.h>
+#import <AsyncDisplayKit/ASCollectionElement.h>
+#import <AsyncDisplayKit/ASCollectionLayoutContext+Private.h>
+#import <AsyncDisplayKit/ASCollectionLayoutDelegate.h>
+#import <AsyncDisplayKit/ASCollectionLayoutState.h>
+#import <AsyncDisplayKit/ASCollectionNode+Beta.h>
+#import <AsyncDisplayKit/ASDisplayNode+FrameworkPrivate.h>
+#import <AsyncDisplayKit/ASElementMap.h>
+#import <AsyncDisplayKit/ASEqualityHelpers.h>
+#import <AsyncDisplayKit/ASThread.h>
+
+@interface ASCollectionLayout () <ASDataControllerLayoutDelegate> {
+  ASDN::Mutex __instanceLock__; // Non-recursive mutex, ftw!
+  
+  // Main thread only.
+  ASCollectionLayoutState *_state;
+  
+  // The pending state calculated ahead of time, if any.
+  ASCollectionLayoutState *_pendingState;
+  // The context used to calculate _pendingState
+  ASCollectionLayoutContext *_layoutContextForPendingState;
+  
+  BOOL _layoutDelegateImplementsAdditionalInfoForLayoutWithElements;
+}
+
+@end
+
+@implementation ASCollectionLayout
+
+- (instancetype)initWithLayoutDelegate:(id<ASCollectionLayoutDelegate>)layoutDelegate
+{
+  self = [super init];
+  if (self) {
+    ASDisplayNodeAssertNotNil(layoutDelegate, @"Collection layout delegate cannot be nil");
+    _layoutDelegate = layoutDelegate;
+    _layoutDelegateImplementsAdditionalInfoForLayoutWithElements = [layoutDelegate respondsToSelector:@selector(additionalInfoForLayoutWithElements:)];
+  }
+  return self;
+}
+
+#pragma mark - ASDataControllerLayoutDelegate
+
+- (id)layoutContextWithElements:(ASElementMap *)elements
+{
+  ASDisplayNodeAssertMainThread();
+  id additionalInfo = nil;
+  if (_layoutDelegateImplementsAdditionalInfoForLayoutWithElements) {
+    additionalInfo = [_layoutDelegate additionalInfoForLayoutWithElements:elements];
+  }
+  return [[ASCollectionLayoutContext alloc] initWithViewportSize:[self viewportSize] elements:elements additionalInfo:additionalInfo];
+}
+
+- (void)prepareLayoutWithContext:(id)context
+{
+  ASCollectionLayoutState *state = [_layoutDelegate calculateLayoutWithContext:context];
+  
+  ASDN::MutexLocker l(__instanceLock__);
+  _pendingState = state;
+  _layoutContextForPendingState = context;
+}
+
+#pragma mark - UICollectionViewLayout overrides
+
+- (void)prepareLayout
+{
+  ASDisplayNodeAssertMainThread();
+  [super prepareLayout];
+  ASCollectionLayoutContext *context = [self layoutContextWithElements:_collectionNode.visibleElements];
+  
+  ASCollectionLayoutState *state = nil;
+  {
+    ASDN::MutexLocker l(__instanceLock__);
+    if (_pendingState != nil && ASObjectIsEqual(_layoutContextForPendingState, context)) {
+      // Looks like we can use the pending state. Great!
+      state = _pendingState;
+      _pendingState = nil;
+      _layoutContextForPendingState = nil;
+    }
+  }
+  
+  if (state == nil) {
+    state = [_layoutDelegate calculateLayoutWithContext:context];
+  }
+  
+  _state = state;
+}
+
+- (void)invalidateLayout
+{
+  ASDisplayNodeAssertMainThread();
+  [super invalidateLayout];
+  _state = nil;
+}
+
+- (CGSize)collectionViewContentSize
+{
+  ASDisplayNodeAssertMainThread();
+  ASDisplayNodeAssertNotNil(_state, @"Collection layout state should not be nil at this point");
+  return _state.contentSize;
+}
+
+- (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect
+{
+  NSMutableArray *attributesInRect = [NSMutableArray array];
+  NSMapTable *attrsMap = _state.elementToLayoutArrtibutesMap;
+  for (ASCollectionElement *element in attrsMap) {
+    UICollectionViewLayoutAttributes *attrs = [attrsMap objectForKey:element];
+    if (CGRectIntersectsRect(rect, attrs.frame)) {
+      [attributesInRect addObject:attrs];
+    }
+  }
+  return attributesInRect;
+}
+
+- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath
+{
+  ASCollectionLayoutState *state = _state;
+  ASCollectionElement *element = [state.elements elementForItemAtIndexPath:indexPath];
+  return [state.elementToLayoutArrtibutesMap objectForKey:element];
+}
+
+- (UICollectionViewLayoutAttributes *)layoutAttributesForSupplementaryViewOfKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
+{
+  ASCollectionLayoutState *state = _state;
+  ASCollectionElement *element = [state.elements supplementaryElementOfKind:elementKind atIndexPath:indexPath];
+  return [state.elementToLayoutArrtibutesMap objectForKey:element];
+}
+
+#pragma mark - Private methods
+
+- (CGSize)viewportSize
+{
+  ASCollectionNode *collectionNode = _collectionNode;
+  if (collectionNode != nil && !collectionNode.isNodeLoaded) {
+    // TODO consider calculatedSize as well
+    return collectionNode.threadSafeBounds.size;
+  } else {
+    ASDisplayNodeAssertMainThread();
+    return self.collectionView.bounds.size;
+  }
+}
+
+@end

Source/Private/ASCollectionLayoutContext+Private.h

@@ -0,0 +1,19 @@
+//
+//  ASCollectionLayoutContext+Private.h
+//  AsyncDisplayKit
+//
+//  Created by Huy Nguyen on 10/4/17.
+//  Copyright © 2017 Facebook. All rights reserved.
+//
+
+#import <AsyncDisplayKit/ASCollectionLayoutContext.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface ASCollectionLayoutContext (Private)
+
+- (instancetype)initWithViewportSize:(CGSize)viewportSize elements:(ASElementMap *)elements additionalInfo:(nullable id)additionalInfo;
+
+@end
+
+NS_ASSUME_NONNULL_END

Source/Private/ASTwoDimensionalArrayUtils.h

@@ -37,6 +37,11 @@ extern void ASDeleteElementsInTwoDimensionalArrayAtIndexPaths(NSMutableArray *mu
  */
 extern NSArray<NSIndexPath *> *ASIndexPathsForTwoDimensionalArray(NSArray<NSArray *>* twoDimensionalArray) AS_WARN_UNUSED_RESULT;
 
+/**
+ * Return all the elements of a two-dimensional array, in ascending order.
+ */
+extern NSArray *ASElementsInTwoDimensionalArray(NSArray<NSArray *>* twoDimensionalArray) AS_WARN_UNUSED_RESULT;
+
 /**
  * Attempt to get the object at the given index path. Returns @c nil if the index path is out of bounds.
  */

Source/Private/ASTwoDimensionalArrayUtils.m

@@ -21,9 +21,10 @@
 NSMutableArray<NSMutableArray *> *ASTwoDimensionalArrayDeepMutableCopy(NSArray<NSArray *> *array)
 {
   NSMutableArray *newArray = [NSMutableArray arrayWithCapacity:array.count];
+  NSInteger i = 0;
   for (NSArray *subarray in array) {
     ASDisplayNodeCAssert([subarray isKindOfClass:[NSArray class]], @"This function expects NSArray<NSArray *> *");
-    [newArray addObject:[subarray mutableCopy]];
+    newArray[i++] = [subarray mutableCopy];
   }
   return newArray;
 }
@@ -65,17 +66,31 @@ NSArray *ASIndexPathsForTwoDimensionalArray(NSArray <NSArray *>* twoDimensionalA
 {
   NSMutableArray *result = [NSMutableArray array];
   NSInteger section = 0;
+  NSInteger i = 0;
   for (NSArray *subarray in twoDimensionalArray) {
     ASDisplayNodeCAssert([subarray isKindOfClass:[NSArray class]], @"This function expects NSArray<NSArray *> *");
     NSInteger itemCount = subarray.count;
     for (NSInteger item = 0; item < itemCount; item++) {
-      [result addObject:[NSIndexPath indexPathForItem:item inSection:section]];
+      result[i++] = [NSIndexPath indexPathForItem:item inSection:section];
     }
     section++;
   }
   return result;
 }
 
+NSArray *ASElementsInTwoDimensionalArray(NSArray <NSArray *>* twoDimensionalArray)
+{
+  NSMutableArray *result = [NSMutableArray array];
+  NSInteger i = 0;
+  for (NSArray *subarray in twoDimensionalArray) {
+    ASDisplayNodeCAssert([subarray isKindOfClass:[NSArray class]], @"This function expects NSArray<NSArray *> *");
+    for (id element in subarray) {
+      result[i++] = element;
+    }
+  }
+  return result;
+}
+
 id ASGetElementInTwoDimensionalArray(NSArray *array, NSIndexPath *indexPath)
 {
   ASDisplayNodeCAssertNotNil(indexPath, @"Expected non-nil index path");

Source/TextKit/ASEqualityHashHelpers.h

@@ -9,6 +9,8 @@
 //
 
 #import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+#import <AsyncDisplayKit/ASBaseDefines.h>
 
 #import <string>
 
@@ -16,7 +18,7 @@
 // This is the Hash128to64 function from Google's cityhash (available
 // under the MIT License).  We use it to reduce multiple 64 bit hashes
 // into a single hash.
-inline uint64_t ASHashCombine(const uint64_t upper, const uint64_t lower) {
+ASDISPLAYNODE_INLINE uint64_t ASHashCombine(const uint64_t upper, const uint64_t lower) {
   // Murmur-inspired hashing.
   const uint64_t kMul = 0x9ddfea08eb382d69ULL;
   uint64_t a = (lower ^ upper) * kMul;
@@ -28,13 +30,13 @@ inline uint64_t ASHashCombine(const uint64_t upper, const uint64_t lower) {
 }
 
 #if __LP64__
-inline size_t ASHash64ToNative(uint64_t key) {
+ASDISPLAYNODE_INLINE size_t ASHash64ToNative(uint64_t key) {
   return key;
 }
 #else
 
 // Thomas Wang downscaling hash function
-inline size_t ASHash64ToNative(uint64_t key) {
+ASDISPLAYNODE_INLINE size_t ASHash64ToNative(uint64_t key) {
   key = (~key) + (key << 18);
   key = key ^ (key >> 31);
   key = key * 21;
@@ -45,6 +47,8 @@ inline size_t ASHash64ToNative(uint64_t key) {
 }
 #endif
 
+NSUInteger ASHashFromCGSize(const CGSize size);
+
 NSUInteger ASIntegerArrayHash(const NSUInteger *subhashes, NSUInteger count);
 
 namespace AS {

Source/TextKit/ASEqualityHashHelpers.mm

@@ -10,6 +10,13 @@
 
 #import <AsyncDisplayKit/ASEqualityHashHelpers.h>
 
+#import <functional>
+
+NSUInteger ASHashFromCGSize(const CGSize size)
+{
+  return ASHash64ToNative(ASHashCombine(std::hash<CGFloat>()(size.width), std::hash<CGFloat>()(size.height)));
+}
+
 NSUInteger ASIntegerArrayHash(const NSUInteger *subhashes, NSUInteger count)
 {
   uint64_t result = subhashes[0];

examples/ASCollectionView/Sample/ViewController.m

@@ -44,25 +44,15 @@
 {
   [super viewDidLoad];
   
-  UICollectionViewFlowLayout *layout = [[UICollectionViewFlowLayout alloc] init];
+  self.collectionNode = [[ASCollectionNode alloc] initWithLayoutDelegate:[[ASCollectionFlowLayoutDelegate alloc] init] layoutFacilitator:nil];
-  layout.headerReferenceSize = CGSizeMake(50.0, 50.0);
-  layout.footerReferenceSize = CGSizeMake(50.0, 50.0);
-  
-  // This method is deprecated because we reccommend using ASCollectionNode instead of ASCollectionView.
-  // This functionality & example project remains for users who insist on using ASCollectionView.
-  self.collectionNode = [[ASCollectionNode alloc] initWithFrame:self.view.bounds collectionViewLayout:layout];
   self.collectionNode.dataSource = self;
   self.collectionNode.delegate = self;
   
   self.collectionNode.autoresizingMask = UIViewAutoresizingFlexibleWidth | UIViewAutoresizingFlexibleHeight;
   self.collectionNode.backgroundColor = [UIColor whiteColor];
   
-  // This method is deprecated because we reccommend using ASCollectionNode instead of ASCollectionView.
-  // This functionality & example project remains for users who insist on using ASCollectionView.
-  [self.collectionNode registerSupplementaryNodeOfKind:UICollectionElementKindSectionHeader];
-  [self.collectionNode registerSupplementaryNodeOfKind:UICollectionElementKindSectionFooter];
-  
   [self.view addSubnode:self.collectionNode];
+  self.collectionNode.frame = self.view.bounds;
   
 #if !SIMULATE_WEB_RESPONSE
   self.navigationItem.leftItemsSupplementBackButton = YES;