Levi's comments

AsyncDisplayKit/ASDisplayNode.mm

@@ -2718,13 +2718,18 @@ static const char *ASDisplayNodeDrawingPriorityKey = "ASDrawingPriority";
   return ASEnvironmentStateTraitCollectionPropagationEnabled();
 }
 
+- (ASEnvironmentTraitCollection)environmentTraitCollection
+{
+  return _environmentState.traitCollection;
+}
+
 ASEnvironmentLayoutOptionsForwarding
 ASEnvironmentLayoutExtensibilityForwarding
 
 - (ASTraitCollection *)asyncTraitCollection
 {
   ASDN::MutexLocker l(_propertyLock);
-  return [ASTraitCollection displayTraitsWithASEnvironmentTraitCollection:_environmentState.traitCollection];
+  return [ASTraitCollection traitCollectionWithASEnvironmentTraitCollection:_environmentState.traitCollection];
 }
 
 #if TARGET_OS_TV

AsyncDisplayKit/ASViewController.h

@@ -22,7 +22,17 @@ typedef ASTraitCollection * _Nonnull (^ASDisplayTraitsForTraitWindowSizeBlock)(C
 
 @property (nonatomic, strong, readonly) DisplayNodeType node;
 
-@property (nonatomic, strong) id displayTraitsContext;
+/**
+ *  An optional context to pass along with an ASTraitCollection.
+ *  This can be used to pass any internal state to all subnodes via the ASTraitCollection that is not
+ *  included in UITraitCollection. This could range from more fine-tuned size classes to a class of
+ *  constants that is based upon the new trait collection.
+ *
+ *  Be aware that internally this context is held by a C struct which cannot retain the pointer. Therefore
+ *  ASVC keeps a strong reference to the context to make sure that it stays alive. If you change this value
+ *  it will propagate the change to the subnodes.
+ */
+@property (nonatomic, strong) id _Nullable traitColectionContext;
 
 /**
  * Set this block to customize the ASDisplayTraits returned when the VC transitions to the given traitCollection.

AsyncDisplayKit/ASViewController.mm

@@ -9,6 +9,7 @@
 #import "ASViewController.h"
 #import "ASAssert.h"
 #import "ASDimension.h"
+#import "ASDisplayNodeInternal.h"
 #import "ASDisplayNode+FrameworkPrivate.h"
 #import "ASDisplayNode+Beta.h"
 #import "ASTraitCollection.h"
@@ -50,9 +51,8 @@
 
 - (void)dealloc
 {
-  if (_displayTraitsContext != nil) {
+  if (_traitColectionContext != nil) {
-    ASEnvironmentTraitCollectionClearDisplayContext(self.node);
+    self.traitColectionContext = nil;
-    _displayTraitsContext = nil;
   }
 }
 
@@ -144,16 +144,26 @@
 
 #pragma mark - ASEnvironmentTraitCollection
 
+- (void)setTraitColectionContext:(id)traitColectionContext
+{
+  if (_traitColectionContext != traitColectionContext) {
+    // propagate first so that nodes aren't hanging around with a dealloc'ed pointer
+    ASEnvironmentTraitCollectionUpdateDisplayContext(self.node, traitColectionContext);
+    
+    _traitColectionContext = traitColectionContext;
+  }
+}
+
 - (ASEnvironmentTraitCollection)displayTraitsForTraitCollection:(UITraitCollection *)traitCollection
 {
   if (self.overrideDisplayTraitsWithTraitCollection) {
     ASTraitCollection *asyncTraitCollection = self.overrideDisplayTraitsWithTraitCollection(traitCollection);
-    asyncTraitCollection.isMutable = NO;
+    self.traitColectionContext = asyncTraitCollection.traitColectionContext;
     return [asyncTraitCollection environmentTraitCollection];
   }
   
   ASEnvironmentTraitCollection asyncTraitCollection = ASEnvironmentTraitCollectionFromUITraitCollection(traitCollection);
-  asyncTraitCollection.displayContext = _displayTraitsContext;
+  asyncTraitCollection.displayContext = self.traitColectionContext;
   return asyncTraitCollection;
 }
 
@@ -161,10 +171,10 @@
 {
   if (self.overrideDisplayTraitsWithWindowSize) {
     ASTraitCollection *traitCollection = self.overrideDisplayTraitsWithWindowSize(windowSize);
-    traitCollection.isMutable = NO;
+    self.traitColectionContext = traitCollection.traitColectionContext;
     return [traitCollection environmentTraitCollection];
   }
-  return self.node.environmentState.traitCollection;
+  return self.node.environmentTraitCollection;
 }
 
 - (void)progagateNewDisplayTraits:(ASEnvironmentTraitCollection)traitCollection
@@ -174,7 +184,7 @@
   
   if (ASEnvironmentTraitCollectionIsEqualToASEnvironmentTraitCollection(traitCollection, oldTraitCollection) == NO) {
     environmentState.traitCollection = traitCollection;
-    [self.node setEnvironmentState:environmentState];
+    self.node.environmentState = environmentState;
     [self.node setNeedsLayout];
     
     NSArray<id<ASEnvironment>> *children = [self.node children];

AsyncDisplayKit/Details/ASDataController.mm

@@ -525,7 +525,7 @@ static void *kASSizingQueueContext = &kASSizingQueueContext;
       ASCellNodeBlock nodeBlockPropagatingDisplayTraits = ^{
         ASCellNode *cellNode = nodeBlock();
         id<ASEnvironment> environment = [self.environmentDelegate dataControllerEnvironment];
-        ASEnvironmentStatePropagateDown(cellNode, [environment environmentState].traitCollection);
+        ASEnvironmentStatePropagateDown(cellNode, [environment environmentTraitCollection]);
         return cellNode;
       };
       

AsyncDisplayKit/Details/ASEnvironment.h

@@ -87,10 +87,10 @@ typedef struct ASEnvironmentTraitCollection {
   id __unsafe_unretained displayContext;
 } ASEnvironmentTraitCollection;
 
-extern void ASEnvironmentTraitCollectionClearDisplayContext(id<ASEnvironment> rootEnvironment);
+extern void ASEnvironmentTraitCollectionUpdateDisplayContext(id<ASEnvironment> rootEnvironment, id _Nullable context);
 
 extern ASEnvironmentTraitCollection ASEnvironmentTraitCollectionFromUITraitCollection(UITraitCollection *traitCollection);
-extern BOOL ASEnvironmentTraitCollectionIsEqualToASEnvironmentTraitCollection(ASEnvironmentTraitCollection displayTraits0, ASEnvironmentTraitCollection displayTraits1);
+extern BOOL ASEnvironmentTraitCollectionIsEqualToASEnvironmentTraitCollection(ASEnvironmentTraitCollection lhs, ASEnvironmentTraitCollection rhs);
 
 #pragma mark - ASEnvironmentState
 
@@ -133,6 +133,10 @@ ASDISPLAYNODE_EXTERN_C_END
 /// Returns an NSObject-representation of the environment's ASEnvironmentDisplayTraits
 - (ASTraitCollection *)asyncTraitCollection;
 
+/// Returns a struct-representation of the environment's ASEnvironmentDisplayTraits. This only exists as a internal
+/// convenience method. Users should access the trait collections through the NSObject based asyncTraitCollection API
+- (ASEnvironmentTraitCollection)environmentTraitCollection;
+
 @end
 
 // ASCollection/TableNodes don't actually have ASCellNodes as subnodes. Because of this we can't rely on display trait

AsyncDisplayKit/Details/ASEnvironment.mm

@@ -26,11 +26,11 @@ ASEnvironmentHierarchyState _ASEnvironmentHierarchyStateMakeDefault()
   };
 }
 
-extern void ASEnvironmentTraitCollectionClearDisplayContext(id<ASEnvironment> rootEnvironment)
+extern void ASEnvironmentTraitCollectionUpdateDisplayContext(id<ASEnvironment> rootEnvironment, id context)
 {
   ASEnvironmentState envState = [rootEnvironment environmentState];
   ASEnvironmentTraitCollection displayTraits = envState.traitCollection;
-  displayTraits.displayContext = nil;
+  displayTraits.displayContext = context;
   envState.traitCollection = displayTraits;
   [rootEnvironment setEnvironmentState:envState];
   
@@ -61,15 +61,15 @@ ASEnvironmentTraitCollection ASEnvironmentTraitCollectionFromUITraitCollection(U
   return asyncTraitCollection;
 }
 
-BOOL ASEnvironmentTraitCollectionIsEqualToASEnvironmentTraitCollection(ASEnvironmentTraitCollection traitCollection0, ASEnvironmentTraitCollection traitCollection1)
+BOOL ASEnvironmentTraitCollectionIsEqualToASEnvironmentTraitCollection(ASEnvironmentTraitCollection lhs, ASEnvironmentTraitCollection rhs)
 {
   return
-    traitCollection0.verticalSizeClass == traitCollection1.verticalSizeClass &&
+    lhs.verticalSizeClass == rhs.verticalSizeClass &&
-    traitCollection0.horizontalSizeClass == traitCollection1.horizontalSizeClass &&
+    lhs.horizontalSizeClass == rhs.horizontalSizeClass &&
-    traitCollection0.displayScale == traitCollection1.displayScale &&
+    lhs.displayScale == rhs.displayScale &&
-    traitCollection0.userInterfaceIdiom == traitCollection1.userInterfaceIdiom &&
+    lhs.userInterfaceIdiom == rhs.userInterfaceIdiom &&
-    traitCollection0.forceTouchCapability == traitCollection1.forceTouchCapability &&
+    lhs.forceTouchCapability == rhs.forceTouchCapability &&
-    traitCollection0.displayContext == traitCollection1.displayContext;
+    lhs.displayContext == rhs.displayContext;
 }
 
 ASEnvironmentState ASEnvironmentStateMakeDefault()

AsyncDisplayKit/Details/ASTraitCollection.h

@@ -13,16 +13,41 @@
 
 @interface ASTraitCollection : NSObject
 
-@property (nonatomic, assign) BOOL isMutable;
+@property (nonatomic, assign, readonly) CGFloat displayScale;
+@property (nonatomic, assign, readonly) UIUserInterfaceSizeClass horizontalSizeClass;
+@property (nonatomic, assign, readonly) UIUserInterfaceIdiom userInterfaceIdiom;
+@property (nonatomic, assign, readonly) UIUserInterfaceSizeClass verticalSizeClass;
+@property (nonatomic, assign, readonly) UIForceTouchCapability forceTouchCapability;
 
-@property (nonatomic, assign) CGFloat displayScale;
+/**
-@property (nonatomic, assign) UIUserInterfaceSizeClass horizontalSizeClass;
+ *  An optional context to pass along with an ASTraitCollection.
-@property (nonatomic, assign) UIUserInterfaceIdiom userInterfaceIdiom;
+ *  This can be used to pass any internal state to all subnodes via the ASTraitCollection that is not
-@property (nonatomic, assign) UIUserInterfaceSizeClass verticalSizeClass;
+ *  included in UITraitCollection. This could range from more fine-tuned size classes to a class of
-@property (nonatomic, assign) UIForceTouchCapability forceTouchCapability;
+ *  constants that is based upon the new trait collection.
+ *
+ *  Be aware that internally this context is held by a C struct which cannot retain the pointer. 
+ *  ASTraitCollection is generally a very short-lived class, existing only to provide a non-struct API
+ *  to trait collections. When an ASTraitCollection is returned via one of ASViewController's 2 
+ *  custom trait collection creation blocks, traitColectionContext is assigned to the VC's traitColectionContext.
+ *  This makes sure that the VC is the owner of the context and ASEnvironmentTraitCollections will not
+ *  have a reference to a dangling pointer.
+ */
+@property (nonatomic, strong, readonly) id traitColectionContext;
+
+
++ (ASTraitCollection *)traitCollectionWithASEnvironmentTraitCollection:(ASEnvironmentTraitCollection)traits;
+
++ (ASTraitCollection *)traitCollectionWithUITraitCollection:(UITraitCollection *)traitCollection
+                                     traitCollectionContext:(id)traitCollectionContext;
+
+
++ (ASTraitCollection *)traitCollectionWithDisplayScale:(CGFloat)displayScale
+                                    userInterfaceIdiom:(UIUserInterfaceIdiom)userInterfaceIdiom
+                                   horizontalSizeClass:(UIUserInterfaceSizeClass)horizontalSizeClass
+                                     verticalSizeClass:(UIUserInterfaceSizeClass)verticalSizeClass
+                                  forceTouchCapability:(UIForceTouchCapability)forceTouchCapability
+                                traitCollectionContext:(id)traitCollectionContext;
 
-+ (ASTraitCollection *)displayTraitsWithASEnvironmentTraitCollection:(ASEnvironmentTraitCollection)traits;
-+ (ASTraitCollection *)displayTraitsWithUITraitCollection:(UITraitCollection *)traitCollection;
 
 - (ASEnvironmentTraitCollection)environmentTraitCollection;
 

AsyncDisplayKit/Details/ASTraitCollection.m

@@ -12,68 +12,74 @@
 
 @implementation ASTraitCollection
 
-- (instancetype)init
+- (instancetype)initWithDisplayScale:(CGFloat)displayScale
+                  userInterfaceIdiom:(UIUserInterfaceIdiom)userInterfaceIdiom
+                 horizontalSizeClass:(UIUserInterfaceSizeClass)horizontalSizeClass
+                   verticalSizeClass:(UIUserInterfaceSizeClass)verticalSizeClass
+                forceTouchCapability:(UIForceTouchCapability)forceTouchCapability
+              traitCollectionContext:(id)traitCollectionContext
 {
     self = [super init];
     if (self) {
-        _isMutable = YES;
+      _displayScale = displayScale;
+      _userInterfaceIdiom = userInterfaceIdiom;
+      _horizontalSizeClass = horizontalSizeClass;
+      _verticalSizeClass = verticalSizeClass;
+      _forceTouchCapability = forceTouchCapability;
+      _traitColectionContext = traitCollectionContext;
     }
     return self;
 }
 
-- (void)setDisplayScale:(CGFloat)displayScale
++ (ASTraitCollection *)traitCollectionWithDisplayScale:(CGFloat)displayScale
+                                    userInterfaceIdiom:(UIUserInterfaceIdiom)userInterfaceIdiom
+                                   horizontalSizeClass:(UIUserInterfaceSizeClass)horizontalSizeClass
+                                     verticalSizeClass:(UIUserInterfaceSizeClass)verticalSizeClass
+                                  forceTouchCapability:(UIForceTouchCapability)forceTouchCapability
+                                traitCollectionContext:(id)traitCollectionContext
 {
-  ASDisplayNodeAssert(self.isMutable, @"ASDisplayTraits is no longer mutable");
+  return [[[self class] alloc] initWithDisplayScale:displayScale
-  _displayScale = displayScale;
+                                 userInterfaceIdiom:userInterfaceIdiom
+                                horizontalSizeClass:horizontalSizeClass
+                                  verticalSizeClass:verticalSizeClass
+                               forceTouchCapability:forceTouchCapability
+                             traitCollectionContext:traitCollectionContext];
 }
 
-- (void)setHorizontalSizeClass:(UIUserInterfaceSizeClass)horizontalSizeClass
++ (ASTraitCollection *)traitCollectionWithASEnvironmentTraitCollection:(ASEnvironmentTraitCollection)traits
 {
-  ASDisplayNodeAssert(self.isMutable, @"ASDisplayTraits is no longer mutable");
+  return [[[self class] alloc] initWithDisplayScale:traits.displayScale
-  _horizontalSizeClass = horizontalSizeClass;
+                                 userInterfaceIdiom:traits.userInterfaceIdiom
+                                horizontalSizeClass:traits.horizontalSizeClass
+                                  verticalSizeClass:traits.verticalSizeClass
+                               forceTouchCapability:traits.forceTouchCapability
+                             traitCollectionContext:traits.displayContext];
+
 }
 
-- (void)setUserInterfaceIdiom:(UIUserInterfaceIdiom)userInterfaceIdiom
++ (ASTraitCollection *)traitCollectionWithUITraitCollection:(UITraitCollection *)traitCollection
+                                     traitCollectionContext:(id)traitCollectionContext
 {
-  ASDisplayNodeAssert(self.isMutable, @"ASDisplayTraits is no longer mutable");
+  ASTraitCollection *asyncTraitCollection = nil;
-  _userInterfaceIdiom = userInterfaceIdiom;
+  if (AS_AT_LEAST_IOS9) {
-}
+    asyncTraitCollection = [[[self class] alloc] initWithDisplayScale:traitCollection.displayScale
-
+                                                   userInterfaceIdiom:traitCollection.userInterfaceIdiom
-- (void)setVerticalSizeClass:(UIUserInterfaceSizeClass)verticalSizeClass
+                                                  horizontalSizeClass:traitCollection.horizontalSizeClass
-{
+                                                    verticalSizeClass:traitCollection.verticalSizeClass
-  ASDisplayNodeAssert(self.isMutable, @"ASDisplayTraits is no longer mutable");
+                                                 forceTouchCapability:traitCollection.forceTouchCapability
-  _verticalSizeClass = verticalSizeClass;
+                                               traitCollectionContext:traitCollectionContext];
-}
-
-- (void)setForceTouchCapability:(UIForceTouchCapability)forceTouchCapability
-{
-  ASDisplayNodeAssert(self.isMutable, @"ASDisplayTraits is no longer mutable");
-  _forceTouchCapability = forceTouchCapability;
-}
-
-+ (ASTraitCollection *)displayTraitsWithASEnvironmentTraitCollection:(ASEnvironmentTraitCollection)traits
-{
-  ASTraitCollection *traitCollection = [[ASTraitCollection alloc] init];
-  traitCollection.displayScale = traits.displayScale;
-  traitCollection.horizontalSizeClass = traits.horizontalSizeClass;
-  traitCollection.verticalSizeClass = traits.verticalSizeClass;
-  traitCollection.userInterfaceIdiom = traits.userInterfaceIdiom;
-  traitCollection.forceTouchCapability = traits.forceTouchCapability;
-  return traitCollection;
-}
-
-+ (ASTraitCollection *)displayTraitsWithUITraitCollection:(UITraitCollection *)traitCollection
-{
-  ASTraitCollection *asyncTraitCollection = [[ASTraitCollection alloc] init];
-  if (AS_AT_LEAST_IOS8) {
-    asyncTraitCollection.displayScale = traitCollection.displayScale;
-    asyncTraitCollection.horizontalSizeClass = traitCollection.horizontalSizeClass;
-    asyncTraitCollection.verticalSizeClass = traitCollection.verticalSizeClass;
-    asyncTraitCollection.userInterfaceIdiom = traitCollection.userInterfaceIdiom;
-    if (AS_AT_LEAST_IOS9) {
-      asyncTraitCollection.forceTouchCapability = traitCollection.forceTouchCapability;
-    }
   }
+  else if (AS_AT_LEAST_IOS8) {
+    asyncTraitCollection = [[[self class] alloc] initWithDisplayScale:traitCollection.displayScale
+                                                   userInterfaceIdiom:traitCollection.userInterfaceIdiom
+                                                  horizontalSizeClass:traitCollection.horizontalSizeClass
+                                                    verticalSizeClass:traitCollection.verticalSizeClass
+                                                 forceTouchCapability:0
+                                               traitCollectionContext:traitCollectionContext];
+  } else {
+    asyncTraitCollection = [[[self class] alloc] init];
+  }
+  
   return asyncTraitCollection;
 }
 
@@ -85,6 +91,7 @@
     .userInterfaceIdiom = self.userInterfaceIdiom,
     .verticalSizeClass = self.verticalSizeClass,
     .forceTouchCapability = self.forceTouchCapability,
+    .displayContext = self.traitColectionContext,
   };
 }
 

