From b6b88dfcfe6da6a6dd0d5771fd1c4fb2594d94a4 Mon Sep 17 00:00:00 2001
From: Emerson Knapp <emerson.b.knapp@gmail.com>
Date: Sun, 23 Aug 2026 13:41:21 -0700
Subject: [PATCH] feat: auto-fix include-what-you-use cpplint items

Signed-off-by: Emerson Knapp <emerson.b.knapp@gmail.com>
---
 polymath_code_standard/checkers/cpp.py |  69 +++++++++-
 tests/test_cpp.py                      | 168 +++++++++++++++++++++++++
 2 files changed, 236 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_cpp.py

diff --git a/polymath_code_standard/checkers/cpp.py b/polymath_code_standard/checkers/cpp.py
index 97a6790..0bc2737 100644
--- a/polymath_code_standard/checkers/cpp.py
+++ b/polymath_code_standard/checkers/cpp.py
@@ -2,12 +2,78 @@
 # SPDX-License-Identifier: Apache-2.0
 import argparse
 import os
+import re
 import shutil
 import subprocess
 import tempfile
 from pathlib import Path
 
-from polymath_code_standard.checker import CONFIG_DIR, CheckerGroup, Result, check_group
+from polymath_code_standard.checker import CONFIG_DIR, CheckerGroup, Result, check_group, tool
+
+_IWYU_RE = re.compile(r'^(.+?):(\d+):\s+Add (#include (?:<[^>]+>|"[^"]+")) for')
+
+
+def _parse_iwyu_output(output: str) -> dict[str, set[str]]:
+    missing: dict[str, set[str]] = {}
+    for line in output.splitlines():
+        m = _IWYU_RE.match(line)
+        if m:
+            filepath, _, directive = m.groups()
+            missing.setdefault(filepath, set()).add(directive)
+    return missing
+
+
+def _insertion_point(lines: list[str]) -> int:
+    """Return the line index after which to insert new #include directives."""
+    for i in range(len(lines) - 1, -1, -1):
+        if lines[i].lstrip().startswith('#include'):
+            return i + 1
+    # No existing includes: skip past the top-of-file header, insert before first real code.
+    in_block_comment = False
+    for i, line in enumerate(lines):
+        stripped = line.strip()
+        if in_block_comment:
+            if '*/' in stripped:
+                in_block_comment = False
+            continue
+        if not stripped or stripped.startswith('//'):
+            continue
+        if stripped.startswith('/*'):
+            if '*/' not in stripped[2:]:
+                in_block_comment = True
+            continue
+        if stripped.startswith(('#pragma once', '#ifndef', '#define')):
+            continue
+        return i
+    return len(lines)
+
+
+def _insert_includes(filepath: str, directives: set[str]) -> None:
+    path = Path(filepath)
+    lines = path.read_text().splitlines(keepends=True)
+    insert_at = _insertion_point(lines)
+    lines[insert_at:insert_at] = [f'{d}\n' for d in sorted(directives)]
+    path.write_text(''.join(lines))
+
+
+def fix_iwyu(cpp_files: list[str], config_path: str) -> Result:
+    """Run cpplint IWYU-only and auto-insert any missing #include directives."""
+    if not cpp_files:
+        return Result(name='iwyu-fix', passed=True, skipped=True)
+    cmd = [tool('cpplint'), f'--config={config_path}', '--filter=-all,+build/include_what_you_use'] + cpp_files
+    proc = subprocess.run(cmd, capture_output=True, text=True)
+    missing = _parse_iwyu_output(proc.stdout + proc.stderr)
+    if not missing:
+        return Result(name='iwyu-fix', passed=True, cmd=cmd)
+    for filepath, directives in missing.items():
+        _insert_includes(filepath, directives)
+    modified = sorted(missing.keys())
+    return Result(
+        name='iwyu-fix',
+        passed=False,
+        output=f'Inserted missing includes in: {", ".join(modified)}\n(files have been modified — please re-stage and recommit)',
+        cmd=cmd,
+    )
 
 
 def run_clang_format(cpp_files: list[str]) -> Result:
@@ -43,6 +109,7 @@ def run(self, args: argparse.Namespace) -> list[Result]:
         try:
             shutil.copy2(CONFIG_DIR / '.cpplint.cfg', tmp_path)
             return [
+                fix_iwyu(args.files, tmp_path.name),
                 run_clang_format(args.files),
                 self._check('cpplint', [f'--config={tmp_path.name}', '--quiet', '--output=sed'], args.files),
             ]
