From d02d88d3d83ed5c558c5c7d06381ea0ee555460d Mon Sep 17 00:00:00 2001 From: vitali87 Date: Sun, 30 Aug 2026 13:45:50 +0100 Subject: [PATCH 1/4] fix: keep one assignment per file when merging chunk groups When a later diff chunk assigned more hunks of a file to a group that already held that file, _merge_chunk_groups appended a second GroupAssignment for the same path. validate_coverage accepted it (each hunk claimed once) but materialize_group_files wrote the file once per assignment, so the last assignment overwrote the first and the earlier hunks were silently missing from the committed branch. - _merge_chunk_groups merges same-path assignments (union of hunks, WHOLE_FILE wins). - materialize_group_files groups assignments by path and writes each file once from the union of their hunks, so duplicates from any source can no longer drop changes. --- pr_split/diff_ops/reconstructor.py | 42 +++++++++------- pr_split/planner/client.py | 28 ++++++++++- tests/test_client.py | 77 ++++++++++++++++++++++++++++++ tests/test_reconstructor.py | 54 +++++++++++++++++++++ 4 files changed, 183 insertions(+), 18 deletions(-) diff --git a/pr_split/diff_ops/reconstructor.py b/pr_split/diff_ops/reconstructor.py index 2132da3..3b70c34 100644 --- a/pr_split/diff_ops/reconstructor.py +++ b/pr_split/diff_ops/reconstructor.py @@ -83,26 +83,39 @@ def apply_hunks(base_content: str, patch_file: PatchedFile, assigned_indices: li return "".join(lines) +def _assigned_hunk_indices( + patch_file: PatchedFile, assignments: list[GroupAssignment] +) -> list[int]: + """Union of the hunks every assignment for one file claims.""" + covered: set[int] = set() + for assignment in assignments: + if assignment.assignment_type is AssignmentType.WHOLE_FILE: + covered.update(range(len(patch_file))) + else: + covered.update(assignment.hunk_indices) + return sorted(covered) + + def materialize_group_files( parsed_diff: ParsedDiff, group: Group, ref: str ) -> dict[str, str | None]: logger.info(logs.MATERIALIZING_FILES.format(count=len(group.assignments), group=group.id)) pf_map = {pf.path: pf for pf in parsed_diff.patch_set} - result: dict[str, str | None] = {} + # Several assignments may name the same file (e.g. merged across diff + # chunks); each file is written once from the union of their hunks. + assignments_by_path: dict[str, list[GroupAssignment]] = {} for assignment in group.assignments: - patch_file = pf_map.get(assignment.file_path) + assignments_by_path.setdefault(assignment.file_path, []).append(assignment) + result: dict[str, str | None] = {} + for file_path, assignments in assignments_by_path.items(): + patch_file = pf_map.get(file_path) if patch_file is None: continue if patch_file.is_removed_file: - result[assignment.file_path] = None + result[file_path] = None continue + indices = _assigned_hunk_indices(patch_file, assignments) if patch_file.is_added_file: - all_indices = list(range(len(patch_file))) - match assignment.assignment_type: - case AssignmentType.WHOLE_FILE: - indices = all_indices - case AssignmentType.PARTIAL_HUNKS: - indices = assignment.hunk_indices target_lines = [] for idx in indices: hunk = patch_file[idx] @@ -111,15 +124,10 @@ def materialize_group_files( target_lines.append(str(line)[1:]) # Lines from unidiff keep their trailing newline, so they are # concatenated as-is; joining on "\n" double-spaces the file. - result[assignment.file_path] = "".join( + result[file_path] = "".join( ln if ln.endswith("\n") else ln + "\n" for ln in target_lines ) continue - base_content = _get_base_file_content(assignment.file_path, ref) - match assignment.assignment_type: - case AssignmentType.WHOLE_FILE: - indices = list(range(len(patch_file))) - case AssignmentType.PARTIAL_HUNKS: - indices = assignment.hunk_indices - result[assignment.file_path] = apply_hunks(base_content, patch_file, indices) + base_content = _get_base_file_content(file_path, ref) + result[file_path] = apply_hunks(base_content, patch_file, indices) return result diff --git a/pr_split/planner/client.py b/pr_split/planner/client.py index 0c933c8..2c2b0d6 100644 --- a/pr_split/planner/client.py +++ b/pr_split/planner/client.py @@ -227,12 +227,38 @@ def _parse_groups(raw: RawToolOutput) -> list[Group]: return groups +def _merge_assignment(existing: GroupAssignment, incoming: GroupAssignment) -> GroupAssignment: + if ( + existing.assignment_type is AssignmentType.WHOLE_FILE + or incoming.assignment_type is AssignmentType.WHOLE_FILE + ): + assignment_type = AssignmentType.WHOLE_FILE + else: + assignment_type = AssignmentType.PARTIAL_HUNKS + return GroupAssignment( + file_path=existing.file_path, + assignment_type=assignment_type, + hunk_indices=sorted(set(existing.hunk_indices) | set(incoming.hunk_indices)), + ) + + def _merge_chunk_groups(accumulated: list[Group], chunk_groups: list[Group]) -> list[Group]: acc_map = {g.id: g for g in accumulated} for cg in chunk_groups: if cg.id in acc_map: existing = acc_map[cg.id] - existing.assignments.extend(cg.assignments) + # A later chunk may assign more hunks of a file this group already + # holds. Keep one assignment per path: materialization writes a + # file once per assignment, so duplicates would drop hunks. + by_path = {a.file_path: a for a in existing.assignments} + for incoming in cg.assignments: + if incoming.file_path in by_path: + by_path[incoming.file_path] = _merge_assignment( + by_path[incoming.file_path], incoming + ) + else: + by_path[incoming.file_path] = incoming + existing.assignments = list(by_path.values()) for dep in cg.depends_on: if dep not in existing.depends_on: existing.depends_on.append(dep) diff --git a/tests/test_client.py b/tests/test_client.py index 1bdfffa..2542c8d 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -986,3 +986,80 @@ def test_large_diff_uses_chunking( result = _plan_split_with_llm(parsed, settings) assert len(result) == 1 mock_chunked.assert_called_once() + + +class TestMergeChunkGroupsSameFile: + def test_same_file_in_later_chunk_is_merged_into_one_assignment(self) -> None: + first = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="f.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ) + ], + ) + later = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="f.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[1], + ) + ], + ) + (merged,) = _merge_chunk_groups([first], [later]) + assert len(merged.assignments) == 1 + assert merged.assignments[0].hunk_indices == [0, 1] + assert merged.assignments[0].assignment_type == AssignmentType.PARTIAL_HUNKS + + def test_whole_file_wins_when_merging(self) -> None: + partial = GroupAssignment( + file_path="f.py", assignment_type=AssignmentType.PARTIAL_HUNKS, hunk_indices=[0] + ) + whole = GroupAssignment( + file_path="f.py", assignment_type=AssignmentType.WHOLE_FILE, hunk_indices=[] + ) + first = Group(id="pr-1", title="t", description="d", assignments=[partial]) + later = Group(id="pr-1", title="t", description="d", assignments=[whole]) + (merged,) = _merge_chunk_groups([first], [later]) + assert len(merged.assignments) == 1 + assert merged.assignments[0].assignment_type == AssignmentType.WHOLE_FILE + assert merged.assignments[0].hunk_indices == [0] + + def test_other_files_untouched(self) -> None: + first = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="a.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ) + ], + ) + later = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="b.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[2], + ) + ], + ) + (merged,) = _merge_chunk_groups([first], [later]) + assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [ + ("a.py", [0]), + ("b.py", [2]), + ] diff --git a/tests/test_reconstructor.py b/tests/test_reconstructor.py index e9495f6..1e9c16b 100644 --- a/tests/test_reconstructor.py +++ b/tests/test_reconstructor.py @@ -379,3 +379,57 @@ def test_added_file_content_is_not_double_spaced(self) -> None: ) result = materialize_group_files(parsed, group, "abc123") assert result["new_file.py"] == 'def hello():\n return "world"\n\n' + + +class TestMaterializeDuplicateAssignments: + @patch("pr_split.diff_ops.reconstructor._get_base_file_content") + def test_two_assignments_for_one_file_apply_both_hunks(self, mock_base: MagicMock) -> None: + mock_base.return_value = _base_content() + parsed = parse_diff(PATCH_TEXT) + group = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="example.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ), + GroupAssignment( + file_path="example.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[1], + ), + ], + ) + result = materialize_group_files(parsed, group, "main") + content = result["example.py"] + assert content is not None + assert "inserted_after_1" in content + assert "inserted_after_11" in content + assert mock_base.call_count == 1 + + def test_duplicate_assignments_on_new_file_apply_both_hunks(self) -> None: + parsed = parse_diff( + "diff --git a/n.py b/n.py\nnew file mode 100644\n--- /dev/null\n+++ b/n.py\n" + "@@ -0,0 +1,2 @@\n+one\n+two\n" + ) + group = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="n.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ), + GroupAssignment( + file_path="n.py", + assignment_type=AssignmentType.WHOLE_FILE, + hunk_indices=[], + ), + ], + ) + assert materialize_group_files(parsed, group, "main")["n.py"] == "one\ntwo\n" From 57af5e8ecbb9b3a389bc9c60d2bc7334748aeaef Mon Sep 17 00:00:00 2001 From: vitali87 Date: Sun, 30 Aug 2026 14:13:45 +0100 Subject: [PATCH 2/4] fix: log the number of files materialized, not assignments --- pr_split/diff_ops/reconstructor.py | 2 +- tests/test_reconstructor.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pr_split/diff_ops/reconstructor.py b/pr_split/diff_ops/reconstructor.py index 3b70c34..68b252c 100644 --- a/pr_split/diff_ops/reconstructor.py +++ b/pr_split/diff_ops/reconstructor.py @@ -99,13 +99,13 @@ def _assigned_hunk_indices( def materialize_group_files( parsed_diff: ParsedDiff, group: Group, ref: str ) -> dict[str, str | None]: - logger.info(logs.MATERIALIZING_FILES.format(count=len(group.assignments), group=group.id)) pf_map = {pf.path: pf for pf in parsed_diff.patch_set} # Several assignments may name the same file (e.g. merged across diff # chunks); each file is written once from the union of their hunks. assignments_by_path: dict[str, list[GroupAssignment]] = {} for assignment in group.assignments: assignments_by_path.setdefault(assignment.file_path, []).append(assignment) + logger.info(logs.MATERIALIZING_FILES.format(count=len(assignments_by_path), group=group.id)) result: dict[str, str | None] = {} for file_path, assignments in assignments_by_path.items(): patch_file = pf_map.get(file_path) diff --git a/tests/test_reconstructor.py b/tests/test_reconstructor.py index 1e9c16b..2f92385 100644 --- a/tests/test_reconstructor.py +++ b/tests/test_reconstructor.py @@ -403,12 +403,14 @@ def test_two_assignments_for_one_file_apply_both_hunks(self, mock_base: MagicMoc ), ], ) - result = materialize_group_files(parsed, group, "main") + with patch("pr_split.diff_ops.reconstructor.logger.info") as mock_log: + result = materialize_group_files(parsed, group, "main") content = result["example.py"] assert content is not None assert "inserted_after_1" in content assert "inserted_after_11" in content assert mock_base.call_count == 1 + assert "Materializing 1 file" in mock_log.call_args[0][0] def test_duplicate_assignments_on_new_file_apply_both_hunks(self) -> None: parsed = parse_diff( From ba513eb9605dca783ae45a23479422744ed93276 Mon Sep 17 00:00:00 2001 From: vitali87 Date: Sun, 30 Aug 2026 14:17:14 +0100 Subject: [PATCH 3/4] fix: keep pre-existing duplicate assignments when merging chunk groups Seeding the per-path map with a dict comprehension kept only the last existing assignment for a path, so duplicates already present in the accumulated group (chunk 1 output is stored verbatim) lost hunks. Fold every existing and incoming assignment through _merge_assignment. --- pr_split/planner/client.py | 12 ++++++------ tests/test_client.py | 25 +++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/pr_split/planner/client.py b/pr_split/planner/client.py index 2c2b0d6..68531df 100644 --- a/pr_split/planner/client.py +++ b/pr_split/planner/client.py @@ -250,14 +250,14 @@ def _merge_chunk_groups(accumulated: list[Group], chunk_groups: list[Group]) -> # A later chunk may assign more hunks of a file this group already # holds. Keep one assignment per path: materialization writes a # file once per assignment, so duplicates would drop hunks. - by_path = {a.file_path: a for a in existing.assignments} - for incoming in cg.assignments: - if incoming.file_path in by_path: - by_path[incoming.file_path] = _merge_assignment( - by_path[incoming.file_path], incoming + by_path: dict[str, GroupAssignment] = {} + for assignment in [*existing.assignments, *cg.assignments]: + if assignment.file_path in by_path: + by_path[assignment.file_path] = _merge_assignment( + by_path[assignment.file_path], assignment ) else: - by_path[incoming.file_path] = incoming + by_path[assignment.file_path] = assignment existing.assignments = list(by_path.values()) for dep in cg.depends_on: if dep not in existing.depends_on: diff --git a/tests/test_client.py b/tests/test_client.py index 2542c8d..a685f91 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1063,3 +1063,28 @@ def test_other_files_untouched(self) -> None: ("a.py", [0]), ("b.py", [2]), ] + + def test_pre_existing_duplicates_in_accumulated_group_are_kept(self) -> None: + def _pa(path: str, hunks: list[int]) -> GroupAssignment: + return GroupAssignment( + file_path=path, + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=hunks, + ) + + first = Group( + id="pr-1", title="t", description="d", assignments=[_pa("f.py", [0]), _pa("f.py", [1])] + ) + later = Group(id="pr-1", title="t", description="d", assignments=[_pa("g.py", [0])]) + (merged,) = _merge_chunk_groups([first], [later]) + assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [ + ("f.py", [0, 1]), + ("g.py", [0]), + ] + + first = Group( + id="pr-1", title="t", description="d", assignments=[_pa("f.py", [0]), _pa("f.py", [1])] + ) + later = Group(id="pr-1", title="t", description="d", assignments=[_pa("f.py", [2])]) + (merged,) = _merge_chunk_groups([first], [later]) + assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [("f.py", [0, 1, 2])] From 10301b44d3551e05d26c41f94a9ba5a455371121 Mon Sep 17 00:00:00 2001 From: vitali87 Date: Sun, 30 Aug 2026 14:20:50 +0100 Subject: [PATCH 4/4] test: make the added-file duplicate-assignment test fail without the fix --- tests/test_reconstructor.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/test_reconstructor.py b/tests/test_reconstructor.py index 2f92385..cb87ecd 100644 --- a/tests/test_reconstructor.py +++ b/tests/test_reconstructor.py @@ -413,10 +413,14 @@ def test_two_assignments_for_one_file_apply_both_hunks(self, mock_base: MagicMoc assert "Materializing 1 file" in mock_log.call_args[0][0] def test_duplicate_assignments_on_new_file_apply_both_hunks(self) -> None: + # Two hunks in a new file (unidiff splits them when the context gap + # is large enough), each claimed by a separate PARTIAL assignment. parsed = parse_diff( "diff --git a/n.py b/n.py\nnew file mode 100644\n--- /dev/null\n+++ b/n.py\n" "@@ -0,0 +1,2 @@\n+one\n+two\n" + "@@ -0,0 +10,1 @@\n+ten\n" ) + assert len(parsed.patch_set[0]) == 2 group = Group( id="pr-1", title="t", @@ -429,9 +433,9 @@ def test_duplicate_assignments_on_new_file_apply_both_hunks(self) -> None: ), GroupAssignment( file_path="n.py", - assignment_type=AssignmentType.WHOLE_FILE, - hunk_indices=[], + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[1], ), ], ) - assert materialize_group_files(parsed, group, "main")["n.py"] == "one\ntwo\n" + assert materialize_group_files(parsed, group, "main")["n.py"] == "one\ntwo\nten\n"