[ASTextNode] Optimize handling of constrained size to almost never recreate NSLayoutManager

This also fixes two fairly subtle but serious bugs, #1076 and #1046.

AsyncDisplayKit/ASCollectionNode+Beta.h

@@ -6,6 +6,7 @@
  * of patent rights can be found in the PATENTS file in the same directory.
  */
 
+#import "ASCollectionNode.h"
 @protocol ASCollectionViewLayoutFacilitatorProtocol;
 
 NS_ASSUME_NONNULL_BEGIN

AsyncDisplayKit/ASTextNode.mm

@@ -35,13 +35,13 @@ static NSString *ASTextNodeTruncationTokenAttributeName = @"ASTextNodeTruncation
 
 - (instancetype)initWithRenderer:(ASTextKitRenderer *)renderer
                       textOrigin:(CGPoint)textOrigin
-                 backgroundColor:(CGColorRef)backgroundColor;
+                 backgroundColor:(UIColor *)backgroundColor;
 
 @property (nonatomic, strong, readonly) ASTextKitRenderer *renderer;
 
 @property (nonatomic, assign, readonly) CGPoint textOrigin;
 
-@property (nonatomic, assign, readonly) CGColorRef backgroundColor;
+@property (nonatomic, strong, readonly) UIColor *backgroundColor;
 
 @end
 
@@ -49,20 +49,18 @@ static NSString *ASTextNodeTruncationTokenAttributeName = @"ASTextNodeTruncation
 
 - (instancetype)initWithRenderer:(ASTextKitRenderer *)renderer
                       textOrigin:(CGPoint)textOrigin
-                 backgroundColor:(CGColorRef)backgroundColor
+                 backgroundColor:(UIColor *)backgroundColor
 {
   if (self = [super init]) {
     _renderer = renderer;
     _textOrigin = textOrigin;
-    _backgroundColor = CGColorRetain(backgroundColor);
+    _backgroundColor = backgroundColor;
   }
   return self;
 }
 
 - (void)dealloc
 {
-  CGColorRelease(_backgroundColor);
-  
   // Destruction of the layout managers/containers/text storage is quite
   // expensive, and can take some time, so we dispatch onto a bg queue to
   // actually dealloc.
@@ -182,7 +180,7 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
   NSString *truncationString = [_composedTruncationString string];
   if (plainString.length > 50)
     plainString = [[plainString substringToIndex:50] stringByAppendingString:@"\u2026"];
-  return [NSString stringWithFormat:@"<%@: %p; text = \"%@\"; truncation string = \"%@\"; frame = %@>", self.class, self, plainString, truncationString, self.nodeLoaded ? NSStringFromCGRect(self.layer.frame) : nil];
+  return [NSString stringWithFormat:@"<%@: %p; text = \"%@\"; truncation string = \"%@\"; frame = %@; renderer = %p>", self.class, self, plainString, truncationString, self.nodeLoaded ? NSStringFromCGRect(self.layer.frame) : nil, _renderer];
 }
 
 #pragma mark - ASDisplayNode
@@ -240,13 +238,13 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
 - (void)setFrame:(CGRect)frame
 {
   [super setFrame:frame];
-  [self _invalidateRendererIfNeeded:frame.size];
+  [self _invalidateRendererIfNeededForBoundsSize:frame.size];
 }
 
 - (void)setBounds:(CGRect)bounds
 {
   [super setBounds:bounds];
-  [self _invalidateRendererIfNeeded:bounds.size];
+  [self _invalidateRendererIfNeededForBoundsSize:bounds.size];
 }
 
 #pragma mark - Renderer Management
@@ -291,12 +289,12 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
 
 - (void)_invalidateRendererIfNeeded
 {
-  [self _invalidateRendererIfNeeded:self.bounds.size];
+  [self _invalidateRendererIfNeededForBoundsSize:self.bounds.size];
 }
 
