Make take_using_weights twice as fast, and 35x for one item

src/prompt_toolkit/layout/containers.py

@@ -448,23 +448,31 @@ def _divide_heights(self, write_position: WritePosition) -> list[int] | None:
 
         i = next(child_generator)
 
+        # The height handed out so far, carried rather than re-added.
+        # These loops run once per unit handed out, and `sum(sizes)`
+        # walked the whole list on every turn to learn what the turn
+        # before it already knew.
+        taken = sum(sizes)
+
         # Increase until we meet at least the 'preferred' size.
         preferred_stop = min(height, sum_dimensions.preferred)
         preferred_dimensions = [d.preferred for d in dimensions]
 
-        while sum(sizes) < preferred_stop:
+        while taken < preferred_stop:
             if sizes[i] < preferred_dimensions[i]:
                 sizes[i] += 1
+                taken += 1
             i = next(child_generator)
 
         # Increase until we use all the available space. (or until "max")
         if not get_app().is_done:
             max_stop = min(height, sum_dimensions.max)
             max_dimensions = [d.max for d in dimensions]
 
-            while sum(sizes) < max_stop:
+            while taken < max_stop:
                 if sizes[i] < max_dimensions[i]:
                     sizes[i] += 1
+                    taken += 1
                 i = next(child_generator)
 
         return sizes
@@ -643,21 +651,30 @@ def _divide_widths(self, width: int) -> list[int] | None:
 
         i = next(child_generator)
 
+        # The width handed out so far, carried rather than re-added.
+        # These loops run about once per column, and `sum(sizes)` walked
+        # the whole list on every turn to learn what the turn before it
+        # already knew.
+        taken = sum(sizes)
+
         # Increase until we meet at least the 'preferred' size.
         preferred_stop = min(width, sum_dimensions.preferred)
+        preferred_dimensions = [d.preferred for d in dimensions]
 
-        while sum(sizes) < preferred_stop:
+        while taken < preferred_stop:
             if sizes[i] < preferred_dimensions[i]:
                 sizes[i] += 1
+                taken += 1
             i = next(child_generator)
 
         # Increase until we use all the available space.
         max_dimensions = [d.max for d in dimensions]
         max_stop = min(width, sum_dimensions.max)
 
-        while sum(sizes) < max_stop:
+        while taken < max_stop:
             if sizes[i] < max_dimensions[i]:
                 sizes[i] += 1
+                taken += 1
             i = next(child_generator)
 
         return sizes

src/prompt_toolkit/utils.py

@@ -249,8 +249,8 @@ def take_using_weights(
     assert len(items) > 0
 
     # Remove items with zero-weight.
-    items2 = []
-    weights2 = []
+    items2: list[_T] = []
+    weights2: list[int] = []
     for item, w in zip(items, weights):
         if w > 0:
             items2.append(item)
@@ -264,8 +264,17 @@ def take_using_weights(
         raise ValueError("Did't got any items with a positive weight.")
 
     #
-    already_taken = [0 for i in items]
     item_count = len(items)
+
+    # One item takes every turn. `VSplit` and `HSplit` ask for a few
+    # thousand items of this per render, and a split of one child is the
+    # common case, so it is worth not building the state below for it.
+    if item_count == 1:
+        item = items[0]
+        while True:
+            yield item
+
+    already_taken = [0] * item_count
     max_weight = max(weights)
 
     i = 0
@@ -275,9 +284,12 @@ def take_using_weights(
         while adding:
             adding = False
 
-            for item_i, item, weight in zip(range(item_count), items, weights):
-                if already_taken[item_i] < i * weight / float(max_weight):
-                    yield item
+            # Reading the two lists by index, rather than building a
+            # `zip` object on every pass of a loop that runs once per
+            # cell handed out. (`range` builds nothing.)
+            for item_i in range(item_count):
+                if already_taken[item_i] < i * weights[item_i] / max_weight:
+                    yield items[item_i]
                     already_taken[item_i] += 1
                     adding = True
 

tests/test_layout.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+from random import Random
+
 import pytest
 
 from prompt_toolkit.layout import InvalidLayoutError, Layout
 from prompt_toolkit.layout.containers import HSplit, VSplit, Window
 from prompt_toolkit.layout.controls import BufferControl
+from prompt_toolkit.layout.dimension import Dimension
 
 
 def test_layout_class():
@@ -51,3 +54,46 @@ def test_layout_class():
 def test_create_invalid_layout():
     with pytest.raises(InvalidLayoutError):
         Layout(HSplit([]))
+
+
+def test_divide_widths_hands_out_exactly_what_fits():
+    """
+    `_divide_widths` grows each child one unit at a time and carries the
+    handed-out total in a local rather than re-adding it. Whatever the
+    loop does, the answer keeps three promises: too little room gives
+    back nothing, every child stays inside its own minimum and maximum,
+    and the sizes add up to everything the split may take. A local that
+    drifted from the true total in either direction would break the
+    third promise, so this pins the carried total as well as the loop.
+    """
+    rng = Random(0)
+
+    for _ in range(200):
+        children = []
+        for _ in range(rng.randint(1, 4)):
+            minimum = rng.randint(0, 5)
+            preferred = minimum + rng.randint(0, 10)
+            children.append(
+                Window(
+                    width=Dimension(
+                        min=minimum,
+                        preferred=preferred,
+                        max=preferred + rng.randint(0, 10),
+                        weight=rng.randint(1, 3),
+                    )
+                )
+            )
+
+        width = rng.randint(0, 30)
+        split = VSplit(children)
+        actual = [c.preferred_width(width) for c in split._all_children]
+        sizes = split._divide_widths(width)
+
+        if sum(d.min for d in actual) > width:
+            assert sizes is None
+            continue
+
+        assert sizes is not None
+        assert len(sizes) == len(actual)
+        assert all(d.min <= s <= d.max for s, d in zip(sizes, actual))
+        assert sum(sizes) == min(width, sum(d.max for d in actual))

tests/test_utils.py

@@ -76,3 +76,29 @@ def take(generator, count):
     # All zero-weight items.
     with pytest.raises(ValueError):
         take(take_using_weights(["A", "B", "C"], [0, 0, 0]), 70)
+
+
+def test_using_weights_yields_one_exact_order():
+    """
+    The order, and not only the proportions.
+
+    `VSplit` and `HSplit` hand out one cell at a time in this order, so
+    two orders with the same proportions still give two layouts: they
+    differ by which child gets the last cell.
+    """
+    for weights, expected in [
+        ([1], "AAAAAAAAAAAA"),
+        ([7], "AAAAAAAAAAAA"),
+        ([1, 1], "ABABABABABABABAB"),
+        ([1, 0], "AAAAAAAAAAAA"),
+        ([5, 10], "ABBABBABBABBABBABBABBABB"),
+        ([1, 1, 1], "ABCABCABCABCABCABCABCABC"),
+        ([3, 1, 2], "ABCACAABCACAABCACAABCACA"),
+        ([20, 10, 5], "ABCAABAABCAABAABCAABAABCAABAABCAABA"),
+        ([1, 2, 3, 4], "ABCDCDBCDDABCDCDBCDDABCDCDBCDD"),
+        ([0, 2, 0, 4], "BDDBDDBDDBDDBDDBDDBDDBDD"),
+    ]:
+        items = list("ABCDEFGH"[: len(weights)])
+        taken = itertools.islice(take_using_weights(items, weights), len(expected))
+
+        assert "".join(taken) == expected, weights