fix: stabilize lane assignments in masonry layout (#1080)

Author: seob717

.changeset/stable-lane-assignments.md

@@ -0,0 +1,14 @@
+---
+'@tanstack/virtual-core': patch
+---
+
+fix: stabilize lane assignments in masonry layout
+
+Added lane assignment caching to prevent items from jumping between lanes when viewport is resized. Previously, items could shift to different lanes during resize due to recalculating "shortest lane" with slightly different heights.
+
+Changes:
+
+- Added `laneAssignments` cache (Map<index, lane>) to persist lane assignments
+- Lane cache is cleared when `lanes` option changes or `measure()` is called
+- Lane cache is cleaned up when `count` decreases (removes stale entries)
+- Lane cache is cleared when virtualizer is disabled

packages/virtual-core/src/index.ts

@@ -361,7 +361,11 @@ export class Virtualizer<
   isScrolling = false
   measurementsCache: Array<VirtualItem> = []
   private itemSizeCache = new Map<Key, number>()
+  private laneAssignments = new Map<number, number>() // index → lane cache
   private pendingMeasuredCacheIndexes: Array<number> = []
+  private prevLanes: number | undefined = undefined
+  private lanesChangedFlag = false
+  private lanesSettling = false
   scrollRect: Rect | null = null
   scrollOffset: number | null = null
   scrollDirection: ScrollDirection | null = null
@@ -617,15 +621,27 @@ export class Virtualizer<
       this.options.scrollMargin,
       this.options.getItemKey,
       this.options.enabled,
+      this.options.lanes,
     ],
-    (count, paddingStart, scrollMargin, getItemKey, enabled) => {
+    (count, paddingStart, scrollMargin, getItemKey, enabled, lanes) => {
+      const lanesChanged =
+        this.prevLanes !== undefined && this.prevLanes !== lanes
+
+      if (lanesChanged) {
+        // Set flag for getMeasurements to handle
+        this.lanesChangedFlag = true
+      }
+
+      this.prevLanes = lanes
       this.pendingMeasuredCacheIndexes = []
+
       return {
         count,
         paddingStart,
         scrollMargin,
         getItemKey,
         enabled,
+        lanes,
       }
     },
     {
@@ -636,41 +652,108 @@ export class Virtualizer<
   private getMeasurements = memo(
     () => [this.getMeasurementOptions(), this.itemSizeCache],
     (
-      { count, paddingStart, scrollMargin, getItemKey, enabled },
+      { count, paddingStart, scrollMargin, getItemKey, enabled, lanes },
       itemSizeCache,
     ) => {
       if (!enabled) {
         this.measurementsCache = []
         this.itemSizeCache.clear()
+        this.laneAssignments.clear()
         return []
       }
 
+      // Clean up stale lane cache entries when count decreases
+      if (this.laneAssignments.size > count) {
+        for (const index of this.laneAssignments.keys()) {
+          if (index >= count) {
+            this.laneAssignments.delete(index)
+          }
+        }
+      }
+
+      // ✅ Force complete recalculation when lanes change
+      if (this.lanesChangedFlag) {
+        this.lanesChangedFlag = false // Reset immediately
+        this.lanesSettling = true // Start settling period
+        this.measurementsCache = []
+        this.itemSizeCache.clear()
+        this.laneAssignments.clear() // Clear lane cache for new lane count
+        // Clear pending indexes to force min = 0
+        this.pendingMeasuredCacheIndexes = []
+      }
+
       if (this.measurementsCache.length === 0) {
         this.measurementsCache = this.options.initialMeasurementsCache
         this.measurementsCache.forEach((item) => {
           this.itemSizeCache.set(item.key, item.size)
         })
       }
 
-      const min =
-        this.pendingMeasuredCacheIndexes.length > 0
+      // ✅ During lanes settling, ignore pendingMeasuredCacheIndexes to prevent repositioning
+      const min = this.lanesSettling
+        ? 0
+        : this.pendingMeasuredCacheIndexes.length > 0
           ? Math.min(...this.pendingMeasuredCacheIndexes)
           : 0
       this.pendingMeasuredCacheIndexes = []
 
+      // ✅ End settling period when cache is fully built
+      if (this.lanesSettling && this.measurementsCache.length === count) {
+        this.lanesSettling = false
+      }
+
       const measurements = this.measurementsCache.slice(0, min)
 
+      // ✅ Performance: Track last item index per lane for O(1) lookup
+      const laneLastIndex: Array<number | undefined> = new Array(lanes).fill(
+        undefined,
+      )
+
+      // Initialize from existing measurements (before min)
+      for (let m = 0; m < min; m++) {
+        const item = measurements[m]
+        if (item) {
+          laneLastIndex[item.lane] = m
+        }
+      }
+
       for (let i = min; i < count; i++) {
         const key = getItemKey(i)
 
-        const furthestMeasurement =
-          this.options.lanes === 1
-            ? measurements[i - 1]
-            : this.getFurthestMeasurement(measurements, i)
-
-        const start = furthestMeasurement
-          ? furthestMeasurement.end + this.options.gap
-          : paddingStart + scrollMargin
+        // Check for cached lane assignment
+        const cachedLane = this.laneAssignments.get(i)
+        let lane: number
+        let start: number
+
+        if (cachedLane !== undefined && this.options.lanes > 1) {
+          // Use cached lane - O(1) lookup for previous item in same lane
+          lane = cachedLane
+          const prevIndex = laneLastIndex[lane]
+          const prevInLane =
+            prevIndex !== undefined ? measurements[prevIndex] : undefined
+          start = prevInLane
+            ? prevInLane.end + this.options.gap
+            : paddingStart + scrollMargin
+        } else {
+          // No cache - use original logic (find shortest lane)
+          const furthestMeasurement =
+            this.options.lanes === 1
+              ? measurements[i - 1]
+              : this.getFurthestMeasurement(measurements, i)
+
+          start = furthestMeasurement
+            ? furthestMeasurement.end + this.options.gap
+            : paddingStart + scrollMargin
+
+          lane = furthestMeasurement
+            ? furthestMeasurement.lane
+            : i % this.options.lanes
+
+          // Cache the lane assignment
+          if (this.options.lanes > 1) {
+            this.laneAssignments.set(i, lane)
+          }
+        }
 
         const measuredSize = itemSizeCache.get(key)
         const size =
@@ -680,10 +763,6 @@ export class Virtualizer<
 
         const end = start + size
 
-        const lane = furthestMeasurement
-          ? furthestMeasurement.lane
-          : i % this.options.lanes
-
         measurements[i] = {
           index: i,
           start,
@@ -692,6 +771,9 @@ export class Virtualizer<
           key,
           lane,
         }
+
+        // ✅ Performance: Update lane's last item index
+        laneLastIndex[lane] = i
       }
 
       this.measurementsCache = measurements
@@ -1077,6 +1159,7 @@ export class Virtualizer<
 
   measure = () => {
     this.itemSizeCache = new Map()
+    this.laneAssignments = new Map() // Clear lane cache for full re-layout
     this.notify(false)
   }
 }