fix: stop cp_sat from putting every file in its own PR

pr_split/planner/partitioning.py

@@ -388,7 +388,11 @@ def _group_units_cp_sat(
     if not units:
         return []
 
-    group_slots = len(units)
+    # The graph plan is a good starting point: it seeds the search and bounds
+    # how many groups the model needs room for. Starting from all singletons,
+    # the solver could stall on a plan with one tiny PR per file.
+    hint = _group_units_graph(units, settings=settings)
+    group_slots = min(len(units), len(hint) + max(2, len(hint) // 4))
     max_total_loc = max(settings.max_loc, sum(unit.loc for unit in units), 1)
 
     model = cp_model.CpModel()
@@ -430,28 +434,6 @@ def _group_units_cp_sat(
     for group_idx in range(group_slots - 1):
         model.Add(y[group_idx] >= y[group_idx + 1])
 
-    pair_terms: list[tuple[int, int, IntVar]] = []
-    for left in range(len(units)):
-        for right in range(left + 1, len(units)):
-            affinity = _affinity_score(units[left], units[right], settings.priority)
-            same_group_terms = []
-            for group_idx in range(group_slots):
-                pair_var = model.NewBoolVar(f"pair_{left}_{right}_{group_idx}")
-                model.Add(pair_var <= x[(left, group_idx)])
-                model.Add(pair_var <= x[(right, group_idx)])
-                model.Add(pair_var >= x[(left, group_idx)] + x[(right, group_idx)] - 1)
-                same_group_terms.append(pair_var)
-            same_group = model.NewBoolVar(f"same_{left}_{right}")
-            model.Add(sum(same_group_terms) == same_group)
-            is_cross_file = int(units[left].file_path != units[right].file_path)
-            pair_terms.append((affinity, is_cross_file, same_group))
-
-    overflow_weight = _CP_SAT_OVERFLOW_WEIGHT
-    group_weight = (
-        _CP_SAT_GROUP_WEIGHT_LOGICAL
-        if settings.priority == Priority.LOGICAL
-        else _CP_SAT_GROUP_WEIGHT_ORTHOGONAL
-    )
     cross_file_penalty = (
         _CP_SAT_CROSS_FILE_PENALTY_LOGICAL
         if settings.priority == Priority.LOGICAL
@@ -462,19 +444,54 @@ def _group_units_cp_sat(
         if settings.priority == Priority.LOGICAL
         else _CP_SAT_AFFINITY_DIVISOR_ORTHOGONAL
     )
+    # One variable per pair that the objective cares about: rewarded when two
+    # related units share a group, penalised when two unrelated files do.
+    # Only unrelated files are kept apart; related ones (same directory,
+    # shared symbols, a test and its module) may share a PR.
+    rewards: list[tuple[int, IntVar]] = []
+    penalties: list[IntVar] = []
+    for left in range(len(units)):
+        for right in range(left + 1, len(units)):
+            affinity = _affinity_score(units[left], units[right], settings.priority)
+            reward = affinity // affinity_divisor if affinity > 0 else 0
+            penalised = (
+                cross_file_penalty > 0
+                and affinity <= 0
+                and units[left].file_path != units[right].file_path
+            )
+            if not reward and not penalised:
+                continue
+            same_group = model.NewBoolVar(f"same_{left}_{right}")
+            for group_idx in range(group_slots):
+                a, b = x[(left, group_idx)], x[(right, group_idx)]
+                if penalised:
+                    # Minimising pushes it down; it must be 1 when both are here.
+                    model.Add(same_group >= a + b - 1)
+                if reward:
+                    # Maximising pushes it up; it must be 0 unless both are here.
+                    model.Add(same_group <= 1 - a + b)
+            if penalised:
+                penalties.append(same_group)
+            if reward:
+                rewards.append((reward, same_group))
+
+    for group_idx, group_units in enumerate(hint[:group_slots]):
+        members = {unit.id for unit in group_units}
+        for unit_idx, unit in enumerate(units):
+            model.AddHint(x[(unit_idx, group_idx)], unit.id in members)
 
+    group_weight = (
+        _CP_SAT_GROUP_WEIGHT_LOGICAL
+        if settings.priority == Priority.LOGICAL
+        else _CP_SAT_GROUP_WEIGHT_ORTHOGONAL
+    )
     underflow_cost = _CP_SAT_UNDERFLOW_WEIGHT * sum(underflow) if underflow else 0
     model.Minimize(
         group_weight * sum(y)
-        + overflow_weight * sum(overflow)
+        + _CP_SAT_OVERFLOW_WEIGHT * sum(overflow)
         + underflow_cost
-        + cross_file_penalty
-        * sum(same_group for _, is_cross_file, same_group in pair_terms if is_cross_file)
-        - sum(
-            (affinity // affinity_divisor) * same_group
-            for affinity, _, same_group in pair_terms
-            if affinity > 0
-        )
+        + cross_file_penalty * sum(penalties)
+        - sum(reward * same_group for reward, same_group in rewards)
     )
 
     solver = cp_model.CpSolver()
@@ -501,6 +518,8 @@ def _group_units_cp_sat(
         for _, group_units in sorted(assignments.items())
         if group_units
     ]
+    # Slot order is arbitrary; number the groups in diff order instead.
+    grouped.sort(key=lambda group_units: group_units[0].position)
     if status == cp_model.FEASIBLE:
         # The solver ran out of time before proving optimality; on larger
         # diffs the plan it returns can be far from the best one (many

pr_split/planner/symbols.py

@@ -222,7 +222,7 @@ def symbol_dependencies(
     owner = {name: next(iter(by.items())) for name, by in defined_by.items() if len(by) == 1}
     # A new field is only useful once something fills it: readers also need
     # every group that writes it (``name=`` or ``"name":``).
-    fields = {name for name in owner if name in _new_fields(added)}
+    fields = set(owner) & _new_fields(added)
     writers: dict[str, set[str]] = defaultdict(set)
     for group in groups:
         for path, lines in added.get(group.id, {}).items():

tests/test_partitioning_extensive.py

@@ -146,6 +146,39 @@ def test_orthogonal_keeps_unrelated_files_separate(
         assert all(len(group.assignments) == 1 for group in groups)
         _assert_valid_plan(groups, UNRELATED_DIFF, 10)
 
+    def test_orthogonal_keeps_related_files_together(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        # Two small files in one package are related: orthogonal keeps only
+        # unrelated files apart, so they share a PR rather than making two tiny ones.
+        pytest.importorskip("ortools.sat.python.cp_model")
+        diff = """\
+diff --git a/pkg/alpha.py b/pkg/alpha.py
+new file mode 100644
+--- /dev/null
++++ b/pkg/alpha.py
+@@ -0,0 +1,2 @@
++ALPHA = 1
++BETA = 2
+diff --git a/pkg/gamma.py b/pkg/gamma.py
+new file mode 100644
+--- /dev/null
++++ b/pkg/gamma.py
+@@ -0,0 +1,2 @@
++GAMMA = 3
++DELTA = 4
+"""
+        settings = _settings(
+            monkeypatch,
+            max_loc=10,
+            partition_strategy=PartitionStrategy.CP_SAT,
+            priority=Priority.ORTHOGONAL,
+        )
+        groups = partition_diff(parse_diff(diff), settings)
+        assert len(groups) == 1
+        assert len(groups[0].assignments) == 2
+        _assert_valid_plan(groups, diff, 10)
+
     def test_logical_groups_related_files_together(self, monkeypatch: pytest.MonkeyPatch) -> None:
         settings = _settings(
             monkeypatch,