diff --git a/pr_split/planner/partitioning.py b/pr_split/planner/partitioning.py index 3fdd974..7c899f8 100644 --- a/pr_split/planner/partitioning.py +++ b/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 diff --git a/pr_split/planner/symbols.py b/pr_split/planner/symbols.py index df82c2b..964a672 100644 --- a/pr_split/planner/symbols.py +++ b/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(): diff --git a/tests/test_partitioning_extensive.py b/tests/test_partitioning_extensive.py index 129bf60..86977c3 100644 --- a/tests/test_partitioning_extensive.py +++ b/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,