diff --git a/tests/test_cpp.py b/tests/test_cpp.py
new file mode 100644
index 0000000..dbc7282
--- /dev/null
+++ b/tests/test_cpp.py
@@ -0,0 +1,168 @@
+# SPDX-FileCopyrightText: 2026 Polymath Robotics, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import shutil
+import tempfile
+import uuid
+from pathlib import Path
+
+import pytest
+
+from polymath_code_standard.checker import CONFIG_DIR
+from polymath_code_standard.checkers.cpp import _insert_includes, _insertion_point, _parse_iwyu_output, fix_iwyu
+
+_PROJECT_ROOT = Path(__file__).parent.parent
+
+
+# --- _parse_iwyu_output ---
+
+
+def test_parse_single_header():
+    output = 'file.cpp:0:  Add #include <string> for string  [build/include_what_you_use] [4]'
+    assert _parse_iwyu_output(output) == {'file.cpp': {'#include <string>'}}
+
+
+def test_parse_multiple_headers_same_file():
+    output = (
+        'a.cpp:0:  Add #include <string> for string  [build/include_what_you_use] [4]\n'
+        'a.cpp:0:  Add #include <vector> for vector<>  [build/include_what_you_use] [4]\n'
+    )
+    assert _parse_iwyu_output(output) == {'a.cpp': {'#include <string>', '#include <vector>'}}
+
+
+def test_parse_multiple_files():
+    output = (
+        'a.cpp:0:  Add #include <string> for string  [build/include_what_you_use] [4]\n'
+        'b.cpp:0:  Add #include <vector> for vector<>  [build/include_what_you_use] [4]\n'
+    )
+    assert _parse_iwyu_output(output) == {'a.cpp': {'#include <string>'}, 'b.cpp': {'#include <vector>'}}
+
+
+def test_parse_quoted_header():
+    output = 'file.cpp:0:  Add #include "my_header.hpp" for MyClass  [build/include_what_you_use] [4]'
+    assert _parse_iwyu_output(output) == {'file.cpp': {'#include "my_header.hpp"'}}
+
+
+def test_parse_ignores_unrelated_errors():
+    output = 'file.cpp:5:  Some other lint error  [some/category] [3]\n'
+    assert _parse_iwyu_output(output) == {}
+
+
+def test_parse_empty_output():
+    assert _parse_iwyu_output('') == {}
+
+
+# --- _insertion_point ---
+
+
+def test_insertion_point_after_last_include():
+    lines = ['#include <iostream>\n', '#include <string>\n', '\n', 'void f() {}\n']
+    assert _insertion_point(lines) == 2
+
+
+def test_insertion_point_single_include():
+    lines = ['#include <iostream>\n', '\n', 'void f() {}\n']
+    assert _insertion_point(lines) == 1
+
+
+def test_insertion_point_no_includes_skips_comment_header():
+    lines = ['// License header\n', '// More comments\n', '\n', 'void f() {}\n']
+    # Should insert before first real code (index 3)
+    assert _insertion_point(lines) == 3
+
+
+def test_insertion_point_no_includes_skips_pragma_once():
+    lines = ['// License\n', '#pragma once\n', '\n', 'void f() {}\n']
+    assert _insertion_point(lines) == 3
+
+
+def test_insertion_point_empty_file():
+    assert _insertion_point([]) == 0
+
+
+# --- _insert_includes ---
+
+
+def test_insert_after_last_include(tmp_path):
+    f = tmp_path / 'test.cpp'
+    f.write_text('#include <iostream>\n\nvoid f() {}\n')
+    _insert_includes(str(f), {'#include <string>'})
+    lines = f.read_text().splitlines()
+    assert lines[0] == '#include <iostream>'
+    assert lines[1] == '#include <string>'
+
+
+def test_insert_multiple_sorted(tmp_path):
+    f = tmp_path / 'test.cpp'
+    f.write_text('#include <iostream>\n\nvoid f() {}\n')
+    _insert_includes(str(f), {'#include <vector>', '#include <string>'})
+    lines = f.read_text().splitlines()
+    inserted = [line for line in lines if line.startswith('#include <s') or line.startswith('#include <v')]
+    assert inserted == ['#include <string>', '#include <vector>']
+
+
+def test_insert_no_existing_includes(tmp_path):
+    f = tmp_path / 'test.cpp'
+    f.write_text('// License header\n\nvoid f() {}\n')
+    _insert_includes(str(f), {'#include <string>'})
+    assert '#include <string>' in f.read_text()
+
+
+# --- fix_iwyu integration ---
+
+
+@pytest.fixture
+def project_files():
+    """Create temp files inside the project root. cpplint --config requires a bare filename
+    with no directory components, so the config and source files must live in cwd (project root)."""
+    created = []
+
+    def _make(name: str, content: str) -> str:
+        d = _PROJECT_ROOT / f'.pytest_tmp_{uuid.uuid4().hex[:8]}'
+        d.mkdir()
+        p = d / name
+        p.write_text(content, encoding='utf-8')
+        created.append(d)
+        return str(p)
+
+    yield _make
+
+    for d in created:
+        shutil.rmtree(d, ignore_errors=True)
+
+
+@pytest.fixture
+def cpplint_cfg():
+    """Temp copy of .cpplint.cfg in the project root (cwd), yielding its bare filename."""
+    fd, path = tempfile.mkstemp(dir=_PROJECT_ROOT, prefix='.cpplint_', suffix='.cfg')
+    os.close(fd)
+    try:
+        shutil.copy2(CONFIG_DIR / '.cpplint.cfg', path)
+        yield Path(path).name
+    finally:
+        Path(path).unlink(missing_ok=True)
+
+
+def test_fix_iwyu_inserts_missing_header(project_files, cpplint_cfg):
+    content = (
+        '// SPDX-FileCopyrightText: 2026 Polymath Robotics, Inc.\n#include <iostream>\nvoid f() { std::string s; }\n'
+    )
+    f = project_files('test.cpp', content)
+    result = fix_iwyu([f], cpplint_cfg)
+    assert not result.passed
+    assert '#include <string>' in Path(f).read_text()
+
+
+def test_fix_iwyu_no_issues(project_files, cpplint_cfg):
+    content = (
+        '// SPDX-FileCopyrightText: 2026 Polymath Robotics, Inc.\n#include <string>\nvoid f() { std::string s; }\n'
+    )
+    f = project_files('test.cpp', content)
+    result = fix_iwyu([f], cpplint_cfg)
+    assert result.passed
+
+
+def test_fix_iwyu_skips_empty_list(cpplint_cfg):
+    result = fix_iwyu([], cpplint_cfg)
+    assert result.passed
+    assert result.skipped