[Refactor] Extract widget layout logic (#2545)

src/scripts/domWidget.ts

@@ -7,6 +7,7 @@ import type {
 } from '@comfyorg/litegraph/dist/types/widgets'
 
 import { useSettingStore } from '@/stores/settingStore'
+import { distributeSpace } from '@/utils/spaceDistribution'
 
 import { ANIM_PREVIEW_WIDGET, app } from './app'
 
@@ -21,7 +22,7 @@ interface Rect {
 
 interface DOMSizeInfo {
   minHeight: number
-  prefHeight: number
+  prefHeight?: number
   w: DOMWidget<HTMLElement, object>
   diff?: number
 }
@@ -138,8 +139,10 @@ function computeSize(this: LGraphNode, size: Size): void {
   let y = this.widgets[0].last_y
   let freeSpace = size[1] - y
 
-  let widgetHeight = 0
-  let dom: DOMSizeInfo[] = []
+  // Collect fixed height widgets first
+  let fixedWidgetHeight = 0
+  const domWidgets: DOMSizeInfo[] = []
+
   for (const w of this.widgets) {
     // @ts-expect-error custom widget type
     if (w.type === 'converted-widget') {
@@ -147,7 +150,7 @@ function computeSize(this: LGraphNode, size: Size): void {
       // @ts-expect-error custom widget type
       delete w.computedHeight
     } else if (w.computeSize) {
-      widgetHeight += w.computeSize()[1] + 4
+      fixedWidgetHeight += w.computeSize()[1] + 4
     } else if (isDomWidget(w)) {
       // Extract DOM widget size info
       const styles = getComputedStyle(w.element)
@@ -174,89 +177,50 @@ function computeSize(this: LGraphNode, size: Size): void {
           minHeight = prefHeight
         }
       }
-      if (isNaN(minHeight)) {
-        minHeight = 50
-      }
-      if (!isNaN(maxHeight)) {
-        if (!isNaN(prefHeight)) {
-          prefHeight = Math.min(prefHeight, maxHeight)
-        } else {
-          prefHeight = maxHeight
-        }
-      }
-      dom.push({
-        minHeight,
-        prefHeight,
+
+      domWidgets.push({
+        minHeight: isNaN(minHeight) ? 50 : minHeight,
+        prefHeight: isNaN(prefHeight)
+          ? undefined
+          : Math.min(prefHeight, maxHeight ?? Infinity),
         w
       })
     } else {
-      widgetHeight += LiteGraph.NODE_WIDGET_HEIGHT + 4
-    }
-  }
-
-  freeSpace -= widgetHeight
-
-  // Calculate sizes with all widgets at their min height
-  const prefGrow = [] // Nodes that want to grow to their prefd size
-  const canGrow = [] // Nodes that can grow to auto size
-  let growBy = 0
-  for (const d of dom) {
-    freeSpace -= d.minHeight
-    if (isNaN(d.prefHeight)) {
-      canGrow.push(d)
-      d.w.computedHeight = d.minHeight
-    } else {
-      const diff = d.prefHeight - d.minHeight
-      if (diff > 0) {
-        prefGrow.push(d)
-        growBy += diff
-        d.diff = diff
-      } else {
-        d.w.computedHeight = d.minHeight
-      }
+      fixedWidgetHeight += LiteGraph.NODE_WIDGET_HEIGHT + 4
     }
   }
 
   if (this.imgs && !this.widgets?.find((w) => w.name === ANIM_PREVIEW_WIDGET)) {
-    freeSpace -= 220
+    fixedWidgetHeight += 220
   }
 
+  // Calculate remaining space for DOM widgets
+  freeSpace -= fixedWidgetHeight
   this.freeWidgetSpace = freeSpace
 
-  if (freeSpace < 0) {
-    // Not enough space for all widgets so we need to grow
-    size[1] -= freeSpace
-    this.graph?.setDirtyCanvas(true)
-  } else {
-    // Share the space between each
-    const growDiff = freeSpace - growBy
-    if (growDiff > 0) {
-      // All pref sizes can be fulfilled
-      freeSpace = growDiff
-      for (const d of prefGrow) {
-        d.w.computedHeight = d.prefHeight
-      }
-    } else {
-      // We need to grow evenly
-      const shared = -growDiff / prefGrow.length
-      for (const d of prefGrow) {
-        d.w.computedHeight = d.prefHeight - shared
-      }
-      freeSpace = 0
-    }
+  // Prepare space requests for distribution
+  const spaceRequests = domWidgets.map((d) => ({
+    minSize: d.minHeight,
+    maxSize: d.prefHeight
+  }))
 
-    if (freeSpace > 0 && canGrow.length) {
-      // Grow any that are auto height
-      const shared = freeSpace / canGrow.length
-      for (const d of canGrow) {
-        if (d.w.computedHeight) {
-          d.w.computedHeight += shared
-        }
-      }
-    }
+  // Distribute space among DOM widgets
+  const allocations = distributeSpace(Math.max(0, freeSpace), spaceRequests)
+
+  // Apply computed heights
+  domWidgets.forEach((d, i) => {
+    d.w.computedHeight = allocations[i]
+  })
+
+  // If we need more space, grow the node
+  const totalNeeded =
+    fixedWidgetHeight + allocations.reduce((sum, h) => sum + h, 0)
+  if (totalNeeded > size[1] - this.widgets[0].last_y) {
+    size[1] = totalNeeded + this.widgets[0].last_y
+    this.graph?.setDirtyCanvas(true)
   }
 
-  // Position each of the widgets
+  // Position widgets
   for (const w of this.widgets) {
     w.y = y
     if (w.computedHeight) {

src/utils/spaceDistribution.ts

@@ -0,0 +1,77 @@
+export interface SpaceRequest {
+  minSize: number
+  maxSize?: number
+}
+
+/**
+ * Distributes available space among items with min/max size constraints
+ * @param totalSpace Total space available to distribute
+ * @param requests Array of space requests with size constraints
+ * @returns Array of space allocations
+ */
+export function distributeSpace(
+  totalSpace: number,
+  requests: SpaceRequest[]
+): number[] {
+  // Handle edge cases
+  if (requests.length === 0) return []
+
+  // Calculate total minimum space needed
+  const totalMinSize = requests.reduce((sum, req) => sum + req.minSize, 0)
+
+  // If we can't meet minimum requirements, return the minimum sizes
+  if (totalSpace < totalMinSize) {
+    return requests.map((req) => req.minSize)
+  }
+
+  // Initialize allocations with minimum sizes
+  let allocations = requests.map((req) => ({
+    computedSize: req.minSize,
+    maxSize: req.maxSize ?? Infinity,
+    remaining: (req.maxSize ?? Infinity) - req.minSize
+  }))
+
+  // Calculate remaining space to distribute
+  let remainingSpace = totalSpace - totalMinSize
+
+  // Distribute remaining space iteratively
+  while (
+    remainingSpace > 0 &&
+    allocations.some((alloc) => alloc.remaining > 0)
+  ) {
+    // Count items that can still grow
+    const growableItems = allocations.filter(
+      (alloc) => alloc.remaining > 0
+    ).length
+
+    if (growableItems === 0) break
+
+    // Calculate fair share per item
+    const sharePerItem = remainingSpace / growableItems
+
+    // Track how much space was actually used in this iteration
+    let spaceUsedThisRound = 0
+
+    // Distribute space
+    allocations = allocations.map((alloc) => {
+      if (alloc.remaining <= 0) return alloc
+
+      const growth = Math.min(sharePerItem, alloc.remaining)
+      spaceUsedThisRound += growth
+
+      return {
+        ...alloc,
+        computedSize: alloc.computedSize + growth,
+        remaining: alloc.remaining - growth
+      }
+    })
+
+    remainingSpace -= spaceUsedThisRound
+
+    // Break if we couldn't distribute any more space
+    if (spaceUsedThisRound === 0) break
+  }
+
+  // Return only the computed sizes
+  return allocations.map(({ computedSize }) => computedSize)
+}

tests-ui/tests/utils/spaceDistrbution.test.ts

@@ -0,0 +1,39 @@
+import { type SpaceRequest, distributeSpace } from '@/utils/spaceDistribution'
+
+describe('distributeSpace', () => {
+  it('should distribute space according to minimum sizes when space is limited', () => {
+    const requests: SpaceRequest[] = [
+      { minSize: 100 },
+      { minSize: 100 },
+      { minSize: 100 }
+    ]
+    expect(distributeSpace(300, requests)).toEqual([100, 100, 100])
+  })
+
+  it('should distribute extra space equally when no maxSize', () => {
+    const requests: SpaceRequest[] = [{ minSize: 100 }, { minSize: 100 }]
+    expect(distributeSpace(400, requests)).toEqual([200, 200])
+  })
+
+  it('should respect maximum sizes', () => {
+    const requests: SpaceRequest[] = [
+      { minSize: 100, maxSize: 150 },
+      { minSize: 100 }
+    ]
+    expect(distributeSpace(400, requests)).toEqual([150, 250])
+  })
+
+  it('should handle empty requests array', () => {
+    expect(distributeSpace(1000, [])).toEqual([])
+  })
+
+  it('should handle negative total space', () => {
+    const requests: SpaceRequest[] = [{ minSize: 100 }, { minSize: 100 }]
+    expect(distributeSpace(-100, requests)).toEqual([100, 100])
+  })
+
+  it('should handle total space smaller than minimum sizes', () => {
+    const requests: SpaceRequest[] = [{ minSize: 100 }, { minSize: 100 }]
+    expect(distributeSpace(100, requests)).toEqual([100, 100])
+  })
+})