-- (void)_invalidateRendererIfNeeded:(CGSize)newSize
+- (void)_invalidateRendererIfNeededForBoundsSize:(CGSize)boundsSize
 {
-  if ([self _needInvalidateRenderer:newSize]) {
+  if ([self _needInvalidateRendererForBoundsSize:boundsSize]) {
     // Our bounds of frame have changed to a size that is not identical to our constraining size,
     // so our previous layout information is invalid, and TextKit may draw at the
     // incorrect origin.
@@ -305,7 +303,7 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
   }
 }
 
-- (BOOL)_needInvalidateRenderer:(CGSize)newSize
+- (BOOL)_needInvalidateRendererForBoundsSize:(CGSize)boundsSize
 {
   if (!_renderer) {
     return YES;
@@ -313,9 +311,9 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
   
   // If the size is not the same as the constraint we provided to the renderer, start out assuming we need
   // a new one.  However, there are common cases where the constrained size doesn't need to be the same as calculated.
-  CGSize oldSize = _renderer.constrainedSize;
+  CGSize rendererConstrainedSize = _renderer.constrainedSize;
   
-  if (CGSizeEqualToSize(newSize, oldSize)) {
+  if (CGSizeEqualToSize(boundsSize, rendererConstrainedSize)) {
     return NO;
   } else {
     // It is very common to have a constrainedSize with a concrete, specific width but +Inf height.
@@ -324,7 +322,12 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
     // experience truncation and don't need to recreate the renderer with the size it already calculated,
     // as this would essentially serve to set its constrainedSize to be its calculatedSize (unnecessary).
     ASLayout *layout = self.calculatedLayout;
-    if (layout != nil && CGSizeEqualToSize(newSize, layout.size)) {
+    if (layout != nil && CGSizeEqualToSize(boundsSize, layout.size)) {
+      if (!CGSizeEqualToSize(boundsSize, rendererConstrainedSize)) {
+        // Don't bother changing _constrainedSize, as ASDisplayNode's -measure: method would have a cache miss
+        // and ask us to recalculate layout if it were called with the same calculatedSize that got us to this point!
+        _renderer.constrainedSize = boundsSize;
+      }
       return NO;
     } else {
       return YES;
@@ -409,12 +412,10 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
 
   // Fill background
   if (!isRasterizing) {
-    CGColorRef backgroundColor = parameters.backgroundColor;
+    UIColor *backgroundColor = parameters.backgroundColor;
     if (backgroundColor) {
-      CGContextSetFillColorWithColor(context, backgroundColor);
+      [backgroundColor setFill];
-      CGContextSetBlendMode(context, kCGBlendModeCopy);
+      UIRectFillUsingBlendMode(CGContextGetClipBoundingBox(context), kCGBlendModeCopy);
-      CGContextFillRect(context, CGContextGetClipBoundingBox(context));
-      CGContextSetBlendMode(context, kCGBlendModeNormal);
     }
   }
 
@@ -430,14 +431,15 @@ static NSArray *DefaultLinkAttributeNames = @[ NSLinkAttributeName ];
 
 - (NSObject *)drawParametersForAsyncLayer:(_ASDisplayLayer *)layer
 {
-  [self _invalidateRendererIfNeeded];
+  CGRect bounds = self.bounds;
+  [self _invalidateRendererIfNeededForBoundsSize:bounds.size];
 
   // Offset the text origin by any shadow padding
   UIEdgeInsets shadowPadding = [self shadowPadding];
-  CGPoint textOrigin = CGPointMake(self.bounds.origin.x - shadowPadding.left, self.bounds.origin.y - shadowPadding.top);
+  CGPoint textOrigin = CGPointMake(bounds.origin.x - shadowPadding.left, bounds.origin.y - shadowPadding.top);
   return [[ASTextNodeDrawParameters alloc] initWithRenderer:[self _renderer]
                                                  textOrigin:textOrigin
-                                            backgroundColor:self.backgroundColor.CGColor];
+                                            backgroundColor:self.backgroundColor];
 }
 
 #pragma mark - Attributes

AsyncDisplayKit/TextKit/ASTextKitContext.h

@@ -30,6 +30,8 @@
                          constrainedSize:(CGSize)constrainedSize
                     layoutManagerFactory:(NSLayoutManager*(*)(void))layoutManagerFactory;
 
+@property (nonatomic, assign, readwrite) CGSize constrainedSize;
+
 /**
  All operations on TextKit values MUST occur within this locked context.  Simultaneous access (even non-mutative) to
  TextKit components may cause crashes.

AsyncDisplayKit/TextKit/ASTextKitContext.mm

@@ -49,6 +49,16 @@
   return self;
 }
 
+- (CGSize)constrainedSize
+{
+  return _textContainer.size;
+}
+
+- (void)setConstrainedSize:(CGSize)constrainedSize
+{
+  _textContainer.size = constrainedSize;
+}
+
 - (void)performBlockWithLockedTextKitComponents:(void (^)(NSLayoutManager *,
                                                           NSTextStorage *,
                                                           NSTextContainer *))block

AsyncDisplayKit/TextKit/ASTextKitRenderer.h

@@ -37,7 +37,6 @@
 
 /**
  Designated Initializer
-dvlkferufedgjnhjjfhldjedlunvtdtv
  @discussion Sizing will occur as a result of initialization, so be careful when/where you use this.
  */
 - (instancetype)initWithTextKitAttributes:(const ASTextKitAttributes &)textComponentAttributes
@@ -51,7 +50,7 @@ dvlkferufedgjnhjjfhldjedlunvtdtv
 
 @property (nonatomic, assign, readonly) ASTextKitAttributes attributes;
 
-@property (nonatomic, assign, readonly) CGSize constrainedSize;
+@property (nonatomic, assign, readwrite) CGSize constrainedSize;
 
 #pragma mark - Drawing
 /*

AsyncDisplayKit/TextKit/ASTextKitRenderer.mm

@@ -17,6 +17,9 @@
 #import "ASTextKitTailTruncater.h"
 #import "ASTextKitTruncating.h"
 
+//#define LOG(...) NSLog(__VA_ARGS__)
+#define LOG(...)
+
 static NSCharacterSet *_defaultAvoidTruncationCharacterSet()
 {
   static NSCharacterSet *truncationCharacterSet;
@@ -65,12 +68,10 @@ static NSCharacterSet *_defaultAvoidTruncationCharacterSet()
 {
   if (!_truncater) {
     ASTextKitAttributes attributes = _attributes;
-    // We must inset the constrained size by the size of the shadower.
+    NSCharacterSet *avoidTailTruncationSet = attributes.avoidTailTruncationSet ? : _defaultAvoidTruncationCharacterSet();
-    CGSize shadowConstrainedSize = [[self shadower] insetSizeWithConstrainedSize:_constrainedSize];
     _truncater = [[ASTextKitTailTruncater alloc] initWithContext:[self context]
                                       truncationAttributedString:attributes.truncationAttributedString
-                                          avoidTailTruncationSet:attributes.avoidTailTruncationSet ?: _defaultAvoidTruncationCharacterSet()
+                                          avoidTailTruncationSet:avoidTailTruncationSet];
-                                                 constrainedSize:shadowConstrainedSize];
   }
   return _truncater;
 }
@@ -79,6 +80,7 @@ static NSCharacterSet *_defaultAvoidTruncationCharacterSet()
 {
   if (!_context) {
     ASTextKitAttributes attributes = _attributes;
+    // We must inset the constrained size by the size of the shadower.
     CGSize shadowConstrainedSize = [[self shadower] insetSizeWithConstrainedSize:_constrainedSize];
     _context = [[ASTextKitContext alloc] initWithAttributedString:attributes.attributedString
                                                     lineBreakMode:attributes.lineBreakMode
@@ -92,6 +94,30 @@ static NSCharacterSet *_defaultAvoidTruncationCharacterSet()
 
 #pragma mark - Sizing
 
+- (CGSize)size
+{
+  if (!_sizeIsCalculated) {
+    [self _calculateSize];
+    _sizeIsCalculated = YES;
+  }
+  return _calculatedSize;
+}
+
+- (void)setConstrainedSize:(CGSize)constrainedSize
+{
+  if (!CGSizeEqualToSize(constrainedSize, _constrainedSize)) {
+    _sizeIsCalculated = NO;
+    _constrainedSize = constrainedSize;
+    // If the context isn't created yet, it will be initialized with the appropriate size when next accessed.
+    if (_context) {
+      // If we're updating an existing context, make sure to use the same inset logic used during initialization.
+      // This codepath allows us to reuse the
+      CGSize shadowConstrainedSize = [[self shadower] insetSizeWithConstrainedSize:constrainedSize];
+      _context.constrainedSize = shadowConstrainedSize;
+    }
+  }
+}
+
 - (void)_calculateSize
 {
   // Force glyph generation and layout, which may not have happened yet (and isn't triggered by
@@ -111,16 +137,7 @@ static NSCharacterSet *_defaultAvoidTruncationCharacterSet()
   // to make sure our width calculations aren't being offset by glyphs going beyond the constrained rect.
   boundingRect = CGRectIntersection(boundingRect, {.size = constrainedRect.size});
 
-  _calculatedSize = [_shadower outsetSizeWithInsetSize:CGSizeMake(boundingRect.size.width + boundingRect.origin.x, boundingRect.size.height + boundingRect.origin.y)];
+  _calculatedSize = [_shadower outsetSizeWithInsetSize:boundingRect.size];
-}
-
-- (CGSize)size
-{
-  if (!_sizeIsCalculated) {
-    [self _calculateSize];
-    _sizeIsCalculated = YES;
-  }
-  return _calculatedSize;
 }
 
 #pragma mark - Drawing
@@ -136,8 +153,12 @@ static NSCharacterSet *_defaultAvoidTruncationCharacterSet()
   [[self shadower] setShadowInContext:context];
   UIGraphicsPushContext(context);
 
+  LOG(@"%@, shadowInsetBounds = %@",self, NSStringFromCGRect(shadowInsetBounds));
+  
   [[self context] performBlockWithLockedTextKitComponents:^(NSLayoutManager *layoutManager, NSTextStorage *textStorage, NSTextContainer *textContainer) {
+    LOG(@"usedRect: %@", NSStringFromCGRect([layoutManager usedRectForTextContainer:textContainer]));
     NSRange glyphRange = [layoutManager glyphRangeForTextContainer:textContainer];
+    LOG(@"boundingRect: %@", NSStringFromCGRect([layoutManager boundingRectForGlyphRange:glyphRange inTextContainer:textContainer]));
     [layoutManager drawBackgroundForGlyphRange:glyphRange atPoint:shadowInsetBounds.origin];
     [layoutManager drawGlyphsForGlyphRange:glyphRange atPoint:shadowInsetBounds.origin];
   }];

AsyncDisplayKit/TextKit/ASTextKitTailTruncater.mm

@@ -18,7 +18,6 @@
   __weak ASTextKitContext *_context;
   NSAttributedString *_truncationAttributedString;
   NSCharacterSet *_avoidTailTruncationSet;
-  CGSize _constrainedSize;
 }
 @synthesize visibleRanges = _visibleRanges;
 @synthesize truncationStringRect = _truncationStringRect;
@@ -26,13 +25,11 @@
 - (instancetype)initWithContext:(ASTextKitContext *)context
      truncationAttributedString:(NSAttributedString *)truncationAttributedString
          avoidTailTruncationSet:(NSCharacterSet *)avoidTailTruncationSet
-                constrainedSize:(CGSize)constrainedSize
 {
   if (self = [super init]) {
     _context = context;
     _truncationAttributedString = truncationAttributedString;
     _avoidTailTruncationSet = avoidTailTruncationSet;
-    _constrainedSize = constrainedSize;
 
     [self _truncate];
   }

AsyncDisplayKit/TextKit/ASTextKitTruncating.h

@@ -31,7 +31,6 @@
  */
 - (instancetype)initWithContext:(ASTextKitContext *)context
      truncationAttributedString:(NSAttributedString *)truncationAttributedString
-         avoidTailTruncationSet:(NSCharacterSet *)avoidTailTruncationSet
+         avoidTailTruncationSet:(NSCharacterSet *)avoidTailTruncationSet;
-                constrainedSize:(CGSize)constrainedSize;
 
 @end

AsyncDisplayKitTests/ASTextKitTruncationTests.mm

@@ -50,8 +50,7 @@
   }];
   ASTextKitTailTruncater *tailTruncater = [[ASTextKitTailTruncater alloc] initWithContext:context
                                                                truncationAttributedString:nil
-                                                                   avoidTailTruncationSet:nil
+                                                                   avoidTailTruncationSet:nil];
-                                                                          constrainedSize:constrainedSize];
   XCTAssert(NSEqualRanges(textKitVisibleRange, tailTruncater.visibleRanges[0]));
 }
 
@@ -67,8 +66,7 @@
                                                             layoutManagerFactory:nil];
   ASTextKitTailTruncater *tailTruncater = [[ASTextKitTailTruncater alloc] initWithContext:context
                                                                truncationAttributedString:[self _simpleTruncationAttributedString]
-                                                                   avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@""]
+                                                                   avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@""]];
-                                                                          constrainedSize:constrainedSize];
   __block NSString *drawnString;
   [context performBlockWithLockedTextKitComponents:^(NSLayoutManager *layoutManager, NSTextStorage *textStorage, NSTextContainer *textContainer) {
     drawnString = textStorage.string;
@@ -90,8 +88,7 @@
                                                             layoutManagerFactory:nil];
   ASTextKitTailTruncater *tailTruncater = [[ASTextKitTailTruncater alloc] initWithContext:context
                                                                truncationAttributedString:[self _simpleTruncationAttributedString]
-                                                                   avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@"."]
+                                                                   avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@"."]];
-                                                                          constrainedSize:constrainedSize];
   (void)tailTruncater;
   __block NSString *drawnString;
   [context performBlockWithLockedTextKitComponents:^(NSLayoutManager *layoutManager, NSTextStorage *textStorage, NSTextContainer *textContainer) {
@@ -114,8 +111,7 @@
                                                             layoutManagerFactory:nil];
   ASTextKitTailTruncater *tailTruncater = [[ASTextKitTailTruncater alloc] initWithContext:context
                                                                truncationAttributedString:[self _simpleTruncationAttributedString]
-                                                                   avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@"."]
+                                                                   avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@"."]];
-                                                                          constrainedSize:constrainedSize];
   // So Xcode doesn't yell at me for an unused var...
   (void)tailTruncater;
   __block NSString *drawnString;
@@ -139,8 +135,7 @@
                                                             layoutManagerFactory:nil];
   XCTAssertNoThrow([[ASTextKitTailTruncater alloc] initWithContext:context
                                         truncationAttributedString:[self _simpleTruncationAttributedString]
-                                            avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@"."]
+                                            avoidTailTruncationSet:[NSCharacterSet characterSetWithCharactersInString:@"."]]);
-                                                   constrainedSize:constrainedSize]);
 }
 
 @end