Try to allow memoization of components with state

ComponentKit/Core/CKComponent.mm

@@ -300,4 +300,9 @@ - (id)scopeFrameToken
   return _scopeHandle ? @(_scopeHandle.globalIdentifier) : nil;
 }
 
+- (CKComponentScopeHandleIdentifier)globalIdentifier
+{
+  return _scopeHandle.globalIdentifier;
+}
+
 @end

ComponentKit/Core/CKComponentInternal.h

@@ -13,6 +13,7 @@
 #import <ComponentKit/ComponentMountContext.h>
 #import <ComponentKit/CKComponent.h>
 #import <ComponentKit/CKComponentLayout.h>
+#import <ComponentKit/CKComponentScopeTypes.h>
 
 @interface CKComponent ()
 
@@ -67,6 +68,9 @@
 /** For internal use only; don't touch this. */
 @property (nonatomic, strong, readonly) id scopeFrameToken;
 
+/** For internal use only; don't touch this. */
+@property (nonatomic, readonly) CKComponentScopeHandleIdentifier globalIdentifier;
+
 @end
 
 // Internal, for CKComponent

ComponentKit/Core/CKComponentMemoizer.mm

@@ -14,11 +14,47 @@
 #import "CKComponentSubclass.h"
 #import "CKMacros.h"
 #import "CKInternalHelpers.h"
+#import "CKThreadLocalComponentScope.h"
+#import "CKComponentScopeHandle.h"
 
 #include <map>
 
+
 static NSString *CKComponentMemoizerThreadKey = @"CKComponentMemoizer";
 
+static bool CKCanUseMemoizedComponent(CKComponent *c) {
+  CKThreadLocalComponentScope *s = CKThreadLocalComponentScope::currentScope();
+
+  NSCAssert(s != nullptr, @"No thread local scope!");
+  if (!s) {
+    return false;
+  }
+
+  auto stackTop = s->stack.top();
+  CKComponentScopeHandleIdentifier oldIdent = c.globalIdentifier;
+
+  // Look for the equivalent child frame
+  auto oldFrameRet = [stackTop.equivalentPreviousFrame childScopeFrameForOldIdentifier:oldIdent];
+  CKComponentScopeFrame *oldScopeFrame = oldFrameRet.frame;
+
+  // Make sure none of the descendents of the oldScopeFrame have pending state updates
+  if (oldScopeFrame && s->stateUpdates.size() > 0) {
+    if ([oldScopeFrame anyDescendentHasPendingUpdates:s->stateUpdates]) {
+      //NSLog(@"Cannot reuse component %@ because state has been updated.", c);
+      return false;
+    }
+  }
+
+  // Copy oldScopeFrame into new scope frame
+  if (oldScopeFrame) {
+    [stackTop.frame injectScopeHierarchy:oldFrameRet];
+  }
+
+  //NSLog(@"Ok to reuse component %@.", c);
+
+  return true;
+}
+
 // Define hash as just pulling out the precomputed hash field
 namespace std {
   template <>
@@ -149,6 +185,11 @@ id CKMemoize(CKMemoizationKey memoizationKey, id (^block)(void))
 {
   _CKComponentMemoizerImpl *impl = [_CKComponentMemoizerImpl currentMemoizer];
   CKComponent *component = [impl dequeueComponentForKey:memoizationKey];
+  if (component) {
+    if (!CKCanUseMemoizedComponent(component)) {
+      component = nil;
+    }
+  }
   if (!component && block) {
     component = block();
   }

ComponentKit/Core/Scope/CKComponentScopeFrame.h

@@ -25,7 +25,7 @@ struct CKComponentScopeFramePair {
 
 + (CKComponentScopeFramePair)childPairForPair:(const CKComponentScopeFramePair &)pair
                                       newRoot:(CKComponentScopeRoot *)newRoot
-                               componentClass:(Class)aClass
+                               componentClass:(Class)componentClass
                                    identifier:(id)identifier
                           initialStateCreator:(id (^)(void))initialStateCreator
                                  stateUpdates:(const CKComponentStateUpdateMap &)stateUpdates;
@@ -34,4 +34,16 @@ struct CKComponentScopeFramePair {
 
 @property (nonatomic, strong, readonly) CKComponentScopeHandle *handle;
 
+struct _CKChildScopeFrameRet {
+  Class cls;
+  id identifier;
+  CKComponentScopeFrame *frame;
+};
+
+- (_CKChildScopeFrameRet)childScopeFrameForOldIdentifier:(CKComponentScopeHandleIdentifier)globalIdentifier;
+
+- (void)injectScopeHierarchy:(_CKChildScopeFrameRet)oldFrame;
+
+- (BOOL)anyDescendentHasPendingUpdates:(const CKComponentStateUpdateMap &)stateUpdates;
+
 @end

ComponentKit/Core/Scope/CKComponentScopeFrame.mm

@@ -23,14 +23,14 @@
 #import "CKMacros.h"
 #import "CKThreadLocalComponentScope.h"
 
-typedef struct _CKStateScopeKey {
+struct _CKStateScopeKey {
   Class __unsafe_unretained componentClass;
   id identifier;
 
   bool operator==(const _CKStateScopeKey &v) const {
     return (CKObjectIsEqual(this->componentClass, v.componentClass) && CKObjectIsEqual(this->identifier, v.identifier));
   }
-} _CKStateScopeKey;
+};
 
 namespace std {
   template <>
@@ -88,4 +88,40 @@ - (instancetype)initWithHandle:(CKComponentScopeHandle *)handle
   return self;
 }
 
+- (void)injectScopeHierarchy:(_CKChildScopeFrameRet)oldFrame
+{
+  Class componentClass = oldFrame.cls;
+  id identifier = oldFrame.identifier;
+  CKAssert([oldFrame.frame isKindOfClass:[CKComponentScopeFrame class]], @"Can't inject with incorrect frame class.");
+
+  auto result = _children.insert({{componentClass, identifier}, oldFrame.frame});
+
+  CKAssert(result.second, @"Scope collision: attempting to create duplicate scope %@:%@", componentClass, identifier);
+}
+
+- (_CKChildScopeFrameRet)childScopeFrameForOldIdentifier:(CKComponentScopeHandleIdentifier)globalIdentifier
+{
+  for (auto &pr : _children) {
+    if (pr.second->_handle.globalIdentifier == globalIdentifier) {
+      return _CKChildScopeFrameRet{pr.first.componentClass, pr.first.identifier, pr.second};
+    }
+  }
+  return {nil, nil, 0};
+}
+
+- (BOOL)anyDescendentHasPendingUpdates:(const CKComponentStateUpdateMap &)stateUpdates
+{
+  auto it = stateUpdates.find(_handle.globalIdentifier);
+  if (it != stateUpdates.end()) {
+    return YES;
+  }
+  // Look in children
+  for (auto &pr : _children) {
+    if ([pr.second anyDescendentHasPendingUpdates:stateUpdates]) {
+      return YES;
+    }
+  }
+  return NO;
+}
+
 @end