AsyncDisplayKit/Layout/ASLayoutSpec.mm

@@ -201,13 +201,18 @@
   }
 }
 
+- (ASEnvironmentTraitCollection)environmentTraitCollection
+{
+  return _environmentState.traitCollection;
+}
+
 ASEnvironmentLayoutOptionsForwarding
 ASEnvironmentLayoutExtensibilityForwarding
 
 - (ASTraitCollection *)asyncTraitCollection
 {
   ASDN::MutexLocker l(_propertyLock);
-  return [ASTraitCollection displayTraitsWithASEnvironmentTraitCollection:_environmentState.traitCollection];
+  return [ASTraitCollection traitCollectionWithASEnvironmentTraitCollection:_environmentState.traitCollection];
 }
 
 @end

AsyncDisplayKit/Private/ASDisplayNodeInternal.h

@@ -228,4 +228,10 @@ FOUNDATION_EXPORT NSString * const ASRenderingEngineDidDisplayNodesScheduledBefo
  */
 - (ASDisplayNode *)_supernodeWithClass:(Class)supernodeClass checkViewHierarchy:(BOOL)checkViewHierarchy;
 
+/**
+ *  Convenience method to access this node's trait collection struct. Externally, users should interact
+ *  with the trait collection via ASTraitCollection
+ */
+- (ASEnvironmentTraitCollection)environmentTraitCollection;
+
 @end

examples/ASTraitCollection/Sample/KittenNode.m

@@ -152,9 +152,12 @@ static const CGFloat kInnerPadding = 10.0f;
   OverrideViewController *overrideVC = [[OverrideViewController alloc] init];
   
   overrideVC.overrideDisplayTraitsWithTraitCollection = ^(UITraitCollection *traitCollection) {
-    ASTraitCollection *asyncTraitCollection = [ASTraitCollection displayTraitsWithUITraitCollection:traitCollection];
+    ASTraitCollection *asyncTraitCollection = [ASTraitCollection traitCollectionWithDisplayScale:traitCollection.displayScale
-    asyncTraitCollection.horizontalSizeClass = UIUserInterfaceSizeClassCompact;
+                                                                              userInterfaceIdiom:traitCollection.userInterfaceIdiom
-    asyncTraitCollection.verticalSizeClass = UIUserInterfaceSizeClassCompact;
+                                                                             horizontalSizeClass:UIUserInterfaceSizeClassCompact
+                                                                               verticalSizeClass:UIUserInterfaceSizeClassCompact
+                                                                            forceTouchCapability:traitCollection.forceTouchCapability
+                                                                          traitCollectionContext:nil];
     return asyncTraitCollection;
   };