ListLayout cache invalidation fixes

packages/@react-aria/virtualizer/src/VirtualizerItem.tsx

@@ -43,7 +43,16 @@ export function layoutInfoToStyle(layoutInfo: LayoutInfo, dir: Direction, parent
   let xProperty = dir === 'rtl' ? 'right' : 'left';
   let cached = cache.get(layoutInfo);
   if (cached && cached[xProperty] != null) {
-    return cached;
+    if (!parent) {
+      return cached;
+    }
+
+    // Invalidate if the parent position changed.
+    let top = layoutInfo.rect.y - parent.rect.y;
+    let x = layoutInfo.rect.x - parent.rect.x;
+    if (cached.top === top && cached[xProperty] === x) {
+      return cached;
+    }
   }
 
   let style: CSSProperties = {
@@ -55,8 +64,8 @@ export function layoutInfoToStyle(layoutInfo: LayoutInfo, dir: Direction, parent
     WebkitTransition: 'all',
     WebkitTransitionDuration: 'inherit',
     transitionDuration: 'inherit',
-    width: layoutInfo.rect.width + 'px',
-    height: layoutInfo.rect.height + 'px',
+    width: layoutInfo.rect.width,
+    height: layoutInfo.rect.height,
     opacity: layoutInfo.opacity,
     zIndex: layoutInfo.zIndex,
     transform: layoutInfo.transform,

packages/@react-stately/layout/src/ListLayout.ts

@@ -28,6 +28,7 @@ type ListLayoutOptions<T> = {
 
 // A wrapper around LayoutInfo that supports heirarchy
 export interface LayoutNode {
+  node?: Node<unknown>,
   layoutInfo: LayoutInfo,
   header?: LayoutInfo,
   children?: LayoutNode[]
@@ -62,8 +63,7 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate {
   protected lastCollection: Collection<Node<T>>;
   protected rootNodes: LayoutNode[];
   protected collator: Intl.Collator;
-  protected cache: WeakMap<Node<T>, LayoutNode> = new WeakMap();
-  protected newCache: WeakMap<Node<T>, LayoutNode> = new WeakMap();
+  protected invalidateEverything: boolean;
 
   /**
    * Creates a new ListLayout with options. See the list of properties below for a description
@@ -116,22 +116,13 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate {
   }
 
   validate(invalidationContext: InvalidationContext<Node<T>, unknown>) {
-    // Initialize new cache
-    this.newCache = new WeakMap();
-
     // Invalidate cache if the size of the collection changed.
     // In this case, we need to recalculate the entire layout.
-    if (invalidationContext.sizeChanged) {
-      this.cache = new WeakMap();
-    }
+    this.invalidateEverything = invalidationContext.sizeChanged;
 
     this.collection = this.virtualizer.collection;
     this.rootNodes = this.buildCollection();
 
-    // Replace old cache with new cache so nodes that are no longer in the collection aren't cached
-    // (for cases like combobox where original collection sticks around for filtering)
-    this.cache = this.newCache;
-
     this.lastWidth = this.virtualizer.visibleRect.width;
     this.lastCollection = this.collection;
   }
@@ -158,12 +149,13 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate {
   }
 
   buildChild(node: Node<T>, x: number, y: number): LayoutNode {
-    let cached = this.cache.get(node);
-    if (cached && y === (cached.header || cached.layoutInfo).rect.y) {
+    let cached = this.layoutNodes.get(node.key);
+    if (!this.invalidateEverything && cached && cached.node === node && y === (cached.header || cached.layoutInfo).rect.y) {
       return cached;
     }
 
     let layoutNode = this.buildNode(node, x, y);
+    layoutNode.node = node;
 
     layoutNode.layoutInfo.parentKey = node.parentKey || null;
     this.layoutInfos.set(layoutNode.layoutInfo.key, layoutNode.layoutInfo);
@@ -172,17 +164,16 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate {
     }
 
     // Remove deleted child layout nodes from key mapping.
-    let prev = this.layoutNodes.get(node.key);
-    if (prev) {
+    if (cached) {
       let childKeys = new Set();
       if (layoutNode.children) {
         for (let child of layoutNode.children) {
           childKeys.add(child.layoutInfo.key);
         }
       }
 
-      if (prev.children) {
-        for (let child of prev.children) {
+      if (cached.children) {
+        for (let child of cached.children) {
           if (!childKeys.has(child.layoutInfo.key)) {
             this.removeLayoutNode(child);
           }
@@ -191,10 +182,6 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate {
     }
 
     this.layoutNodes.set(node.key, layoutNode);
-
-    // Update new cache rather than old cache so that we can get rid of cached nodes that don't exist in collection any more
-    this.newCache.set(node, layoutNode);
-
     return layoutNode;
   }
 
@@ -318,14 +305,17 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate {
     let layoutInfo = this.layoutInfos.get(key);
     layoutInfo.estimatedSize = false;
     if (layoutInfo.rect.height !== size.height) {
-      layoutInfo.rect.height = size.height;
+      // Copy layout info rather than mutating so that later caches are invalidated.
+      let newLayoutInfo = layoutInfo.copy();
+      newLayoutInfo.rect.height = size.height;
+      this.layoutInfos.set(key, newLayoutInfo);
 
       // Invalidate layout for this layout node and all parents
-      this.cache.delete(this.collection.getItem(key));
+      this.updateLayoutNode(key, layoutInfo, newLayoutInfo);
 
       let node = this.collection.getItem(layoutInfo.parentKey);
       while (node) {
-        this.cache.delete(this.collection.getItem(node.key));
+        this.updateLayoutNode(node.key, layoutInfo, newLayoutInfo);
         node = this.collection.getItem(node.parentKey);
       }
 
@@ -335,6 +325,21 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate {
     return false;
   }
 
+  updateLayoutNode(key: Key, oldLayoutInfo: LayoutInfo, newLayoutInfo: LayoutInfo) {
+    let n = this.layoutNodes.get(key);
+    if (n) {
+      // Invalidate by clearing node.
+      n.node = null;
+
+      // Replace layout info in LayoutNode
+      if (n.header === oldLayoutInfo) {
+        n.header = newLayoutInfo;
+      } else if (n.layoutInfo === oldLayoutInfo) {
+        n.layoutInfo = newLayoutInfo;
+      }
+    }
+  }
+
   getContentSize() {
     return this.contentSize;
   }

packages/@react-stately/layout/src/TableLayout.ts

@@ -28,7 +28,8 @@ export class TableLayout<T> extends ListLayout<T> {
       this.collection.columns.length !== this.lastCollection.columns.length ||
       this.collection.columns.some((c, i) => c.key !== this.lastCollection.columns[i].key)
     ) {
-      this.cache = new WeakMap();
+      // Invalidate everything in this layout pass. Will be reset in ListLayout on the next pass.
+      this.invalidateEverything = true;
     }
 
     this.buildColumnWidths();